Merge pull request #447 from 9seconds/handshake-timeout

Add separate handshake timeout

mtglib/init.go

@@ -81,6 +81,10 @@ const (
 	// avoid racing with MTProto ping_delay_disconnect (~60s interval).
 	DefaultIdleTimeout = 5 * time.Minute
 
+	// DefaultHandshakeTimeout defines a time period during which the
+	// all handshake ceremonies must be completed.
+	DefaultHandshakeTimeout = 10 * time.Second
+
 	// DefaultTolerateTimeSkewness is a default timeout for time skewness on a
 	// faketls timeout verification.
 	DefaultTolerateTimeSkewness = 3 * time.Second

mtglib/internal/tls/fake/client_side.go

@@ -47,11 +47,6 @@ func ReadClientHello(
 	hostname string,
 	tolerateTimeSkewness time.Duration,
 ) (*ClientHello, error) {
-	if err := conn.SetReadDeadline(time.Now().Add(ClientHelloReadTimeout)); err != nil {
-		return nil, fmt.Errorf("cannot set read deadline: %w", err)
-	}
-	defer conn.SetReadDeadline(resetDeadline) //nolint: errcheck
-
 	// This is how FakeTLS is organized:
 	//  1. We create sha256 HMAC with a given secret
 	//  2. We dump there a whole TLS frame except of the fact that random

mtglib/internal/tls/fake/client_side_snapshot_test.go

@@ -12,7 +12,6 @@ import (
 	"github.com/9seconds/mtg/v2/mtglib"
 	"github.com/9seconds/mtg/v2/mtglib/internal/tls/fake"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 )
@@ -71,11 +70,6 @@ func (suite *ParseClientHelloSnapshotTestSuite) makeConn(data []byte) *parseClie
 		readBuf: readBuf,
 	}
 
-	connMock.
-		On("SetReadDeadline", mock.AnythingOfType("time.Time")).
-		Twice().
-		Return(nil)
-
 	return connMock
 }
 

mtglib/internal/tls/fake/client_side_test.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"encoding/binary"
 	"encoding/json"
-	"errors"
 	"io"
 	"os"
 	"testing"
@@ -14,7 +13,6 @@ import (
 	"github.com/9seconds/mtg/v2/mtglib"
 	"github.com/9seconds/mtg/v2/mtglib/internal/tls"
 	"github.com/9seconds/mtg/v2/mtglib/internal/tls/fake"
-	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 )
@@ -53,11 +51,6 @@ func (suite *ParseClientHelloTestSuite) SetupTest() {
 	suite.connMock = &parseClientHelloConnMock{
 		readBuf: suite.readBuf,
 	}
-
-	suite.connMock.
-		On("SetReadDeadline", mock.AnythingOfType("time.Time")).
-		Twice().
-		Return(nil)
 }
 
 func (suite *ParseClientHelloTestSuite) TearDownTest() {
@@ -69,23 +62,11 @@ type ParseClientHello_TLSHeaderTestSuite struct {
 }
 
 func (suite *ParseClientHello_TLSHeaderTestSuite) TestEmpty() {
-	suite.connMock.ExpectedCalls = []*mock.Call{}
-	suite.connMock.
-		On("SetReadDeadline", mock.AnythingOfType("time.Time")).
-		Once().
-		Return(errors.New("fail"))
-
 	_, err := fake.ReadClientHello(suite.connMock, suite.secret.Key[:], suite.secret.Host, TolerateTime)
-	suite.ErrorContains(err, "fail")
+	suite.ErrorContains(err, "cannot read client hello")
 }
 
 func (suite *ParseClientHello_TLSHeaderTestSuite) TestNothing() {
-	suite.connMock.ExpectedCalls = []*mock.Call{}
-	suite.connMock.
-		On("SetReadDeadline", mock.AnythingOfType("time.Time")).
-		Twice().
-		Return(nil)
-
 	_, err := fake.ReadClientHello(suite.connMock, suite.secret.Key[:], suite.secret.Host, TolerateTime)
 	suite.ErrorIs(err, io.EOF)
 }
@@ -479,11 +460,6 @@ func (s *ParseClientHelloFragmentedTestSuite) makeConn(data []byte) *parseClient
 		readBuf: readBuf,
 	}
 
-	connMock.
-		On("SetReadDeadline", mock.AnythingOfType("time.Time")).
-		Twice().
-		Return(nil)
-
 	return connMock
 }
 

mtglib/internal/tls/fake/init.go

@@ -2,15 +2,6 @@ package fake
 
 import (
 	"errors"
-	"time"
 )
 
-const (
-	ClientHelloReadTimeout = 5 * time.Second
-)
-
-var (
-	resetDeadline time.Time
-
-	ErrBadDigest = errors.New("incorrect client random")
-)
+var ErrBadDigest = errors.New("incorrect client random")

mtglib/proxy.go

@@ -28,6 +28,7 @@ type Proxy struct {
 	allowFallbackOnUnknownDC    bool
 	tolerateTimeSkewness        time.Duration
 	idleTimeout                 time.Duration
+	handshakeTimeout            time.Duration
 	domainFrontingPort          int
 	domainFrontingIP            string
 	domainFrontingProxyProtocol bool
@@ -66,6 +67,11 @@ func (p *Proxy) ServeConn(conn essentials.Conn) {
 	ctx := newStreamContext(p.ctx, p.logger, conn)
 	defer ctx.Close()
 
+	if err := ctx.clientConn.SetDeadline(time.Now().Add(p.handshakeTimeout)); err != nil {
+		ctx.logger.WarningError("cannot set handshake timeout", err)
+		return
+	}
+
 	stop := context.AfterFunc(ctx, func() {
 		ctx.Close()
 	})
@@ -97,6 +103,11 @@ func (p *Proxy) ServeConn(conn essentials.Conn) {
 		return
 	}
 
+	if err := ctx.clientConn.SetDeadline(time.Time{}); err != nil {
+		ctx.logger.WarningError("cannot set deadline", err)
+		return
+	}
+
 	if err := p.doTelegramCall(ctx); err != nil {
 		ctx.logger.WarningError("cannot dial to telegram", err)
 		return
@@ -346,6 +357,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		domainFrontingIP:         opts.DomainFrontingIP,
 		tolerateTimeSkewness:     opts.getTolerateTimeSkewness(),
 		idleTimeout:              opts.getIdleTimeout(),
+		handshakeTimeout:         opts.getHandshakeTimeout(),
 		allowFallbackOnUnknownDC: opts.AllowFallbackOnUnknownDC,
 		telegram:                 tg,
 		doppelGanger: doppel.NewGanger(

mtglib/proxy_opts.go

@@ -70,6 +70,12 @@ type ProxyOpts struct {
 	// This is an optional setting.
 	IdleTimeout time.Duration
 
+	// HandshakeTimeout is a timeout during which all handshake ceremonies must
+	// be completed, otherwise this process will be aborted
+	//
+	// This is an optional setting.
+	HandshakeTimeout time.Duration
+
 	// TolerateTimeSkewness is a time boundary that defines a time range where
 	// faketls timestamp is acceptable.
 	//
@@ -215,6 +221,14 @@ func (p ProxyOpts) getPreferIP() string {
 	return p.PreferIP
 }
 
+func (p ProxyOpts) getHandshakeTimeout() time.Duration {
+	if p.HandshakeTimeout == 0 {
+		return DefaultHandshakeTimeout
+	}
+
+	return p.HandshakeTimeout
+}
+
 func (p ProxyOpts) getIdleTimeout() time.Duration {
 	if p.IdleTimeout == 0 {
 		return DefaultIdleTimeout