Repository: valyala/fasthttp
Pull request #1045 (2021-06-13 17:27:03)

Defined Transport for the client
Solves https://github.com/valyala/fasthttp/issues/972

dgrr: Notice that this implementation only works with the `HostClient`
erikdubbelboer: What is the reason for not adding it to `Client` as well?
dgrr: > What is the reason for not adding it to `Client` as well?

I'll in the future, but using the HostClient.Transport, and also I need more time for that. Doing so now in the Client is a little bit problematic. The ideal scenario would be to have something like the Server's NextProtos in the Client. So, on every new connection the HostClient creates, we'll need to check the TLS ALPN fields and if there's any extension for that has been added to fasthttp through a 3rd party library, then put it as TransportFunc.

The main problem is that the Client doesn't own the connections but the HostClient does. And the HTTP/2 lib manages it's own connections, therefore we assume all the `Transport` implementations are going to manage their own connections. Because at the end, the Transport call should only be used to send requests and get a response back.

As said, I'll try to support Transport in Client, but I'll need more time.
dgrr: @erikdubbelboer 

Patch:
diff --git a/client.go b/client.go
index f135ec3..4736321 100644
--- a/client.go
+++ b/client.go
@@ -294,9 +294,9 @@ type Client struct {
 	// By default will use isIdempotent function
 	RetryIf RetryIfFunc
 
-	mLock sync.Mutex
-	m     map[string]*HostClient
-	ms    map[string]*HostClient
+	mLock      sync.Mutex
+	m          map[string]*HostClient
+	ms         map[string]*HostClient
 	readerPool sync.Pool
 	writerPool sync.Pool
 }
@@ -594,6 +594,9 @@ type DialFunc func(addr string) (net.Conn, error)
 // Request argument passed to RetryIfFunc, if there are any request errors.
 type RetryIfFunc func(request *Request) bool
 
+// TransportFunc wraps every request/response.
+type TransportFunc func(*Request, *Response) error
+
 // HostClient balances http requests among hosts listed in Addr.
 //
 // HostClient may be used for balancing load among multiple upstream hosts.
@@ -745,6 +748,9 @@ type HostClient struct {
 	// By default will use isIdempotent function
 	RetryIf RetryIfFunc
 
+	// Transport defines a transport-like mechanism that wraps every request/response.
+	Transport TransportFunc
+
 	clientName  atomic.Value
 	lastUseTime uint32
 
@@ -1360,8 +1366,16 @@ func (c *HostClient) doNonNilReqResp(req *Request, resp *Response) (bool, error)
 	resp.Reset()
 	resp.SkipBody = customSkipBody
 
-	if c.DisablePathNormalizing {
-		req.URI().DisablePathNormalizing = true
+	req.URI().DisablePathNormalizing = c.DisablePathNormalizing
+
+	userAgentOld := req.Header.UserAgent()
+	if len(userAgentOld) == 0 {
+		req.Header.userAgent = append(req.Header.userAgent[:0], c.getClientName()...)
+	}
+
+	if c.Transport != nil {
+		err := c.Transport(req, resp)
+		return err == nil, err
 	}
 
 	cc, err := c.acquireConn(req.timeout, req.ConnectionClose())
@@ -1388,10 +1402,6 @@ func (c *HostClient) doNonNilReqResp(req *Request, resp *Response) (bool, error)
 		resetConnection = true
 	}
 
-	userAgentOld := req.Header.UserAgent()
-	if len(userAgentOld) == 0 {
-		req.Header.userAgent = append(req.Header.userAgent[:0], c.getClientName()...)
-	}
 	bw := c.acquireWriter(conn)
 	err = req.Write(bw)
 


Test patch:
diff --git a/client_test.go b/client_test.go
index a6be61e..0a61f0d 100644
--- a/client_test.go
+++ b/client_test.go
@@ -1952,6 +1952,70 @@ func TestClientRetryRequestWithCustomDecider(t *testing.T) {
 	}
 }
 
+func TestHostClientTransport(t *testing.T) {
+	t.Parallel()
+
+	ln := fasthttputil.NewInmemoryListener()
+
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			ctx.WriteString("abcd") //nolint:errcheck
+		},
+	}
+	serverStopCh := make(chan struct{})
+	go func() {
+		if err := s.Serve(ln); err != nil {
+			t.Errorf("unexpected error: %s", err)
+		}
+		close(serverStopCh)
+	}()
+
+	c := &HostClient{
+		Addr: "foobar",
+		Transport: func() TransportFunc {
+			c, _ := ln.Dial()
+
+			br := bufio.NewReader(c)
+			bw := bufio.NewWriter(c)
+
+			return func(req *Request, res *Response) error {
+				if err := req.Write(bw); err != nil {
+					return err
+				}
+
+				if err := bw.Flush(); err != nil {
+					return err
+				}
+
+				return res.Read(br)
+			}
+		}(),
+	}
+
+	for i := 0; i < 5; i++ {
+		statusCode, body, err := c.Get(nil, "http://aaaa.com/bbb/cc")
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		if statusCode != StatusOK {
+			t.Fatalf("unexpected status code %d. Expecting %d", statusCode, StatusOK)
+		}
+		if string(body) != "abcd" {
+			t.Fatalf("unexpected body %q. Expecting %q", body, "abcd")
+		}
+	}
+
+	if err := ln.Close(); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	select {
+	case <-serverStopCh:
+	case <-time.After(time.Second):
+		t.Fatalf("timeout")
+	}
+}
+
 type writeErrorConn struct {
 	net.Conn
 }