Repository: valyala/fasthttp
Pull request #537 (2019-02-12 05:10:48)

add conn's address info in Response
for #535 

erikdubbelboer: If you rebase on master, I have just [added some tests to make sure we have 0 allocations](https://github.com/valyala/fasthttp/commit/61039c93c2a31d3be661c50269d5800090f44175).
xuecai: I learned a lot. Thanks for allocations test.
i'm sorry for reset commits, i made a mistake when I  rebase code
xuecai: I think I have done it in [Response.Reset()](https://github.com/valyala/fasthttp/pull/537/files#diff-31dcde2645a9daeff4e0064653331de4R846)?
erikdubbelboer: Sorry I missed that.

@kirillDanshin what do you think? If you agree then you can merge it.
xuecai: @erikdubbelboer @kirillDanshin i'm sorry, i find a bug on response copy and i fix it. 
please reivew.
xuecai: i find CopyTo miss tests in fasthttp.
may i add then by another pr later?
xuecai: hi, Is there any other advice for this pr? @erikdubbelboer @kirillDanshin  

Patch:
diff --git a/client.go b/client.go
index 5c0d062..89e9808 100644
--- a/client.go
+++ b/client.go
@@ -1146,6 +1146,8 @@ func (c *HostClient) doNonNilReqResp(req *Request, resp *Response) (bool, error)
 	}
 	conn := cc.c
 
+	resp.parseNetConn(conn)
+
 	if c.WriteTimeout > 0 {
 		// Set Deadline every time, since golang has fixed the performance issue
 		// See https://github.com/golang/go/issues/15133#issuecomment-271571395 for details
@@ -2087,6 +2089,8 @@ func (c *pipelineConnClient) writer(conn net.Conn, stopCh <-chan struct{}) error
 			continue
 		}
 
+		w.resp.parseNetConn(conn)
+
 		if writeTimeout > 0 {
 			// Set Deadline every time, since golang has fixed the performance issue
 			// See https://github.com/golang/go/issues/15133#issuecomment-271571395 for details
diff --git a/cookie.go b/cookie.go
index 1d9861d..8137643 100644
--- a/cookie.go
+++ b/cookie.go
@@ -21,6 +21,7 @@ var (
 // CookieSameSite is an enum for the mode in which the SameSite flag should be set for the given cookie.
 // See https://tools.ietf.org/html/draft-ietf-httpbis-cookie-same-site-00 for details.
 type CookieSameSite int
+
 const (
 	// CookieSameSiteDisabled removes the SameSite flag
 	CookieSameSiteDisabled CookieSameSite = iota
diff --git a/http.go b/http.go
index 1458602..10dc465 100644
--- a/http.go
+++ b/http.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"mime/multipart"
+	"net"
 	"os"
 	"sync"
 
@@ -75,6 +76,11 @@ type Response struct {
 	SkipBody bool
 
 	keepBodyBuffer bool
+
+	// Remote TCPAddr from concurrently net.Conn
+	raddr net.Addr
+	// Local TCPAddr from concurrently net.Conn
+	laddr net.Addr
 }
 
 // SetHost sets host for the request.
@@ -281,6 +287,23 @@ func (w *requestBodyWriter) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
+func (resp *Response) parseNetConn(conn net.Conn) {
+	resp.raddr = conn.RemoteAddr()
+	resp.laddr = conn.LocalAddr()
+}
+
+// RemoteAddr returns the remote network address. The Addr returned is shared
+// by all invocations of RemoteAddr, so do not modify it.
+func (resp *Response) RemoteAddr() net.Addr {
+	return resp.raddr
+}
+
+// LocalAddr returns the local network address. The Addr returned is shared
+// by all invocations of LocalAddr, so do not modify it.
+func (resp *Response) LocalAddr() net.Addr {
+	return resp.laddr
+}
+
 // Body returns response body.
 //
 // The returned body is valid until the response modification.
@@ -627,6 +650,8 @@ func (resp *Response) copyToSkipBody(dst *Response) {
 	dst.Reset()
 	resp.Header.CopyTo(&dst.Header)
 	dst.SkipBody = resp.SkipBody
+	dst.raddr = resp.raddr
+	dst.laddr = resp.laddr
 }
 
 func swapRequestBody(a, b *Request) {
@@ -820,6 +845,8 @@ func (resp *Response) Reset() {
 	resp.Header.Reset()
 	resp.resetSkipHeader()
 	resp.SkipBody = false
+	resp.raddr = nil
+	resp.laddr = nil
 }
 
 func (resp *Response) resetSkipHeader() {


Test patch:
diff --git a/client_test.go b/client_test.go
index a6310a0..2eb7e41 100644
--- a/client_test.go
+++ b/client_test.go
@@ -8,6 +8,7 @@ import (
 	"net"
 	"net/url"
 	"os"
+	"regexp"
 	"runtime"
 	"strings"
 	"sync"
@@ -18,6 +19,38 @@ import (
 	"github.com/valyala/fasthttp/fasthttputil"
 )
 
+func TestClientParseConn(t *testing.T) {
+	network := "tcp"
+	ln, _ := net.Listen(network, "127.0.0.1:0")
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			return
+		},
+	}
+	go s.Serve(ln)
+	host := ln.Addr().String()
+	c := &Client{}
+	req, res := AcquireRequest(), AcquireResponse()
+	defer func() {
+		ReleaseRequest(req)
+		ReleaseResponse(res)
+	}()
+	req.SetRequestURI("http://" + host + "")
+	c.Do(req, res)
+
+	if res.RemoteAddr().Network() != network {
+		t.Fatalf("req RemoteAddr parse network fail: %s, hope: %s", res.RemoteAddr().Network(), network)
+	}
+	if host != res.RemoteAddr().String() {
+		t.Fatalf("req RemoteAddr parse addr fail: %s, hope: %s", res.RemoteAddr().String(), host)
+	}
+
+	if !regexp.MustCompile(`^127\.0\.0\.1:[0-9]{4,5}$`).MatchString(res.LocalAddr().String()) {
+		t.Fatalf("res LocalAddr addr match fail: %s, hope match: %s", res.LocalAddr().String(), "^127.0.0.1:[0-9]{4,5}$")
+	}
+
+}
+
 func TestClientPostArgs(t *testing.T) {
 	ln := fasthttputil.NewInmemoryListener()
 	s := &Server{
@@ -36,6 +69,10 @@ func TestClientPostArgs(t *testing.T) {
 		},
 	}
 	req, res := AcquireRequest(), AcquireResponse()
+	defer func() {
+		ReleaseRequest(req)
+		ReleaseResponse(res)
+	}()
 	args := req.PostArgs()
 	args.Add("addhttp2", "support")
 	args.Add("fast", "http")
@@ -1157,6 +1194,14 @@ func (r *readTimeoutConn) Close() error {
 	return nil
 }
 
+func (r *readTimeoutConn) LocalAddr() net.Addr {
+	return nil
+}
+
+func (r *readTimeoutConn) RemoteAddr() net.Addr {
+	return nil
+}
+
 func TestClientNonIdempotentRetry(t *testing.T) {
 	dialsCount := 0
 	c := &Client{
@@ -1271,6 +1316,14 @@ func (w *writeErrorConn) Close() error {
 	return nil
 }
 
+func (r *writeErrorConn) LocalAddr() net.Addr {
+	return nil
+}
+
+func (r *writeErrorConn) RemoteAddr() net.Addr {
+	return nil
+}
+
 type readErrorConn struct {
 	net.Conn
 }
@@ -1287,6 +1340,14 @@ func (r *readErrorConn) Close() error {
 	return nil
 }
 
+func (r *readErrorConn) LocalAddr() net.Addr {
+	return nil
+}
+
+func (r *readErrorConn) RemoteAddr() net.Addr {
+	return nil
+}
+
 type singleReadConn struct {
 	net.Conn
 	s string
@@ -1310,6 +1371,14 @@ func (r *singleReadConn) Close() error {
 	return nil
 }
 
+func (r *singleReadConn) LocalAddr() net.Addr {
+	return nil
+}
+
+func (r *singleReadConn) RemoteAddr() net.Addr {
+	return nil
+}
+
 func TestClientHTTPSInvalidServerName(t *testing.T) {
 	sHTTPS := startEchoServerTLS(t, "tcp", "127.0.0.1:")
 	defer sHTTPS.Stop()
diff --git a/http_test.go b/http_test.go
index 16a101a..dd4d3ca 100644
--- a/http_test.go
+++ b/http_test.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"io/ioutil"
 	"mime/multipart"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -14,6 +15,24 @@ import (
 	"github.com/valyala/bytebufferpool"
 )
 
+func TestResponseCopyTo(t *testing.T) {
+	resp := &Response{}
+	copyResp := &Response{}
+
+	// init resp
+	resp.laddr = zeroTCPAddr
+	resp.SkipBody = true
+	resp.Header.SetStatusCode(200)
+	resp.SetBodyString("test")
+
+	resp.CopyTo(copyResp)
+
+	if !reflect.DeepEqual(resp, copyResp) {
+		t.Fatal("ResponseCopyTo fail")
+	}
+
+}
+
 func TestResponseBodyStreamDeflate(t *testing.T) {
 	body := createFixedBody(1e5)