Explorar el Código

Add new TCPBufferSize parameter to network

tags/v2.0.0-rc1
9seconds hace 5 años
padre
commit
57cb1b5aa0

+ 1
- 0
mtglib/init.go Ver fichero

@@ -27,6 +27,7 @@ type Network interface {
27 27
 	DialContext(ctx context.Context, network, address string) (net.Conn, error)
28 28
 	MakeHTTPClient(func(ctx context.Context, network, address string) (net.Conn, error)) *http.Client
29 29
 	IdleTimeout() time.Duration
30
+	TCPBufferSize() int
30 31
 }
31 32
 
32 33
 type AntiReplayCache interface {

+ 4
- 0
network/default.go Ver fichero

@@ -60,6 +60,10 @@ func (d *defaultDialer) DialContext(ctx context.Context, network, address string
60 60
 	return tcpConn, nil
61 61
 }
62 62
 
63
+func (d *defaultDialer) TCPBufferSize() int {
64
+	return d.bufferSize
65
+}
66
+
63 67
 func NewDefaultDialer(timeout time.Duration, bufferSize int) (Dialer, error) {
64 68
 	switch {
65 69
 	case timeout < 0:

+ 4
- 0
network/default_test.go Ver fichero

@@ -71,6 +71,10 @@ func (suite *DefaultDialerTestSuite) TestHTTPRequest() {
71 71
 	suite.Equal(http.StatusOK, resp.StatusCode)
72 72
 }
73 73
 
74
+func (suite *DefaultDialerTestSuite) TestTCPBufferSize() {
75
+	suite.Equal(network.DefaultBufferSize, suite.d.TCPBufferSize())
76
+}
77
+
74 78
 func TestDefaultDialer(t *testing.T) {
75 79
 	t.Parallel()
76 80
 	suite.Run(t, &DefaultDialerTestSuite{})

+ 1
- 0
network/init.go Ver fichero

@@ -29,4 +29,5 @@ var (
29 29
 type Dialer interface {
30 30
 	Dial(network, address string) (net.Conn, error)
31 31
 	DialContext(ctx context.Context, network, address string) (net.Conn, error)
32
+	TCPBufferSize() int
32 33
 }

+ 4
- 0
network/init_internal_test.go Ver fichero

@@ -22,3 +22,7 @@ func (d *DialerMock) DialContext(ctx context.Context, network, address string) (
22 22
 
23 23
 	return args.Get(0).(net.Conn), args.Error(1)
24 24
 }
25
+
26
+func (d *DialerMock) TCPBufferSize() int {
27
+	return d.Called().Int(0)
28
+}

+ 4
- 0
network/init_test.go Ver fichero

@@ -30,6 +30,10 @@ func (d *DialerMock) DialContext(ctx context.Context, network, address string) (
30 30
 	return args.Get(0).(net.Conn), args.Error(1)
31 31
 }
32 32
 
33
+func (d *DialerMock) TCPBufferSize() int {
34
+	return d.Called().Int(0)
35
+}
36
+
33 37
 type HTTPServerTestSuite struct {
34 38
 	httpServer *httptest.Server
35 39
 }

+ 8
- 2
network/load_balanced_socks5.go Ver fichero

@@ -9,13 +9,18 @@ import (
9 9
 )
10 10
 
11 11
 type loadBalancedSocks5Dialer struct {
12
-	dialers []Dialer
12
+	dialers    []Dialer
13
+	bufferSize int
13 14
 }
14 15
 
15 16
 func (l loadBalancedSocks5Dialer) Dial(network, address string) (net.Conn, error) {
16 17
 	return l.DialContext(context.Background(), network, address)
17 18
 }
18 19
 
20
+func (l loadBalancedSocks5Dialer) TCPBufferSize() int {
21
+	return l.bufferSize
22
+}
23
+
19 24
 func (l loadBalancedSocks5Dialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
20 25
 	length := len(l.dialers)
21 26
 	start := rand.Intn(length)
@@ -45,6 +50,7 @@ func NewLoadBalancedSocks5Dialer(baseDialer Dialer, proxyURLs []*url.URL) (Diale
45 50
 	}
46 51
 
47 52
 	return loadBalancedSocks5Dialer{
48
-		dialers: dialers,
53
+		dialers:    dialers,
54
+		bufferSize: baseDialer.TCPBufferSize(),
49 55
 	}, nil
50 56
 }

+ 1
- 0
network/load_balanced_socks5_test.go Ver fichero

@@ -57,6 +57,7 @@ func (suite *LoadBalancedSocks5TestSuite) TestCannotDial() {
57 57
 	baseDialer.On("DialContext", mock.Anything, "tcp", "127.0.0.2:1080").
58 58
 		Times(network.ProxyDialerOpenThreshold).
59 59
 		Return(&net.TCPConn{}, io.EOF)
60
+	baseDialer.On("TCPBufferSize").Return(network.DefaultBufferSize)
60 61
 
61 62
 	lbDialer, err := network.NewLoadBalancedSocks5Dialer(baseDialer, []*url.URL{
62 63
 		{Scheme: "socks5", User: url.UserPassword("user", "password"), Host: "127.0.0.1:1080"},

+ 4
- 0
network/network.go Ver fichero

@@ -75,6 +75,10 @@ func (n *network) IdleTimeout() time.Duration {
75 75
 	return n.idleTimeout
76 76
 }
77 77
 
78
+func (n *network) TCPBufferSize() int {
79
+	return n.dialer.TCPBufferSize()
80
+}
81
+
78 82
 func (n *network) dnsResolve(protocol, address string) ([]string, error) {
79 83
 	if net.ParseIP(address) != nil {
80 84
 		return []string{address}, nil

+ 19
- 1
network/socks5.go Ver fichero

@@ -2,16 +2,34 @@ package network
2 2
 
3 3
 import (
4 4
 	"fmt"
5
+	"net"
5 6
 	"net/url"
6 7
 
7 8
 	"golang.org/x/net/proxy"
8 9
 )
9 10
 
11
+type socks5Dialer struct {
12
+	proxy.ContextDialer
13
+
14
+	bufferSize int
15
+}
16
+
17
+func (s socks5Dialer) Dial(protocol, address string) (net.Conn, error) {
18
+	return s.ContextDialer.(proxy.Dialer).Dial(protocol, address)
19
+}
20
+
21
+func (s socks5Dialer) TCPBufferSize() int {
22
+	return s.bufferSize
23
+}
24
+
10 25
 func NewSocks5Dialer(baseDialer Dialer, proxyURL *url.URL) (Dialer, error) {
11 26
 	rv, err := proxy.FromURL(proxyURL, baseDialer)
12 27
 	if err != nil {
13 28
 		return nil, fmt.Errorf("cannot initialize socks5 proxy dialer: %w", err)
14 29
 	}
15 30
 
16
-	return rv.(Dialer), nil
31
+	return socks5Dialer{
32
+		ContextDialer: rv.(proxy.ContextDialer),
33
+		bufferSize:    baseDialer.TCPBufferSize(),
34
+	}, nil
17 35
 }

+ 4
- 0
testlib/mtglib_network_mock.go Ver fichero

@@ -33,3 +33,7 @@ func (m *MtglibNetworkMock) MakeHTTPClient(dialFunc func(ctx context.Context,
33 33
 func (m *MtglibNetworkMock) IdleTimeout() time.Duration {
34 34
 	return m.Called().Get(0).(time.Duration)
35 35
 }
36
+
37
+func (m *MtglibNetworkMock) TCPBufferSize() int {
38
+	return m.Called().Int(0)
39
+}

Loading…
Cancelar
Guardar