diff --git a/mtglib/init.go b/mtglib/init.go index e663c62..5f68661 100644 --- a/mtglib/init.go +++ b/mtglib/init.go @@ -27,6 +27,7 @@ type Network interface { DialContext(ctx context.Context, network, address string) (net.Conn, error) MakeHTTPClient(func(ctx context.Context, network, address string) (net.Conn, error)) *http.Client IdleTimeout() time.Duration + TCPBufferSize() int } type AntiReplayCache interface { diff --git a/network/default.go b/network/default.go index 72b3da5..db104f0 100644 --- a/network/default.go +++ b/network/default.go @@ -60,6 +60,10 @@ func (d *defaultDialer) DialContext(ctx context.Context, network, address string return tcpConn, nil } +func (d *defaultDialer) TCPBufferSize() int { + return d.bufferSize +} + func NewDefaultDialer(timeout time.Duration, bufferSize int) (Dialer, error) { switch { case timeout < 0: diff --git a/network/default_test.go b/network/default_test.go index a6b24c7..4136770 100644 --- a/network/default_test.go +++ b/network/default_test.go @@ -71,6 +71,10 @@ func (suite *DefaultDialerTestSuite) TestHTTPRequest() { suite.Equal(http.StatusOK, resp.StatusCode) } +func (suite *DefaultDialerTestSuite) TestTCPBufferSize() { + suite.Equal(network.DefaultBufferSize, suite.d.TCPBufferSize()) +} + func TestDefaultDialer(t *testing.T) { t.Parallel() suite.Run(t, &DefaultDialerTestSuite{}) diff --git a/network/init.go b/network/init.go index 3bab936..bdb4b26 100644 --- a/network/init.go +++ b/network/init.go @@ -29,4 +29,5 @@ var ( type Dialer interface { Dial(network, address string) (net.Conn, error) DialContext(ctx context.Context, network, address string) (net.Conn, error) + TCPBufferSize() int } diff --git a/network/init_internal_test.go b/network/init_internal_test.go index 9335818..b908a01 100644 --- a/network/init_internal_test.go +++ b/network/init_internal_test.go @@ -22,3 +22,7 @@ func (d *DialerMock) DialContext(ctx context.Context, network, address string) ( return args.Get(0).(net.Conn), args.Error(1) } + +func (d *DialerMock) TCPBufferSize() int { + return d.Called().Int(0) +} diff --git a/network/init_test.go b/network/init_test.go index f42d1c3..b75ff11 100644 --- a/network/init_test.go +++ b/network/init_test.go @@ -30,6 +30,10 @@ func (d *DialerMock) DialContext(ctx context.Context, network, address string) ( return args.Get(0).(net.Conn), args.Error(1) } +func (d *DialerMock) TCPBufferSize() int { + return d.Called().Int(0) +} + type HTTPServerTestSuite struct { httpServer *httptest.Server } diff --git a/network/load_balanced_socks5.go b/network/load_balanced_socks5.go index 2a60d69..41331a5 100644 --- a/network/load_balanced_socks5.go +++ b/network/load_balanced_socks5.go @@ -9,13 +9,18 @@ import ( ) type loadBalancedSocks5Dialer struct { - dialers []Dialer + dialers []Dialer + bufferSize int } func (l loadBalancedSocks5Dialer) Dial(network, address string) (net.Conn, error) { return l.DialContext(context.Background(), network, address) } +func (l loadBalancedSocks5Dialer) TCPBufferSize() int { + return l.bufferSize +} + func (l loadBalancedSocks5Dialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) { length := len(l.dialers) start := rand.Intn(length) @@ -45,6 +50,7 @@ func NewLoadBalancedSocks5Dialer(baseDialer Dialer, proxyURLs []*url.URL) (Diale } return loadBalancedSocks5Dialer{ - dialers: dialers, + dialers: dialers, + bufferSize: baseDialer.TCPBufferSize(), }, nil } diff --git a/network/load_balanced_socks5_test.go b/network/load_balanced_socks5_test.go index b6983c8..25ca3b3 100644 --- a/network/load_balanced_socks5_test.go +++ b/network/load_balanced_socks5_test.go @@ -57,6 +57,7 @@ func (suite *LoadBalancedSocks5TestSuite) TestCannotDial() { baseDialer.On("DialContext", mock.Anything, "tcp", "127.0.0.2:1080"). Times(network.ProxyDialerOpenThreshold). Return(&net.TCPConn{}, io.EOF) + baseDialer.On("TCPBufferSize").Return(network.DefaultBufferSize) lbDialer, err := network.NewLoadBalancedSocks5Dialer(baseDialer, []*url.URL{ {Scheme: "socks5", User: url.UserPassword("user", "password"), Host: "127.0.0.1:1080"}, diff --git a/network/network.go b/network/network.go index 585ff3c..f2cb044 100644 --- a/network/network.go +++ b/network/network.go @@ -75,6 +75,10 @@ func (n *network) IdleTimeout() time.Duration { return n.idleTimeout } +func (n *network) TCPBufferSize() int { + return n.dialer.TCPBufferSize() +} + func (n *network) dnsResolve(protocol, address string) ([]string, error) { if net.ParseIP(address) != nil { return []string{address}, nil diff --git a/network/socks5.go b/network/socks5.go index 91c9c61..66c1f23 100644 --- a/network/socks5.go +++ b/network/socks5.go @@ -2,16 +2,34 @@ package network import ( "fmt" + "net" "net/url" "golang.org/x/net/proxy" ) +type socks5Dialer struct { + proxy.ContextDialer + + bufferSize int +} + +func (s socks5Dialer) Dial(protocol, address string) (net.Conn, error) { + return s.ContextDialer.(proxy.Dialer).Dial(protocol, address) +} + +func (s socks5Dialer) TCPBufferSize() int { + return s.bufferSize +} + func NewSocks5Dialer(baseDialer Dialer, proxyURL *url.URL) (Dialer, error) { rv, err := proxy.FromURL(proxyURL, baseDialer) if err != nil { return nil, fmt.Errorf("cannot initialize socks5 proxy dialer: %w", err) } - return rv.(Dialer), nil + return socks5Dialer{ + ContextDialer: rv.(proxy.ContextDialer), + bufferSize: baseDialer.TCPBufferSize(), + }, nil } diff --git a/testlib/mtglib_network_mock.go b/testlib/mtglib_network_mock.go index 23fa8ba..3e89fdd 100644 --- a/testlib/mtglib_network_mock.go +++ b/testlib/mtglib_network_mock.go @@ -33,3 +33,7 @@ func (m *MtglibNetworkMock) MakeHTTPClient(dialFunc func(ctx context.Context, func (m *MtglibNetworkMock) IdleTimeout() time.Duration { return m.Called().Get(0).(time.Duration) } + +func (m *MtglibNetworkMock) TCPBufferSize() int { + return m.Called().Int(0) +}