Merge pull request #230 from 9seconds/simplify-sockopts

Simplify sockopts

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

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",

mtglib/init.go

@@ -23,6 +23,8 @@ import (
 	"net"
 	"net/http"
 	"time"
+
+	"github.com/9seconds/mtg/v2/essentials"
 )
 
 var (
@@ -61,6 +63,8 @@ const (
 	DefaultConcurrency = 4096
 
 	// DefaultBufferSize is a default size of a copy buffer.
+	//
+	// Deprecated: this setting no longer makes any effect.
 	DefaultBufferSize = 16 * 1024 // 16 kib
 
 	// DefaultDomainFrontingPort is a default port (HTTPS) to connect to in
@@ -114,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

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/conn.go

@@ -1,19 +0,0 @@
-package relay
-
-import (
-	"fmt"
-	"net"
-	"time"
-)
-
-type conn struct {
-	net.Conn
-}
-
-func (c conn) Read(p []byte) (int, error) {
-	if err := c.SetReadDeadline(time.Now().Add(getTimeout())); err != nil {
-		return 0, fmt.Errorf("cannot set read deadline: %w", err)
-	}
-
-	return c.Conn.Read(p) // nolint: wrapcheck
-}

mtglib/internal/relay/init.go

@@ -3,10 +3,9 @@ package relay
 import "time"
 
 const (
-	ConnectionTimeToLiveMin = 2 * time.Minute
-	ConnectionTimeToLiveMax = 10 * time.Minute
-	TimeoutMin              = 20 * time.Second
-	TimeoutMax              = time.Minute
+	copyBufferSize   = 64 * 1024
+	writerBufferSize = 128 * 1024
+	readTimeout      = 10 * time.Millisecond
 )
 
 type Logger interface {

mtglib/internal/relay/pools.go

@@ -1,32 +1,31 @@
 package relay
 
-import "sync"
+import (
+	"bufio"
+	"io"
+	"net"
+	"sync"
+)
 
-type eastWest struct {
-	east []byte
-	west []byte
-}
-
-var eastWestPool = sync.Pool{
+var syncPairPool = sync.Pool{
 	New: func() interface{} {
-		return &eastWest{}
+		return &syncPair{
+			writer:  bufio.NewWriterSize(nil, writerBufferSize),
+			copyBuf: make([]byte, copyBufferSize),
+		}
 	},
 }
 
-func acquireEastWest(bufferSize int) *eastWest {
-	wanted := eastWestPool.Get().(*eastWest) // nolint: forcetypeassert
+func acquireSyncPair(reader net.Conn, writer io.Writer) *syncPair {
+	sp := syncPairPool.Get().(*syncPair) // nolint: forcetypeassert
+	sp.writer.Reset(writer)
+	sp.reader = reader
 
-	if len(wanted.east) != bufferSize {
-		wanted.east = make([]byte, bufferSize)
-	}
-
-	if len(wanted.west) != bufferSize {
-		wanted.west = make([]byte, bufferSize)
-	}
-
-	return wanted
+	return sp
 }
 
-func releaseEastWest(ew *eastWest) {
-	eastWestPool.Put(ew)
+func releaseSyncPair(sp *syncPair) {
+	sp.writer.Reset(nil)
+	sp.reader = nil
+	syncPairPool.Put(sp)
 }

mtglib/internal/relay/relay.go

@@ -2,17 +2,18 @@ package relay
 
 import (
 	"context"
+	"errors"
 	"io"
-	"net"
 	"sync"
+
+	"github.com/9seconds/mtg/v2/essentials"
 )
 
-func Relay(ctx context.Context, log Logger, bufferSize int,
-	telegramConn net.Conn, clientConn io.ReadWriteCloser) {
+func Relay(ctx context.Context, log Logger, telegramConn, clientConn essentials.Conn) {
 	defer telegramConn.Close()
 	defer clientConn.Close()
 
-	ctx, cancel := context.WithTimeout(ctx, getConnectionTimeToLive())
+	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	go func() {
@@ -21,30 +22,35 @@ func Relay(ctx context.Context, log Logger, bufferSize int,
 		clientConn.Close()
 	}()
 
-	buffers := acquireEastWest(bufferSize)
-	defer releaseEastWest(buffers)
-
-	telegramConn = conn{
-		Conn: telegramConn,
-	}
-
 	wg := &sync.WaitGroup{}
 	wg.Add(2) // nolint: gomnd
 
-	go pump(log, telegramConn, clientConn, wg, buffers.east, "east -> west")
+	go pump(log, telegramConn, clientConn, wg, "client -> telegram")
 
-	pump(log, clientConn, telegramConn, wg, buffers.west, "west -> east")
+	pump(log, clientConn, telegramConn, wg, "telegram -> client")
 
 	wg.Wait()
 }
 
-func pump(log Logger, src io.ReadCloser, dst io.WriteCloser, wg *sync.WaitGroup,
-	buf []byte, direction string) {
-	defer wg.Done()
-	defer src.Close()
-	defer dst.Close()
+func pump(log Logger, src, dst essentials.Conn, wg *sync.WaitGroup, direction string) {
+	syncer := acquireSyncPair(src, dst)
 
-	if n, err := io.CopyBuffer(dst, src, buf); err != nil {
-		log.Printf("cannot pump %s (written %d bytes): %w", direction, n, err)
+	defer func() {
+		syncer.Flush()
+		releaseSyncPair(syncer)
+		src.CloseRead()  // nolint: errcheck
+		dst.CloseWrite() // nolint: errcheck
+		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() {
@@ -37,17 +37,21 @@ func (suite *RelayTestSuite) TearDownTest() {
 }
 
 func (suite *RelayTestSuite) TestExit() {
-	suite.telegramConnMock.On("SetReadDeadline", mock.Anything).Return(nil)
 	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, 1024,
-		suite.telegramConnMock, suite.clientConnMock)
+	relay.Relay(suite.ctx, suite.loggerMock, suite.telegramConnMock, suite.clientConnMock)
 }
 
 func TestRelay(t *testing.T) {

mtglib/internal/relay/sync_pair.go

@@ -0,0 +1,77 @@
+package relay
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"os"
+	"sync"
+	"time"
+)
+
+type syncPair struct {
+	writer  *bufio.Writer
+	copyBuf []byte
+
+	mutex  sync.Mutex
+	reader net.Conn
+}
+
+func (s *syncPair) Sync() (int64, error) {
+	return io.CopyBuffer(s, s, s.copyBuf) // nolint: wrapcheck
+}
+
+func (s *syncPair) Read(p []byte) (int, error) {
+	n, err := s.readBlocking(p, false)
+
+	// nothing has been delivered for readTimeout time. Let's flush.
+	if errors.Is(err, os.ErrDeadlineExceeded) {
+		if err := s.Flush(); err != nil {
+			return 0, fmt.Errorf("cannot flush writer hand-side: %w", err)
+		}
+
+		return s.readBlocking(p, true)
+	}
+
+	return n, err
+}
+
+func (s *syncPair) Write(p []byte) (int, error) {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	n, err := s.writer.Write(p)
+
+	// optimization for a case when we have a small package and want to avoid a
+	// delay in readTimeout. In that case, we assume that peer has finished to
+	// sent a data it wants to send so we can flush without waiting for anything
+	// else.
+	if err == nil && n < copyBufferSize {
+		err = s.writer.Flush()
+	}
+
+	return n, err // nolint: wrapcheck
+}
+
+func (s *syncPair) Flush() error {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	return s.writer.Flush() // nolint: wrapcheck
+}
+
+func (s *syncPair) readBlocking(p []byte, blocking bool) (int, error) {
+	var deadline time.Time
+
+	if !blocking {
+		deadline = time.Now().Add(readTimeout)
+	}
+
+	if err := s.reader.SetReadDeadline(deadline); err != nil {
+		return 0, fmt.Errorf("cannot set read deadline: %w", err)
+	}
+
+	return s.reader.Read(p) // nolint: wrapcheck
+}

mtglib/internal/relay/timeouts.go

@@ -1,22 +0,0 @@
-package relay
-
-import (
-	"math/rand"
-	"time"
-)
-
-func getConnectionTimeToLive() time.Duration {
-	return getTime(ConnectionTimeToLiveMin, ConnectionTimeToLiveMax)
-}
-
-func getTimeout() time.Duration {
-	return getTime(TimeoutMin, TimeoutMax)
-}
-
-func getTime(minDuration, maxDuration time.Duration) time.Duration {
-	minDurationInSeconds := int(minDuration.Seconds())
-	maxDurationInSeconds := int(maxDuration.Seconds())
-	number := minDurationInSeconds + rand.Intn(maxDurationInSeconds-minDurationInSeconds)
-
-	return time.Duration(number) * time.Second
-}

mtglib/internal/relay/timeouts_internal_test.go

@@ -1,37 +0,0 @@
-package relay
-
-import (
-	"fmt"
-	"testing"
-
-	"github.com/stretchr/testify/suite"
-)
-
-type TimeoutsTestSuite struct {
-	suite.Suite
-}
-
-func (suite *TimeoutsTestSuite) TestGetConnectionTimeToLive() {
-	for i := 0; i < 100; i++ {
-		value := getConnectionTimeToLive()
-		message := fmt.Sprintf("generated value is %v", value)
-
-		suite.GreaterOrEqual(value, ConnectionTimeToLiveMin, message)
-		suite.LessOrEqual(value, ConnectionTimeToLiveMax, message)
-	}
-}
-
-func (suite *TimeoutsTestSuite) TestGetTimeout() {
-	for i := 0; i < 100; i++ {
-		value := getTimeout()
-		message := fmt.Sprintf("generated value is %v", value)
-
-		suite.GreaterOrEqual(value, TimeoutMin, message)
-		suite.LessOrEqual(value, TimeoutMax, message)
-	}
-}
-
-func TestTimeouts(t *testing.T) {
-	t.Parallel()
-	suite.Run(t, &TimeoutsTestSuite{})
-}

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
 	)
 

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"
@@ -25,7 +26,6 @@ type Proxy struct {
 
 	allowFallbackOnUnknownDC bool
 	tolerateTimeSkewness     time.Duration
-	bufferSize               int
 	domainFrontingPort       int
 	workerPool               *ants.PoolWithFunc
 	telegram                 *telegram.Telegram
@@ -46,7 +46,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()
 
@@ -85,7 +85,6 @@ func (p *Proxy) ServeConn(conn net.Conn) {
 	relay.Relay(
 		ctx,
 		ctx.logger.Named("relay"),
-		p.bufferSize,
 		ctx.telegramConn,
 		ctx.clientConn,
 	)
@@ -276,7 +275,6 @@ func (p *Proxy) doDomainFronting(ctx *streamContext, conn *connRewind) {
 	relay.Relay(
 		ctx,
 		ctx.logger.Named("domain-fronting"),
-		p.bufferSize,
 		frontConn,
 		conn,
 	)
@@ -306,14 +304,13 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		logger:                   opts.getLogger("proxy"),
 		domainFrontingPort:       opts.getDomainFrontingPort(),
 		tolerateTimeSkewness:     opts.getTolerateTimeSkewness(),
-		bufferSize:               opts.getBufferSize(),
 		allowFallbackOnUnknownDC: opts.AllowFallbackOnUnknownDC,
 		telegram:                 tg,
 	}
 
 	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))

mtglib/proxy_opts.go

@@ -50,6 +50,8 @@ type ProxyOpts struct {
 	// buffers: to and from.
 	//
 	// This is an optional setting.
+	//
+	// Deprecated: this setting is no longer makes any effect.
 	BufferSize uint
 
 	// Concurrency is a size of the worker pool for connection management.
@@ -134,14 +136,6 @@ func (p ProxyOpts) valid() error {
 	return nil
 }
 
-func (p ProxyOpts) getBufferSize() int {
-	if p.BufferSize < 1 {
-		return DefaultBufferSize
-	}
-
-	return int(p.BufferSize)
-}
-
 func (p ProxyOpts) getConcurrency() int {
 	if p.Concurrency == 0 {
 		return DefaultConcurrency

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 {

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