package main
import (
"bytes"
"net/url"
"testing"
)
type AttrTestCase struct {
AttrName []byte
AttrValue []byte
ExpectedOutput []byte
}
type StringTestCase struct {
Input string
ExpectedOutput string
}
var attrTestData []*AttrTestCase = []*AttrTestCase{
&AttrTestCase{
[]byte("href"),
[]byte("./x"),
[]byte(` href="./?mortyurl=http%3A%2F%2F127.0.0.1%2Fx"`),
},
&AttrTestCase{
[]byte("src"),
[]byte("http://x.com/y"),
[]byte(` src="./?mortyurl=http%3A%2F%2Fx.com%2Fy"`),
},
&AttrTestCase{
[]byte("action"),
[]byte("/z"),
[]byte(` action="./?mortyurl=http%3A%2F%2F127.0.0.1%2Fz"`),
},
&AttrTestCase{
[]byte("onclick"),
[]byte("console.log(document.cookies)"),
nil,
},
}
var urlTestData []*StringTestCase = []*StringTestCase{
&StringTestCase{
"http://x.com/",
"./?mortyurl=http%3A%2F%2Fx.com%2F",
},
&StringTestCase{
"http://a@x.com/",
"./?mortyurl=http%3A%2F%2Fa%40x.com%2F",
},
}
func TestAttrSanitizer(t *testing.T) {
u, _ := url.Parse("http://127.0.0.1/")
rc := &RequestConfig{BaseURL: u}
for _, testCase := range attrTestData {
out := bytes.NewBuffer(nil)
sanitizeAttr(rc, out, testCase.AttrName, testCase.AttrValue, testCase.AttrValue)
res, _ := out.ReadBytes(byte(0))
if !bytes.Equal(res, testCase.ExpectedOutput) {
t.Errorf(
`Attribute parse error. Name: "%s", Value: "%s", Expected: %s, Got: "%s"`,
testCase.AttrName,
testCase.AttrValue,
testCase.ExpectedOutput,
res,
)
}
}
}
func TestURLProxifier(t *testing.T) {
u, _ := url.Parse("http://127.0.0.1/")
rc := &RequestConfig{BaseURL: u}
for _, testCase := range urlTestData {
newUrl, err := rc.ProxifyURI(testCase.Input)
if err != nil {
t.Errorf("Failed to parse URL: %s", testCase.Input)
}
if newUrl != testCase.ExpectedOutput {
t.Errorf(
`URL proxifier error. Expected: "%s", Got: "%s"`,
testCase.ExpectedOutput,
newUrl,
)
}
}
}
var BENCH_SIMPLE_HTML []byte = []byte(`
test
Test heading
`)
func BenchmarkSanitizeSimpleHTML(b *testing.B) {
u, _ := url.Parse("http://127.0.0.1/")
rc := &RequestConfig{BaseURL: u}
b.ResetTimer()
for i := 0; i < b.N; i++ {
out := bytes.NewBuffer(nil)
sanitizeHTML(rc, out, BENCH_SIMPLE_HTML)
}
}
var BENCH_COMPLEX_HTML []byte = []byte(`
test 2
Test heading
`)
func BenchmarkSanitizeComplexHTML(b *testing.B) {
u, _ := url.Parse("http://127.0.0.1/")
rc := &RequestConfig{BaseURL: u}
b.ResetTimer()
for i := 0; i < b.N; i++ {
out := bytes.NewBuffer(nil)
sanitizeHTML(rc, out, BENCH_COMPLEX_HTML)
}
}