diff --git a/essentials/conns.go b/essentials/conns.go new file mode 100644 index 0000000..d52d761 --- /dev/null +++ b/essentials/conns.go @@ -0,0 +1,17 @@ +package essentials + +import "net" + +type CloseableReader interface { + CloseRead() error +} + +type CloseableWriter interface { + CloseWrite() error +} + +type Conn interface { + net.Conn + CloseableReader + CloseableWriter +} diff --git a/go.mod b/go.mod index 90dbe91..6d9e206 100644 --- a/go.mod +++ b/go.mod @@ -38,9 +38,13 @@ require ( github.com/gotd/ige v0.1.5 // indirect github.com/gotd/xor v0.1.1 // indirect github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect + github.com/patrickmn/go-cache v2.1.0+incompatible // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/prometheus/client_model v0.2.0 // indirect + github.com/txthinking/runnergroup v0.0.0-20210608031112-152c7c4432bf // indirect + github.com/txthinking/socks5 v0.0.0-20211121111206-e03c1217a50b // indirect + github.com/txthinking/x v0.0.0-20210326105829-476fab902fbe // indirect go.uber.org/atomic v1.7.0 // indirect go.uber.org/multierr v1.6.0 // indirect go.uber.org/zap v1.16.0 // indirect diff --git a/go.sum b/go.sum index eaf855f..90dbc59 100644 --- a/go.sum +++ b/go.sum @@ -195,6 +195,8 @@ github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRW github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U= github.com/panjf2000/ants/v2 v2.4.6 h1:drmj9mcygn2gawZ155dRbo+NfXEfAssjZNU1qoIb4gQ= github.com/panjf2000/ants/v2 v2.4.6/go.mod h1:f6F0NZVFsGCp5A7QW/Zj/m92atWwOkY0OIhFxRNFr4A= +github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc= +github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ= github.com/pelletier/go-toml v1.9.4 h1:tjENF6MfZAg8e4ZmZTeWaWiT2vXtsoO6+iuOjFhECwM= github.com/pelletier/go-toml v1.9.4/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= @@ -250,6 +252,12 @@ github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81P github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/txthinking/runnergroup v0.0.0-20210608031112-152c7c4432bf h1:7PflaKRtU4np/epFxRXlFhlzLXZzKFrH5/I4so5Ove0= +github.com/txthinking/runnergroup v0.0.0-20210608031112-152c7c4432bf/go.mod h1:CLUSJbazqETbaR+i0YAhXBICV9TrKH93pziccMhmhpM= +github.com/txthinking/socks5 v0.0.0-20211121111206-e03c1217a50b h1:6J/38A0Xmdnjacfie0Udams7OP/GdoExyTipKwuQWjY= +github.com/txthinking/socks5 v0.0.0-20211121111206-e03c1217a50b/go.mod h1:7NloQcrxaZYKURWph5HLxVDlIwMHJXCPkeWPtpftsIg= +github.com/txthinking/x v0.0.0-20210326105829-476fab902fbe h1:gMWxZxBFRAXqoGkwkYlPX2zvyyKNWJpxOxCrjqJkm5A= +github.com/txthinking/x v0.0.0-20210326105829-476fab902fbe/go.mod h1:WgqbSEmUYSjEV3B1qmee/PpP2NYEz4bL9/+mF1ma+s4= github.com/tylertreat/BoomFilters v0.0.0-20210315201527-1a82519a3e43 h1:QEePdg0ty2r0t1+qwfZmQ4OOl/MB2UXIeJSpIZv56lg= github.com/tylertreat/BoomFilters v0.0.0-20210315201527-1a82519a3e43/go.mod h1:OYRfF6eb5wY9VRFkXJH8FFBi3plw2v+giaIu7P054pM= github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= diff --git a/internal/cli/access.go b/internal/cli/access.go index 23a47ce..2c9a031 100644 --- a/internal/cli/access.go +++ b/internal/cli/access.go @@ -13,6 +13,7 @@ import ( "strings" "sync" + "github.com/9seconds/mtg/v2/essentials" "github.com/9seconds/mtg/v2/internal/config" "github.com/9seconds/mtg/v2/internal/utils" "github.com/9seconds/mtg/v2/mtglib" @@ -106,7 +107,7 @@ func (a *Access) Run(cli *CLI, version string) error { } func (a *Access) getIP(ntw mtglib.Network, protocol string) net.IP { - client := ntw.MakeHTTPClient(func(ctx context.Context, network, address string) (net.Conn, error) { + client := ntw.MakeHTTPClient(func(ctx context.Context, network, address string) (essentials.Conn, error) { return ntw.DialContext(ctx, protocol, address) // nolint: wrapcheck }) diff --git a/internal/testlib/mtglib_network_mock.go b/internal/testlib/mtglib_network_mock.go index 0f58a25..97bd69a 100644 --- a/internal/testlib/mtglib_network_mock.go +++ b/internal/testlib/mtglib_network_mock.go @@ -2,9 +2,9 @@ package testlib import ( "context" - "net" "net/http" + "github.com/9seconds/mtg/v2/essentials" "github.com/stretchr/testify/mock" ) @@ -12,19 +12,19 @@ type MtglibNetworkMock struct { mock.Mock } -func (m *MtglibNetworkMock) Dial(network, address string) (net.Conn, error) { +func (m *MtglibNetworkMock) Dial(network, address string) (essentials.Conn, error) { args := m.Called(network, address) - return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck + return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck } -func (m *MtglibNetworkMock) DialContext(ctx context.Context, network, address string) (net.Conn, error) { +func (m *MtglibNetworkMock) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) { args := m.Called(ctx, network, address) - return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck + return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck } func (m *MtglibNetworkMock) MakeHTTPClient(dialFunc func(ctx context.Context, - network, address string) (net.Conn, error)) *http.Client { + network, address string) (essentials.Conn, error)) *http.Client { return m.Called(dialFunc).Get(0).(*http.Client) } diff --git a/internal/testlib/net_conn_mock.go b/internal/testlib/net_conn_mock.go index e167665..476fe09 100644 --- a/internal/testlib/net_conn_mock.go +++ b/internal/testlib/net_conn_mock.go @@ -7,42 +7,50 @@ import ( "github.com/stretchr/testify/mock" ) -type NetConnMock struct { +type EssentialsConnMock struct { mock.Mock } -func (n *NetConnMock) Read(b []byte) (int, error) { +func (n *EssentialsConnMock) Read(b []byte) (int, error) { args := n.Called(b) return args.Int(0), args.Error(1) } -func (n *NetConnMock) Write(b []byte) (int, error) { +func (n *EssentialsConnMock) Write(b []byte) (int, error) { args := n.Called(b) return args.Int(0), args.Error(1) } -func (n *NetConnMock) Close() error { +func (n *EssentialsConnMock) Close() error { return n.Called().Error(0) // nolint: wrapcheck } -func (n *NetConnMock) LocalAddr() net.Addr { +func (n *EssentialsConnMock) CloseRead() error { + return n.Called().Error(0) // nolint: wrapcheck +} + +func (n *EssentialsConnMock) CloseWrite() error { + return n.Called().Error(0) // nolint: wrapcheck +} + +func (n *EssentialsConnMock) LocalAddr() net.Addr { return n.Called().Get(0).(net.Addr) } -func (n *NetConnMock) RemoteAddr() net.Addr { +func (n *EssentialsConnMock) RemoteAddr() net.Addr { return n.Called().Get(0).(net.Addr) } -func (n *NetConnMock) SetDeadline(t time.Time) error { +func (n *EssentialsConnMock) SetDeadline(t time.Time) error { return n.Called(t).Error(0) // nolint: wrapcheck } -func (n *NetConnMock) SetReadDeadline(t time.Time) error { +func (n *EssentialsConnMock) SetReadDeadline(t time.Time) error { return n.Called(t).Error(0) // nolint: wrapcheck } -func (n *NetConnMock) SetWriteDeadline(t time.Time) error { +func (n *EssentialsConnMock) SetWriteDeadline(t time.Time) error { return n.Called(t).Error(0) // nolint: wrapcheck } diff --git a/mtglib/conns.go b/mtglib/conns.go index 8dfb3a4..129ef52 100644 --- a/mtglib/conns.go +++ b/mtglib/conns.go @@ -4,12 +4,13 @@ import ( "bytes" "context" "io" - "net" "sync" + + "github.com/9seconds/mtg/v2/essentials" ) type connTraffic struct { - net.Conn + essentials.Conn streamID string stream EventStream @@ -37,7 +38,7 @@ func (c connTraffic) Write(b []byte) (int, error) { } type connRewind struct { - net.Conn + essentials.Conn active io.Reader buf bytes.Buffer @@ -58,7 +59,7 @@ func (c *connRewind) Rewind() { c.active = io.MultiReader(&c.buf, c.Conn) } -func newConnRewind(conn net.Conn) *connRewind { +func newConnRewind(conn essentials.Conn) *connRewind { rv := &connRewind{ Conn: conn, } diff --git a/mtglib/conns_internal_test.go b/mtglib/conns_internal_test.go index 8149b0d..ea46d73 100644 --- a/mtglib/conns_internal_test.go +++ b/mtglib/conns_internal_test.go @@ -14,7 +14,7 @@ import ( ) type ConnRewindBaseConn struct { - testlib.NetConnMock + testlib.EssentialsConnMock readBuffer bytes.Buffer } @@ -29,13 +29,13 @@ type ConnTrafficTestSuite struct { suite.Suite eventStreamMock *EventStreamMock - connMock *testlib.NetConnMock + connMock *testlib.EssentialsConnMock conn io.ReadWriter } func (suite *ConnTrafficTestSuite) SetupTest() { suite.eventStreamMock = &EventStreamMock{} - suite.connMock = &testlib.NetConnMock{} + suite.connMock = &testlib.EssentialsConnMock{} suite.conn = connTraffic{ Conn: suite.connMock, streamID: "CONNID", diff --git a/mtglib/init.go b/mtglib/init.go index 13a1be4..2f34ff7 100644 --- a/mtglib/init.go +++ b/mtglib/init.go @@ -23,6 +23,8 @@ import ( "net" "net/http" "time" + + "github.com/9seconds/mtg/v2/essentials" ) var ( @@ -116,16 +118,16 @@ const ( // 3. Doing HTTP requests (for example, for FireHOL ipblocklist). type Network interface { // Dial establishes context-free TCP connections. - Dial(network, address string) (net.Conn, error) + Dial(network, address string) (essentials.Conn, error) // DialContext dials using a context. This is a preferrable // way of establishing TCP connections. - DialContext(ctx context.Context, network, address string) (net.Conn, error) + DialContext(ctx context.Context, network, address string) (essentials.Conn, error) // MakeHTTPClient build an HTTP client with given dial function. If // nothing is provided, then DialContext of this interface is going // to be used. - MakeHTTPClient(func(ctx context.Context, network, address string) (net.Conn, error)) *http.Client + MakeHTTPClient(func(ctx context.Context, network, address string) (essentials.Conn, error)) *http.Client } // AntiReplayCache is an interface that is used to detect replay attacks diff --git a/mtglib/internal/faketls/conn.go b/mtglib/internal/faketls/conn.go index d7e7837..74c8020 100644 --- a/mtglib/internal/faketls/conn.go +++ b/mtglib/internal/faketls/conn.go @@ -4,13 +4,13 @@ import ( "bytes" "fmt" "math/rand" - "net" + "github.com/9seconds/mtg/v2/essentials" "github.com/9seconds/mtg/v2/mtglib/internal/faketls/record" ) type Conn struct { - net.Conn + essentials.Conn readBuffer bytes.Buffer } diff --git a/mtglib/internal/faketls/conn_test.go b/mtglib/internal/faketls/conn_test.go index affcec0..e7f311a 100644 --- a/mtglib/internal/faketls/conn_test.go +++ b/mtglib/internal/faketls/conn_test.go @@ -15,7 +15,7 @@ import ( ) type ConnMock struct { - testlib.NetConnMock + testlib.EssentialsConnMock readBuffer bytes.Buffer writeBuffer bytes.Buffer diff --git a/mtglib/internal/obfuscated2/client_handshake_test.go b/mtglib/internal/obfuscated2/client_handshake_test.go index 6d4d8c2..c310e96 100644 --- a/mtglib/internal/obfuscated2/client_handshake_test.go +++ b/mtglib/internal/obfuscated2/client_handshake_test.go @@ -42,7 +42,7 @@ func (suite *ClientHandshakeTestSuite) TestOk() { writeData := make([]byte, len(snapshot.Encrypted.Text.data)) readData := make([]byte, len(snapshot.Decrypted.Text.data)) - connMock := &testlib.NetConnMock{} + connMock := &testlib.EssentialsConnMock{} connMock.On("Read", mock.Anything). Once(). Return(len(snapshot.Decrypted.Text.data), nil). diff --git a/mtglib/internal/obfuscated2/conn.go b/mtglib/internal/obfuscated2/conn.go index 511ba62..b6ecbf4 100644 --- a/mtglib/internal/obfuscated2/conn.go +++ b/mtglib/internal/obfuscated2/conn.go @@ -2,11 +2,12 @@ package obfuscated2 import ( "crypto/cipher" - "net" + + "github.com/9seconds/mtg/v2/essentials" ) type Conn struct { - net.Conn + essentials.Conn Encryptor cipher.Stream Decryptor cipher.Stream diff --git a/mtglib/internal/obfuscated2/server_handshake_test.go b/mtglib/internal/obfuscated2/server_handshake_test.go index 418af0c..89719bb 100644 --- a/mtglib/internal/obfuscated2/server_handshake_test.go +++ b/mtglib/internal/obfuscated2/server_handshake_test.go @@ -16,7 +16,7 @@ import ( type ServerHandshakeTestSuite struct { suite.Suite - connMock *testlib.NetConnMock + connMock *testlib.EssentialsConnMock proxyConn obfuscated2.Conn encryptor cipher.Stream decryptor cipher.Stream @@ -24,7 +24,7 @@ type ServerHandshakeTestSuite struct { func (suite *ServerHandshakeTestSuite) SetupTest() { buf := &bytes.Buffer{} - suite.connMock = &testlib.NetConnMock{} + suite.connMock = &testlib.EssentialsConnMock{} encryptor, decryptor, err := obfuscated2.ServerHandshake(buf) suite.NoError(err) diff --git a/mtglib/internal/relay/relay.go b/mtglib/internal/relay/relay.go index d4abe87..bf3bc03 100644 --- a/mtglib/internal/relay/relay.go +++ b/mtglib/internal/relay/relay.go @@ -2,11 +2,14 @@ package relay import ( "context" - "net" + "errors" + "io" "sync" + + "github.com/9seconds/mtg/v2/essentials" ) -func Relay(ctx context.Context, log Logger, telegramConn, clientConn net.Conn) { +func Relay(ctx context.Context, log Logger, telegramConn, clientConn essentials.Conn) { defer telegramConn.Close() defer clientConn.Close() @@ -29,14 +32,25 @@ func Relay(ctx context.Context, log Logger, telegramConn, clientConn net.Conn) { wg.Wait() } -func pump(log Logger, src, dst net.Conn, wg *sync.WaitGroup, direction string) { - defer wg.Done() - +func pump(log Logger, src, dst essentials.Conn, wg *sync.WaitGroup, direction string) { syncer := acquireSyncPair(src, dst) - defer releaseSyncPair(syncer) - defer syncer.Flush() - if n, err := syncer.Sync(); err != nil { - log.Printf("cannot pump %s (written %d bytes): %v", direction, n, err) + defer func() { + syncer.Flush() + releaseSyncPair(syncer) + src.CloseRead() + dst.CloseWrite() + wg.Done() + }() + + n, err := syncer.Sync() + + switch { + case err == nil: + log.Printf("%s has been finished", direction) + case errors.Is(err, io.EOF): + log.Printf("%s has been finished because of EOF. Written %d bytes", direction, n) + default: + log.Printf("%s has been finished (written %d bytes): %v", direction, n, err) } } diff --git a/mtglib/internal/relay/relay_test.go b/mtglib/internal/relay/relay_test.go index 538622a..368469f 100644 --- a/mtglib/internal/relay/relay_test.go +++ b/mtglib/internal/relay/relay_test.go @@ -17,8 +17,8 @@ type RelayTestSuite struct { loggerMock relay.Logger ctx context.Context ctxCancel context.CancelFunc - telegramConnMock *testlib.NetConnMock - clientConnMock *testlib.NetConnMock + telegramConnMock *testlib.EssentialsConnMock + clientConnMock *testlib.EssentialsConnMock } func (suite *RelayTestSuite) SetupTest() { @@ -26,8 +26,8 @@ func (suite *RelayTestSuite) SetupTest() { suite.ctx = ctx suite.ctxCancel = cancel suite.loggerMock = &loggerMock{} - suite.telegramConnMock = &testlib.NetConnMock{} - suite.clientConnMock = &testlib.NetConnMock{} + suite.telegramConnMock = &testlib.EssentialsConnMock{} + suite.clientConnMock = &testlib.EssentialsConnMock{} } func (suite *RelayTestSuite) TearDownTest() { @@ -38,12 +38,18 @@ func (suite *RelayTestSuite) TearDownTest() { func (suite *RelayTestSuite) TestExit() { suite.telegramConnMock.On("Close").Return(nil) + suite.telegramConnMock.On("CloseRead").Return(nil).Once() + suite.telegramConnMock.On("CloseWrite").Return(nil).Once() suite.telegramConnMock.On("Read", mock.Anything).Return(10, io.EOF).Once() suite.telegramConnMock.On("Write", mock.Anything).Return(10, io.EOF).Maybe() + suite.telegramConnMock.On("SetReadDeadline", mock.Anything).Return(nil).Maybe() suite.clientConnMock.On("Read", mock.Anything).Return(0, io.EOF).Once() suite.clientConnMock.On("Write", mock.Anything).Return(10, io.EOF).Maybe() suite.clientConnMock.On("Close").Return(nil) + suite.clientConnMock.On("CloseRead").Return(nil).Once() + suite.clientConnMock.On("CloseWrite").Return(nil).Once() + suite.clientConnMock.On("SetReadDeadline", mock.Anything).Return(nil).Maybe() relay.Relay(suite.ctx, suite.loggerMock, suite.telegramConnMock, suite.clientConnMock) } diff --git a/mtglib/internal/relay/sync_pair.go b/mtglib/internal/relay/sync_pair.go index 68c1e62..72293e7 100644 --- a/mtglib/internal/relay/sync_pair.go +++ b/mtglib/internal/relay/sync_pair.go @@ -48,7 +48,7 @@ func (s *syncPair) Flush() error { s.mutex.Lock() defer s.mutex.Unlock() - return s.writer.Flush() + return s.writer.Flush() // nolint: wrapcheck } func (s *syncPair) readBlocking(p []byte, blocking bool) (int, error) { diff --git a/mtglib/internal/telegram/init.go b/mtglib/internal/telegram/init.go index 469478f..448120b 100644 --- a/mtglib/internal/telegram/init.go +++ b/mtglib/internal/telegram/init.go @@ -2,7 +2,8 @@ package telegram import ( "context" - "net" + + "github.com/9seconds/mtg/v2/essentials" ) type preferIP uint8 @@ -82,5 +83,5 @@ var ( ) type Dialer interface { - DialContext(ctx context.Context, network, address string) (net.Conn, error) + DialContext(ctx context.Context, network, address string) (essentials.Conn, error) } diff --git a/mtglib/internal/telegram/telegram.go b/mtglib/internal/telegram/telegram.go index 43d4867..395a9b7 100644 --- a/mtglib/internal/telegram/telegram.go +++ b/mtglib/internal/telegram/telegram.go @@ -3,8 +3,9 @@ package telegram import ( "context" "fmt" - "net" "strings" + + "github.com/9seconds/mtg/v2/essentials" ) type Telegram struct { @@ -13,7 +14,7 @@ type Telegram struct { pool addressPool } -func (t Telegram) Dial(ctx context.Context, dc int) (net.Conn, error) { +func (t Telegram) Dial(ctx context.Context, dc int) (essentials.Conn, error) { var addresses []tgAddr switch t.preferIP { @@ -28,7 +29,7 @@ func (t Telegram) Dial(ctx context.Context, dc int) (net.Conn, error) { } var ( - conn net.Conn + conn essentials.Conn err error ) diff --git a/mtglib/proxy.go b/mtglib/proxy.go index ebf978c..454ebe9 100644 --- a/mtglib/proxy.go +++ b/mtglib/proxy.go @@ -9,6 +9,7 @@ import ( "sync" "time" + "github.com/9seconds/mtg/v2/essentials" "github.com/9seconds/mtg/v2/mtglib/internal/faketls" "github.com/9seconds/mtg/v2/mtglib/internal/faketls/record" "github.com/9seconds/mtg/v2/mtglib/internal/obfuscated2" @@ -44,7 +45,7 @@ func (p *Proxy) DomainFrontingAddress() string { // ServeConn serves a connection. We do not check IP blocklist and // concurrency limit here. -func (p *Proxy) ServeConn(conn net.Conn) { +func (p *Proxy) ServeConn(conn essentials.Conn) { p.streamWaitGroup.Add(1) defer p.streamWaitGroup.Done() @@ -299,7 +300,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { pool, err := ants.NewPoolWithFunc(opts.getConcurrency(), func(arg interface{}) { - proxy.ServeConn(arg.(net.Conn)) + proxy.ServeConn(arg.(essentials.Conn)) }, ants.WithLogger(opts.getLogger("ants")), ants.WithNonblocking(true)) diff --git a/mtglib/stream_context.go b/mtglib/stream_context.go index 7ee8071..e1f2319 100644 --- a/mtglib/stream_context.go +++ b/mtglib/stream_context.go @@ -6,13 +6,15 @@ import ( "encoding/base64" "net" "time" + + "github.com/9seconds/mtg/v2/essentials" ) type streamContext struct { ctx context.Context ctxCancel context.CancelFunc - clientConn net.Conn - telegramConn net.Conn + clientConn essentials.Conn + telegramConn essentials.Conn streamID string dc int logger Logger @@ -50,7 +52,7 @@ func (s *streamContext) ClientIP() net.IP { return s.clientConn.RemoteAddr().(*net.TCPAddr).IP } -func newStreamContext(ctx context.Context, logger Logger, clientConn net.Conn) *streamContext { +func newStreamContext(ctx context.Context, logger Logger, clientConn essentials.Conn) *streamContext { connIDBytes := make([]byte, ConnectionIDBytesLength) if _, err := rand.Read(connIDBytes); err != nil { diff --git a/mtglib/stream_context_internal_test.go b/mtglib/stream_context_internal_test.go index 17f90ea..52b5d4a 100644 --- a/mtglib/stream_context_internal_test.go +++ b/mtglib/stream_context_internal_test.go @@ -12,7 +12,7 @@ import ( type StreamContextTestSuite struct { suite.Suite - connMock *testlib.NetConnMock + connMock *testlib.EssentialsConnMock logger NoopLogger ctx *streamContext ctxCancel context.CancelFunc @@ -27,7 +27,7 @@ func (suite *StreamContextTestSuite) SetupTest() { ctx = context.WithValue(ctx, "key", "value") // nolint: golint, revive, staticcheck suite.ctxCancel = cancel - suite.connMock = &testlib.NetConnMock{} + suite.connMock = &testlib.EssentialsConnMock{} addr := &net.TCPAddr{ IP: net.ParseIP("10.0.0.10"), @@ -73,7 +73,7 @@ func (suite *StreamContextTestSuite) TestClientIP() { func (suite *StreamContextTestSuite) TestClose() { suite.connMock.On("Close").Once().Return(nil) - tgConnMock := &testlib.NetConnMock{} + tgConnMock := &testlib.EssentialsConnMock{} tgConnMock.On("Close").Once().Return(nil) suite.ctx.telegramConn = tgConnMock diff --git a/network/circuit_breaker.go b/network/circuit_breaker.go index 4c74a9b..91e1317 100644 --- a/network/circuit_breaker.go +++ b/network/circuit_breaker.go @@ -2,9 +2,10 @@ package network import ( "context" - "net" "sync/atomic" "time" + + "github.com/9seconds/mtg/v2/essentials" ) const ( @@ -30,12 +31,12 @@ type circuitBreakerDialer struct { resetFailuresTimeout time.Duration } -func (c *circuitBreakerDialer) Dial(network, address string) (net.Conn, error) { +func (c *circuitBreakerDialer) Dial(network, address string) (essentials.Conn, error) { return c.DialContext(context.Background(), network, address) } func (c *circuitBreakerDialer) DialContext(ctx context.Context, - network, address string) (net.Conn, error) { + network, address string) (essentials.Conn, error) { switch atomic.LoadUint32(&c.state) { case circuitBreakerStateClosed: return c.doClosed(ctx, network, address) @@ -47,7 +48,7 @@ func (c *circuitBreakerDialer) DialContext(ctx context.Context, } func (c *circuitBreakerDialer) doClosed(ctx context.Context, - network, address string) (net.Conn, error) { + network, address string) (essentials.Conn, error) { conn, err := c.Dialer.DialContext(ctx, network, address) select { @@ -78,7 +79,8 @@ func (c *circuitBreakerDialer) doClosed(ctx context.Context, return conn, err // nolint: wrapcheck } -func (c *circuitBreakerDialer) doHalfOpened(ctx context.Context, network, address string) (net.Conn, error) { +func (c *circuitBreakerDialer) doHalfOpened(ctx context.Context, + network, address string) (essentials.Conn, error) { if !atomic.CompareAndSwapUint32(&c.halfOpenAttempts, 0, 1) { return nil, ErrCircuitBreakerOpened } diff --git a/network/circuit_breaker_internal_test.go b/network/circuit_breaker_internal_test.go index a26f594..d300d68 100644 --- a/network/circuit_breaker_internal_test.go +++ b/network/circuit_breaker_internal_test.go @@ -21,7 +21,7 @@ type CircuitBreakerTestSuite struct { mutex sync.Mutex ctx context.Context ctxCancel context.CancelFunc - connMock *testlib.NetConnMock + connMock *testlib.EssentialsConnMock baseDialerMock *DialerMock } @@ -29,7 +29,7 @@ func (suite *CircuitBreakerTestSuite) SetupTest() { suite.mutex = sync.Mutex{} suite.ctx, suite.ctxCancel = context.WithCancel(context.Background()) suite.baseDialerMock = &DialerMock{} - suite.connMock = &testlib.NetConnMock{} + suite.connMock = &testlib.EssentialsConnMock{} suite.d = newCircuitBreakerDialer(suite.baseDialerMock, 3, 100*time.Millisecond, 50*time.Millisecond) } diff --git a/network/default.go b/network/default.go index 7500bd8..50855d9 100644 --- a/network/default.go +++ b/network/default.go @@ -5,17 +5,19 @@ import ( "fmt" "net" "time" + + "github.com/9seconds/mtg/v2/essentials" ) type defaultDialer struct { net.Dialer } -func (d *defaultDialer) Dial(network, address string) (net.Conn, error) { +func (d *defaultDialer) Dial(network, address string) (essentials.Conn, error) { return d.DialContext(context.Background(), network, address) } -func (d *defaultDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) { +func (d *defaultDialer) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) { switch network { case "tcp", "tcp4", "tcp6": // nolint: goconst default: @@ -34,7 +36,7 @@ func (d *defaultDialer) DialContext(ctx context.Context, network, address string return nil, fmt.Errorf("cannot set socket options: %w", err) } - return conn, nil + return conn.(essentials.Conn), nil } // NewDefaultDialer build a new dialer which dials bypassing proxies diff --git a/network/init.go b/network/init.go index 45884f5..a1426b3 100644 --- a/network/init.go +++ b/network/init.go @@ -20,8 +20,9 @@ package network import ( "context" "errors" - "net" "time" + + "github.com/9seconds/mtg/v2/essentials" ) const ( @@ -95,6 +96,6 @@ var ( // Dialer defines an interface which is required to bootstrap a network // instance from. type Dialer interface { - Dial(network, address string) (net.Conn, error) - DialContext(ctx context.Context, network, address string) (net.Conn, error) + Dial(network, address string) (essentials.Conn, error) + DialContext(ctx context.Context, network, address string) (essentials.Conn, error) } diff --git a/network/init_internal_test.go b/network/init_internal_test.go index 365a4e5..6e63532 100644 --- a/network/init_internal_test.go +++ b/network/init_internal_test.go @@ -2,8 +2,8 @@ package network import ( "context" - "net" + "github.com/9seconds/mtg/v2/essentials" "github.com/stretchr/testify/mock" ) @@ -11,14 +11,14 @@ type DialerMock struct { mock.Mock } -func (d *DialerMock) Dial(network, address string) (net.Conn, error) { +func (d *DialerMock) Dial(network, address string) (essentials.Conn, error) { args := d.Called(network, address) - return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck + return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck } -func (d *DialerMock) DialContext(ctx context.Context, network, address string) (net.Conn, error) { +func (d *DialerMock) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) { args := d.Called(ctx, network, address) - return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck + return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck } diff --git a/network/init_test.go b/network/init_test.go index f1692e0..6e79a48 100644 --- a/network/init_test.go +++ b/network/init_test.go @@ -8,6 +8,7 @@ import ( "net/url" "strings" + "github.com/9seconds/mtg/v2/essentials" "github.com/9seconds/mtg/v2/network" socks5 "github.com/armon/go-socks5" "github.com/mccutchen/go-httpbin/httpbin" @@ -18,16 +19,16 @@ type DialerMock struct { mock.Mock } -func (d *DialerMock) Dial(network, address string) (net.Conn, error) { +func (d *DialerMock) Dial(network, address string) (essentials.Conn, error) { args := d.Called(network, address) - return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck + return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck } -func (d *DialerMock) DialContext(ctx context.Context, network, address string) (net.Conn, error) { +func (d *DialerMock) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) { args := d.Called(ctx, network, address) - return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck + return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck } type HTTPServerTestSuite struct { @@ -53,7 +54,9 @@ func (suite *HTTPServerTestSuite) MakeURL(path string) string { func (suite *HTTPServerTestSuite) MakeHTTPClient(dialer network.Dialer) *http.Client { return &http.Client{ Transport: &http.Transport{ - DialContext: dialer.DialContext, + DialContext: func(ctx context.Context, network, address string) (net.Conn, error) { + return dialer.DialContext(ctx, network, address) // nolint: wrapcheck + }, }, } } diff --git a/network/load_balanced_socks5.go b/network/load_balanced_socks5.go index a52004d..37cbf28 100644 --- a/network/load_balanced_socks5.go +++ b/network/load_balanced_socks5.go @@ -4,19 +4,20 @@ import ( "context" "fmt" "math/rand" - "net" "net/url" + + "github.com/9seconds/mtg/v2/essentials" ) type loadBalancedSocks5Dialer struct { dialers []Dialer } -func (l loadBalancedSocks5Dialer) Dial(network, address string) (net.Conn, error) { +func (l loadBalancedSocks5Dialer) Dial(network, address string) (essentials.Conn, error) { return l.DialContext(context.Background(), network, address) } -func (l loadBalancedSocks5Dialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) { +func (l loadBalancedSocks5Dialer) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) { length := len(l.dialers) start := rand.Intn(length) moved := false diff --git a/network/network.go b/network/network.go index cac18fb..e42508f 100644 --- a/network/network.go +++ b/network/network.go @@ -9,6 +9,7 @@ import ( "sync" "time" + "github.com/9seconds/mtg/v2/essentials" "github.com/9seconds/mtg/v2/mtglib" ) @@ -30,11 +31,11 @@ type network struct { dns *dnsResolver } -func (n *network) Dial(protocol, address string) (net.Conn, error) { +func (n *network) Dial(protocol, address string) (essentials.Conn, error) { return n.DialContext(context.Background(), protocol, address) } -func (n *network) DialContext(ctx context.Context, protocol, address string) (net.Conn, error) { +func (n *network) DialContext(ctx context.Context, protocol, address string) (essentials.Conn, error) { host, port, _ := net.SplitHostPort(address) ips, err := n.dnsResolve(protocol, host) @@ -46,7 +47,8 @@ func (n *network) DialContext(ctx context.Context, protocol, address string) (ne ips[i], ips[j] = ips[j], ips[i] }) - var conn net.Conn + var conn essentials.Conn + for _, v := range ips { conn, err = n.dialer.DialContext(ctx, protocol, net.JoinHostPort(v, port)) @@ -59,7 +61,7 @@ func (n *network) DialContext(ctx context.Context, protocol, address string) (ne } func (n *network) MakeHTTPClient(dialFunc func(ctx context.Context, - network, address string) (net.Conn, error)) *http.Client { + network, address string) (essentials.Conn, error)) *http.Client { if dialFunc == nil { dialFunc = n.DialContext } @@ -144,13 +146,15 @@ func NewNetwork(dialer Dialer, func makeHTTPClient(userAgent string, timeout time.Duration, - dialFunc func(ctx context.Context, network, address string) (net.Conn, error)) *http.Client { + dialFunc func(ctx context.Context, network, address string) (essentials.Conn, error)) *http.Client { return &http.Client{ Timeout: timeout, Transport: networkHTTPTransport{ userAgent: userAgent, next: &http.Transport{ - DialContext: dialFunc, + DialContext: func(ctx context.Context, network, address string) (net.Conn, error) { + return dialFunc(ctx, network, address) + }, }, }, } diff --git a/network/socks5.go b/network/socks5.go index 43aac4a..459f64b 100644 --- a/network/socks5.go +++ b/network/socks5.go @@ -1,21 +1,162 @@ package network import ( + "context" "fmt" + "io" + "net" "net/url" - "golang.org/x/net/proxy" + "github.com/9seconds/mtg/v2/essentials" + "github.com/txthinking/socks5" ) +type socks5Dialer struct { + Dialer + + username []byte + password []byte + proxyAddress string +} + +func (s socks5Dialer) Dial(network, address string) (essentials.Conn, error) { + return s.DialContext(context.Background(), network, address) +} + +func (s socks5Dialer) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) { + switch network { + case "tcp", "tcp4", "tcp6": + default: + return nil, fmt.Errorf("%s network type is not supported", network) + } + + conn, err := s.Dialer.DialContext(ctx, network, s.proxyAddress) + if err != nil { + return nil, fmt.Errorf("cannot dial to the proxy: %w", err) + } + + if err := s.handshake(conn); err != nil { + conn.Close() + + return nil, fmt.Errorf("cannot perform a handshake: %w", err) + } + + if err := s.connect(conn, address); err != nil { + conn.Close() + + return nil, fmt.Errorf("cannot connect to a destination host %s: %w", address, err) + } + + return conn, nil +} + +func (s socks5Dialer) handshake(conn io.ReadWriter) error { + authMethod := socks5.MethodUsernamePassword + if len(s.username)+len(s.password) == 0 { + authMethod = socks5.MethodNone + } + + if err := s.handshakeNegotiation(conn, authMethod); err != nil { + return fmt.Errorf("cannot perform negotiation: %w", err) + } + + if authMethod == socks5.MethodNone { + return nil + } + + if err := s.handshakeAuth(conn); err != nil { + return fmt.Errorf("cannot authenticate: %w", err) + } + + return nil +} + +func (s socks5Dialer) handshakeNegotiation(conn io.ReadWriter, authMethod byte) error { + request := socks5.NewNegotiationRequest([]byte{authMethod}) + if _, err := request.WriteTo(conn); err != nil { + return fmt.Errorf("cannot send request: %w", err) + } + + response, err := socks5.NewNegotiationReplyFrom(conn) + if err != nil { + return fmt.Errorf("cannot read response: %w", err) + } + + if response.Method != authMethod { + return fmt.Errorf("%v is unsupported auth method", authMethod) + } + + return nil +} + +func (s socks5Dialer) handshakeAuth(conn io.ReadWriter) error { + request := socks5.NewUserPassNegotiationRequest(s.username, s.password) + + if _, err := request.WriteTo(conn); err != nil { + return fmt.Errorf("cannot send a request: %w", err) + } + + response, err := socks5.NewUserPassNegotiationReplyFrom(conn) + if err != nil { + return fmt.Errorf("cannot read a response: %w", err) + } + + if response.Status != socks5.UserPassStatusSuccess { + return fmt.Errorf("authenticate has failed: %v", response.Status) + } + + return nil +} + +func (s socks5Dialer) connect(conn io.ReadWriter, address string) error { + addrType, host, port, err := socks5.ParseAddress(address) + if err != nil { + return fmt.Errorf("cannot parse address: %w", err) + } + + if addrType == socks5.ATYPDomain { + host = host[1:] + } + + request := socks5.NewRequest(socks5.CmdConnect, addrType, host, port) + + if _, err := request.WriteTo(conn); err != nil { + return fmt.Errorf("cannot send a request: %w", err) + } + + response, err := socks5.NewReplyFrom(conn) + if err != nil { + return fmt.Errorf("cannot read a response: %w", err) + } + + if response.Rep != socks5.RepSuccess { + return fmt.Errorf("unsuccessful request: %v", response.Rep) + } + + return nil +} + // NewSocks5Dialer build a new dialer from a given one (so, in theory // you can chain here). Proxy parameters are passed with URI in a form of: // // socks5://[user:[password]]@host:port 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) + if _, _, err := net.SplitHostPort(proxyURL.Host); err != nil { + return nil, fmt.Errorf("incorrect url %s", proxyURL.Redacted()) } - return rv.(Dialer), nil + dialer := socks5Dialer{ + Dialer: baseDialer, + proxyAddress: proxyURL.Host, + } + + if proxyURL.User != nil { + password, isSet := proxyURL.User.Password() + if isSet { + dialer.username = []byte(proxyURL.User.Username()) + dialer.password = []byte(password) + } + } + + return dialer, nil } diff --git a/network/socks5_test.go b/network/socks5_test.go index 3460f4d..5cb56c3 100644 --- a/network/socks5_test.go +++ b/network/socks5_test.go @@ -55,7 +55,7 @@ func (suite *Socks5TestSuite) TestRequestOk() { suite.Equal(http.StatusOK, resp.StatusCode) } -func TestSocks5TestSuite(t *testing.T) { +func TestSocks5(t *testing.T) { t.Parallel() suite.Run(t, &Socks5TestSuite{}) }