ghttp

go.mod

@@ -13,9 +13,10 @@ require (
 	github.com/klauspost/pgzip v1.2.5
 	github.com/magiconair/properties v1.8.1
 	github.com/mattn/go-sqlite3 v1.14.16
+	github.com/panjf2000/ants/v2 v2.9.0
 	github.com/pkg/errors v0.8.1
 	github.com/spf13/viper v1.7.1
-	github.com/stretchr/testify v1.8.0
+	github.com/stretchr/testify v1.8.2
 	github.com/ulikunitz/xz v0.5.8
 	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9
 	golang.org/x/net v0.0.0-20201224014010-6772e930b67b

util/http/client.go

@@ -102,6 +102,8 @@ type HTTPClient struct {
 	preHandler        func(r *http.Request)
 	keepalive         bool
 	proxyUsed         *url.URL
+	beforeDo          []BeforeDoClientFunc
+	afterDo           []AfterDoClientFunc
 }
 
 // NewHTTPClient new a HTTPClient, all request shared one http.Client object, keep cookies, keepalive etc.
@@ -113,6 +115,64 @@ func NewHTTPClient() *HTTPClient {
 	}
 }
 
+// SetBeforeDo sets callback call before request sent.
+func (s *HTTPClient) SetBeforeDo(beforeDo BeforeDoClientFunc) *HTTPClient {
+	return s.setBeforeDo(beforeDo, true)
+}
+
+// AppendBeforeDo add a callback call before request sent.
+func (s *HTTPClient) AppendBeforeDo(beforeDo BeforeDoClientFunc) *HTTPClient {
+	return s.setBeforeDo(beforeDo, false)
+}
+
+func (s *HTTPClient) setBeforeDo(beforeDo BeforeDoClientFunc, isSet bool) *HTTPClient {
+	if isSet {
+		if beforeDo != nil {
+			s.beforeDo = []BeforeDoClientFunc{beforeDo}
+		} else {
+			s.beforeDo = []BeforeDoClientFunc{}
+		}
+	} else if beforeDo != nil {
+		s.beforeDo = append(s.beforeDo, beforeDo)
+	}
+	return s
+}
+
+// SetAfterDo sets callback call after request sent.
+func (s *HTTPClient) SetAfterDo(afterDo AfterDoClientFunc) *HTTPClient {
+	return s.setAfterDo(afterDo, true)
+}
+
+// AppendAfterDo add a callback call after request sent.
+func (s *HTTPClient) AppendAfterDo(afterDo AfterDoClientFunc) *HTTPClient {
+	return s.setAfterDo(afterDo, false)
+}
+
+func (s *HTTPClient) setAfterDo(afterDo AfterDoClientFunc, isSet bool) *HTTPClient {
+	if isSet {
+		if afterDo != nil {
+			s.afterDo = []AfterDoClientFunc{afterDo}
+		} else {
+			s.afterDo = []AfterDoClientFunc{}
+		}
+	} else if afterDo != nil {
+		s.afterDo = append(s.afterDo, afterDo)
+	}
+	return s
+}
+
+func (s *HTTPClient) callBeforeDo(req *http.Request) {
+	for _, f := range s.beforeDo {
+		f(req)
+	}
+}
+
+func (s *HTTPClient) callAfterDo(req *http.Request, resp *http.Response, err error) {
+	for _, f := range s.afterDo {
+		f(req, resp, err)
+	}
+}
+
 // ProxyUsed returns the final proxy URL used by connect to target
 func (s *HTTPClient) ProxyUsed() *url.URL {
 	return s.proxyUsed
@@ -319,7 +379,9 @@ func (s *HTTPClient) Do(req *http.Request, timeout time.Duration) (resp *http.Re
 	if s.preHandler != nil {
 		s.preHandler(req)
 	}
+	s.callBeforeDo(req)
 	resp, err = client.Do(req)
+	s.callAfterDo(req, resp, err)
 	return
 }
 
@@ -401,13 +463,13 @@ func (s *HTTPClient) UploadOfReader(u, fieldName string, filename string, reader
 	return
 }
 
-//Download gets url bytes contents.
+// Download gets url bytes contents.
 func (s *HTTPClient) Download(u string, timeout time.Duration, queryData, header map[string]string) (data []byte, resp *http.Response, err error) {
 	data, _, resp, err = s.Get(u, timeout, queryData, header)
 	return
 }
 
-//DownloadToFile gets url bytes contents and save to the file.
+// DownloadToFile gets url bytes contents and save to the file.
 func (s *HTTPClient) DownloadToFile(u string, timeout time.Duration, queryData, header map[string]string, file string) (resp *http.Response, err error) {
 	f, err := os.OpenFile(file, os.O_CREATE|os.O_RDWR, 0755)
 	if err != nil {
@@ -417,7 +479,7 @@ func (s *HTTPClient) DownloadToFile(u string, timeout time.Duration, queryData,
 	return s.DownloadToWriter(u, timeout, queryData, header, f)
 }
 
-//DownloadToWriter gets url bytes contents and copy to the writer at realtime.
+// DownloadToWriter gets url bytes contents and copy to the writer at realtime.
 func (s *HTTPClient) DownloadToWriter(u string, timeout time.Duration, queryData, header map[string]string, writer io.Writer) (resp *http.Response, err error) {
 	req, cancel, err := NewGet(u, timeout, queryData, header)
 	defer cancel()

util/http/client_test.go

@@ -10,6 +10,7 @@ import (
 	"crypto/md5"
 	"fmt"
 	assert2 "github.com/stretchr/testify/assert"
+	"net/http"
 	"os"
 	"reflect"
 	"strings"
@@ -145,6 +146,65 @@ func TestHTTPClient_Get6(t *testing.T) {
 	assert.Contains(string(GetResponseBody(resp)), "hello")
 }
 
+func TestHTTPClient_Before(t *testing.T) {
+	assert := assert2.New(t)
+	client := NewHTTPClient()
+	client.SetClientCert(cert, key)
+	var r *http.Request
+	client.AppendBeforeDo(func(req *http.Request) {
+		r = req
+		req.Header.Set("test", "abc")
+	})
+	client.AppendAfterDo(func(req *http.Request, resp *http.Response, err error) {
+		resp.Header.Set("test", "def")
+	})
+	body, _, resp, _ := client.Get(httpsServerURL2+"/hello", time.Second, nil, map[string]string{"token": "200"})
+	assert.Contains(string(body), "hello")
+	assert.Contains(string(GetResponseBody(resp)), "hello")
+	assert.Equal(r.Header.Get("test"), "abc")
+	assert.Equal(resp.Header.Get("test"), "def")
+}
+
+func TestHTTPClient_BeforeSet(t *testing.T) {
+	assert := assert2.New(t)
+	client := NewHTTPClient()
+	client.SetClientCert(cert, key)
+	var r *http.Request
+	client.SetBeforeDo(func(req *http.Request) {
+		r = req
+		req.Header.Set("test", "abc")
+	})
+	client.SetAfterDo(func(req *http.Request, resp *http.Response, err error) {
+		resp.Header.Set("test", "def")
+	})
+	body, _, resp, _ := client.Get(httpsServerURL2+"/hello", time.Second, nil, map[string]string{"token": "200"})
+	assert.Contains(string(body), "hello")
+	assert.Contains(string(GetResponseBody(resp)), "hello")
+	assert.Equal(r.Header.Get("test"), "abc")
+	assert.Equal(resp.Header.Get("test"), "def")
+}
+
+func TestHTTPClient_BeforeSetNil(t *testing.T) {
+	assert := assert2.New(t)
+	client := NewHTTPClient()
+	client.SetClientCert(cert, key)
+	var r *http.Request
+	client.SetBeforeDo(func(req *http.Request) {
+		r = req
+		req.Header.Set("test", "abc")
+	})
+	client.SetBeforeDo(nil)
+	client.SetAfterDo(func(req *http.Request, resp *http.Response, err error) {
+		resp.Header.Set("test", "def")
+	})
+	client.SetAfterDo(nil)
+	body, _, resp, _ := client.Get(httpsServerURL2+"/hello", time.Second, nil, map[string]string{"token": "200"})
+	assert.Contains(string(body), "hello")
+	assert.Contains(string(GetResponseBody(resp)), "hello")
+	assert.Nil(r)
+	assert.NotEqual(resp.Header.Get("test"), "def")
+}
+
 func TestHTTPClient_SetProxyFromEnv(t *testing.T) {
 	assert := assert2.New(t)
 	os.Setenv("HTTP_PROXY", "ftp://127.0.0.1:1000")

util/http/util.go

@@ -12,8 +12,13 @@ import (
 )
 
 type BeforeDoFunc func(idx int, req *http.Request)
+
 type AfterDoFunc func(resp *Response)
 
+type BeforeDoClientFunc func(req *http.Request)
+
+type AfterDoClientFunc func(req *http.Request, resp *http.Response, err error)
+
 func NewRequest(method, URL string, timeout time.Duration, data, header map[string]string) (req *http.Request, cancel context.CancelFunc, err error) {
 	ctx, cancel := getTimeoutContext(timeout)
 	req, err = NewRequestWithContext(ctx, method, URL, data, header)