Repository: valyala/fasthttp
Pull request #990 (2021-03-05 05:25:50)

Improve socks proxy
Use `proxy.FromURL` to support auth


Patch:
diff --git a/client.go b/client.go
index 80684c3..7e74443 100644
--- a/client.go
+++ b/client.go
@@ -2152,6 +2152,24 @@ type PipelineClient struct {
 	// since unfortunately ipv6 remains broken in many networks worldwide :)
 	DialDualStack bool
 
+	// Response header names are passed as-is without normalization
+	// if this option is set.
+	//
+	// Disabled header names' normalization may be useful only for proxying
+	// responses to other clients expecting case-sensitive
+	// header names. See https://github.com/valyala/fasthttp/issues/57
+	// for details.
+	//
+	// By default request and response header names are normalized, i.e.
+	// The first letter and the first letters following dashes
+	// are uppercased, while all the other letters are lowercased.
+	// Examples:
+	//
+	//     * HOST -> Host
+	//     * content-type -> Content-Type
+	//     * cONTENT-lenGTH -> Content-Length
+	DisableHeaderNamesNormalizing bool
+
 	// Path values are sent as-is without normalization
 	//
 	// Disabled path normalization may be useful for proxying incoming requests
@@ -2206,20 +2224,21 @@ type PipelineClient struct {
 type pipelineConnClient struct {
 	noCopy noCopy //nolint:unused,structcheck
 
-	Addr                   string
-	MaxPendingRequests     int
-	MaxBatchDelay          time.Duration
-	Dial                   DialFunc
-	DialDualStack          bool
-	DisablePathNormalizing bool
-	IsTLS                  bool
-	TLSConfig              *tls.Config
-	MaxIdleConnDuration    time.Duration
-	ReadBufferSize         int
-	WriteBufferSize        int
-	ReadTimeout            time.Duration
-	WriteTimeout           time.Duration
-	Logger                 Logger
+	Addr                          string
+	MaxPendingRequests            int
+	MaxBatchDelay                 time.Duration
+	Dial                          DialFunc
+	DialDualStack                 bool
+	DisableHeaderNamesNormalizing bool
+	DisablePathNormalizing        bool
+	IsTLS                         bool
+	TLSConfig                     *tls.Config
+	MaxIdleConnDuration           time.Duration
+	ReadBufferSize                int
+	WriteBufferSize               int
+	ReadTimeout                   time.Duration
+	WriteTimeout                  time.Duration
+	Logger                        Logger
 
 	workPool sync.Pool
 
@@ -2298,6 +2317,7 @@ func (c *pipelineConnClient) DoDeadline(req *Request, resp *Response, deadline t
 	}
 
 	w := acquirePipelineWork(&c.workPool, timeout)
+	w.respCopy.Header.disableNormalizing = c.DisableHeaderNamesNormalizing
 	w.req = &w.reqCopy
 	w.resp = &w.respCopy
 
@@ -2354,6 +2374,8 @@ func (c *PipelineClient) Do(req *Request, resp *Response) error {
 func (c *pipelineConnClient) Do(req *Request, resp *Response) error {
 	c.init()
 
+	resp.Header.disableNormalizing = c.DisableHeaderNamesNormalizing
+
 	if c.DisablePathNormalizing {
 		req.URI().DisablePathNormalizing = true
 	}
@@ -2436,20 +2458,21 @@ func (c *PipelineClient) getConnClientUnlocked() *pipelineConnClient {
 
 func (c *PipelineClient) newConnClient() *pipelineConnClient {
 	cc := &pipelineConnClient{
-		Addr:                   c.Addr,
-		MaxPendingRequests:     c.MaxPendingRequests,
-		MaxBatchDelay:          c.MaxBatchDelay,
-		Dial:                   c.Dial,
-		DialDualStack:          c.DialDualStack,
-		DisablePathNormalizing: c.DisablePathNormalizing,
-		IsTLS:                  c.IsTLS,
-		TLSConfig:              c.TLSConfig,
-		MaxIdleConnDuration:    c.MaxIdleConnDuration,
-		ReadBufferSize:         c.ReadBufferSize,
-		WriteBufferSize:        c.WriteBufferSize,
-		ReadTimeout:            c.ReadTimeout,
-		WriteTimeout:           c.WriteTimeout,
-		Logger:                 c.Logger,
+		Addr:                          c.Addr,
+		MaxPendingRequests:            c.MaxPendingRequests,
+		MaxBatchDelay:                 c.MaxBatchDelay,
+		Dial:                          c.Dial,
+		DialDualStack:                 c.DialDualStack,
+		DisableHeaderNamesNormalizing: c.DisableHeaderNamesNormalizing,
+		DisablePathNormalizing:        c.DisablePathNormalizing,
+		IsTLS:                         c.IsTLS,
+		TLSConfig:                     c.TLSConfig,
+		MaxIdleConnDuration:           c.MaxIdleConnDuration,
+		ReadBufferSize:                c.ReadBufferSize,
+		WriteBufferSize:               c.WriteBufferSize,
+		ReadTimeout:                   c.ReadTimeout,
+		WriteTimeout:                  c.WriteTimeout,
+		Logger:                        c.Logger,
 	}
 	c.connClients = append(c.connClients, cc)
 	return cc
diff --git a/fasthttpproxy/socks5.go b/fasthttpproxy/socks5.go
index 5752704..b3339a4 100644
--- a/fasthttpproxy/socks5.go
+++ b/fasthttpproxy/socks5.go
@@ -2,6 +2,7 @@ package fasthttpproxy
 
 import (
 	"net"
+	"net/url"
 
 	"github.com/valyala/fasthttp"
 	"golang.org/x/net/proxy"
@@ -15,7 +16,14 @@ import (
 //		Dial: fasthttpproxy.FasthttpSocksDialer("localhost:9050"),
 //	}
 func FasthttpSocksDialer(proxyAddr string) fasthttp.DialFunc {
-	dialer, err := proxy.SOCKS5("tcp", proxyAddr, nil, proxy.Direct)
+	var (
+		u      *url.URL
+		err    error
+		dialer proxy.Dialer
+	)
+	if u, err = url.Parse(proxyAddr); err == nil {
+		dialer, err = proxy.FromURL(u, proxy.Direct)
+	}
 	// It would be nice if we could return the error here. But we can't
 	// change our API so just keep returning it in the returned Dial function.
 	// Besides the implementation of proxy.SOCKS5() at the time of writing this


Test patch:
diff --git a/client_test.go b/client_test.go
index e269553..48b5b85 100644
--- a/client_test.go
+++ b/client_test.go
@@ -932,6 +932,76 @@ func testPipelineClientDo(t *testing.T, c *PipelineClient) {
 	ReleaseResponse(resp)
 }
 
+func TestPipelineClientDoDisableHeaderNamesNormalizing(t *testing.T) {
+	t.Parallel()
+
+	testPipelineClientDisableHeaderNamesNormalizing(t, 0)
+}
+
+func TestPipelineClientDoTimeoutDisableHeaderNamesNormalizing(t *testing.T) {
+	t.Parallel()
+
+	testPipelineClientDisableHeaderNamesNormalizing(t, time.Second)
+}
+
+func testPipelineClientDisableHeaderNamesNormalizing(t *testing.T, timeout time.Duration) {
+	ln := fasthttputil.NewInmemoryListener()
+
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			ctx.Response.Header.Set("foo-BAR", "baz")
+		},
+		DisableHeaderNamesNormalizing: true,
+	}
+
+	serverStopCh := make(chan struct{})
+	go func() {
+		if err := s.Serve(ln); err != nil {
+			t.Errorf("unexpected error: %s", err)
+		}
+		close(serverStopCh)
+	}()
+
+	c := &PipelineClient{
+		Dial: func(addr string) (net.Conn, error) {
+			return ln.Dial()
+		},
+		DisableHeaderNamesNormalizing: true,
+	}
+
+	var req Request
+	req.SetRequestURI("http://aaaai.com/bsdf?sddfsd")
+	var resp Response
+	for i := 0; i < 5; i++ {
+		if timeout > 0 {
+			if err := c.DoTimeout(&req, &resp, timeout); err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+		} else {
+			if err := c.Do(&req, &resp); err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+		}
+		hv := resp.Header.Peek("foo-BAR")
+		if string(hv) != "baz" {
+			t.Fatalf("unexpected header value: %q. Expecting %q", hv, "baz")
+		}
+		hv = resp.Header.Peek("Foo-Bar")
+		if len(hv) > 0 {
+			t.Fatalf("unexpected non-empty header value %q", hv)
+		}
+	}
+
+	if err := ln.Close(); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	select {
+	case <-serverStopCh:
+	case <-time.After(time.Second):
+		t.Fatalf("timeout")
+	}
+}
+
 func TestClientDoTimeoutDisableHeaderNamesNormalizing(t *testing.T) {
 	t.Parallel()