Use CloseRead and CloseWrites

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
 }

mtglib/internal/faketls/conn_test.go

@@ -15,7 +15,7 @@ import (
 )
 
 type ConnMock struct {
-	testlib.NetConnMock
+	testlib.EssentialsConnMock
 
 	readBuffer  bytes.Buffer
 	writeBuffer bytes.Buffer

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).

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

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)

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)
 	}
 }

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)
 }

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) {

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)
 }

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
 	)