feat: opt-in per-IP limit on pending handshakes

events/event_stream.go

@@ -100,6 +100,8 @@ func eventStreamProcessor(ctx context.Context, eventChan <-chan mtglib.Event, ob
 				observer.EventIPBlocklisted(typedEvt)
 			case mtglib.EventConcurrencyLimited:
 				observer.EventConcurrencyLimited(typedEvt)
+			case mtglib.EventPendingHandshakeLimit:
+				observer.EventPendingHandshakeLimit(typedEvt)
 			case mtglib.EventReplayAttack:
 				observer.EventReplayAttack(typedEvt)
 			case mtglib.EventIPListSize:

events/init.go

@@ -47,6 +47,10 @@ type Observer interface {
 	// mtglib.EventConcurrencyLimited event.
 	EventConcurrencyLimited(mtglib.EventConcurrencyLimited)
 
+	// EventPendingHandshakeLimit reacts on incoming
+	// mtglib.EventPendingHandshakeLimit event.
+	EventPendingHandshakeLimit(mtglib.EventPendingHandshakeLimit)
+
 	// EventIPBlocklisted reacts on incoming mtglib.EventIPBlocklisted event.
 	EventIPBlocklisted(mtglib.EventIPBlocklisted)
 

events/init_test.go

@@ -33,6 +33,10 @@ func (o *ObserverMock) EventConcurrencyLimited(evt mtglib.EventConcurrencyLimite
 	o.Called(evt)
 }
 
+func (o *ObserverMock) EventPendingHandshakeLimit(evt mtglib.EventPendingHandshakeLimit) {
+	o.Called(evt)
+}
+
 func (o *ObserverMock) EventIPBlocklisted(evt mtglib.EventIPBlocklisted) {
 	o.Called(evt)
 }

events/multi_observer.go

@@ -70,6 +70,18 @@ func (m multiObserver) EventFinish(evt mtglib.EventFinish) {
 	wg.Wait()
 }
 
+func (m multiObserver) EventPendingHandshakeLimit(evt mtglib.EventPendingHandshakeLimit) {
+	wg := &sync.WaitGroup{}
+
+	for _, v := range m.observers {
+		wg.Go(func() {
+			v.EventPendingHandshakeLimit(evt)
+		})
+	}
+
+	wg.Wait()
+}
+
 func (m multiObserver) EventConcurrencyLimited(evt mtglib.EventConcurrencyLimited) {
 	wg := &sync.WaitGroup{}
 

events/noop.go

@@ -17,16 +17,17 @@ func NewNoopStream() mtglib.EventStream {
 
 type noopObserver struct{}
 
-func (n noopObserver) EventStart(_ mtglib.EventStart)                           {}
-func (n noopObserver) EventConnectedToDC(_ mtglib.EventConnectedToDC)           {}
-func (n noopObserver) EventDomainFronting(_ mtglib.EventDomainFronting)         {}
-func (n noopObserver) EventTraffic(_ mtglib.EventTraffic)                       {}
-func (n noopObserver) EventFinish(_ mtglib.EventFinish)                         {}
-func (n noopObserver) EventConcurrencyLimited(_ mtglib.EventConcurrencyLimited) {}
-func (n noopObserver) EventIPBlocklisted(_ mtglib.EventIPBlocklisted)           {}
-func (n noopObserver) EventReplayAttack(_ mtglib.EventReplayAttack)             {}
-func (n noopObserver) EventIPListSize(_ mtglib.EventIPListSize)                 {}
-func (n noopObserver) Shutdown()                                                {}
+func (n noopObserver) EventStart(_ mtglib.EventStart)                                 {}
+func (n noopObserver) EventConnectedToDC(_ mtglib.EventConnectedToDC)                 {}
+func (n noopObserver) EventDomainFronting(_ mtglib.EventDomainFronting)               {}
+func (n noopObserver) EventTraffic(_ mtglib.EventTraffic)                             {}
+func (n noopObserver) EventFinish(_ mtglib.EventFinish)                               {}
+func (n noopObserver) EventConcurrencyLimited(_ mtglib.EventConcurrencyLimited)       {}
+func (n noopObserver) EventIPBlocklisted(_ mtglib.EventIPBlocklisted)                 {}
+func (n noopObserver) EventReplayAttack(_ mtglib.EventReplayAttack)                   {}
+func (n noopObserver) EventIPListSize(_ mtglib.EventIPListSize)                       {}
+func (n noopObserver) EventPendingHandshakeLimit(_ mtglib.EventPendingHandshakeLimit) {}
+func (n noopObserver) Shutdown()                                                      {}
 
 // NewNoopObserver creates an observer which discards each message.
 func NewNoopObserver() Observer {

example.config.toml

@@ -313,6 +313,21 @@ raid-each = "6h"
 #      https://github.com/cloudflare/sslconfig/blob/master/patches/nginx__dynamic_tls_records.patch
 drs = false
 
+# Limit on concurrent unauthenticated handshakes from one client IP. A
+# connection that never completes the handshake holds a worker until the
+# handshake timeout; one address that opens hundreds of them can starve
+# everyone else. Only connections that have not proven a secret yet are
+# counted: the slot is released as soon as the secret is verified or the
+# connection goes to the fronting domain, so authenticated sessions and
+# bursts of short media connections are not affected. Excess connections are
+# closed. Clients behind a carrier-grade NAT share an address: measure with
+# dry-run first. Reported as the pending_handshake_limit metric.
+[defense.pending-handshakes]
+# 0 (default) disables the limit.
+max-per-ip = 0
+# Only count connections over the limit, do not close them.
+dry-run = false
+
 # Some countries do active probing on Telegram connections. This technique
 # allows to protect from such effort.
 #

internal/cli/run_proxy.go

@@ -326,6 +326,9 @@ func runProxy(conf *config.Config, version string) error { //nolint: funlen, cyc
 		DoppelGangerPerRaid: conf.Defense.Doppelganger.Repeats.Get(mtglib.DoppelGangerPerRaid),
 		DoppelGangerEach:    conf.Defense.Doppelganger.UpdateEach.Get(mtglib.DoppelGangerEach),
 		DoppelGangerDRS:     conf.Defense.Doppelganger.DRS.Get(false),
+
+		PendingHandshakesPerIP:  conf.Defense.PendingHandshakes.MaxPerIP.Get(0),
+		PendingHandshakesDryRun: conf.Defense.PendingHandshakes.DryRun.Get(false),
 	}
 
 	proxy, err := mtglib.NewProxy(opts)

internal/config/config.go

@@ -43,6 +43,10 @@ type Config struct {
 		ProxyProtocol TypeBool `json:"proxyProtocol"`
 	} `json:"domainFronting"`
 	Defense struct {
+		PendingHandshakes struct {
+			MaxPerIP TypeConcurrency `json:"maxPerIp"`
+			DryRun   TypeBool        `json:"dryRun"`
+		} `json:"pendingHandshakes"`
 		AntiReplay struct {
 			Optional
 

internal/config/config_test.go

@@ -115,6 +115,21 @@ func (suite *ConfigTestSuite) TestDomainFrontingNotSet() {
 	suite.Equal("", conf.GetDomainFrontingHost())
 }
 
+func (suite *ConfigTestSuite) TestPendingHandshakesDisabledByDefault() {
+	conf, err := config.Parse(suite.ReadConfig("minimal.toml"))
+	suite.NoError(err)
+	suite.EqualValues(0, conf.Defense.PendingHandshakes.MaxPerIP.Get(0))
+	suite.False(conf.Defense.PendingHandshakes.DryRun.Get(false))
+}
+
+func (suite *ConfigTestSuite) TestPendingHandshakes() {
+	conf, err := config.Parse(suite.ReadConfig("pending_handshakes.toml"))
+	suite.NoError(err)
+	suite.NoError(conf.Validate())
+	suite.EqualValues(32, conf.Defense.PendingHandshakes.MaxPerIP.Get(0))
+	suite.True(conf.Defense.PendingHandshakes.DryRun.Get(false))
+}
+
 func TestConfig(t *testing.T) {
 	t.Parallel()
 	suite.Run(t, &ConfigTestSuite{})

internal/config/parse.go

@@ -30,6 +30,10 @@ type tomlConfig struct {
 		ProxyProtocol bool   `toml:"proxy-protocol" json:"proxyProtocol,omitempty"`
 	} `toml:"domain-fronting" json:"domainFronting,omitempty"`
 	Defense struct {
+		PendingHandshakes struct {
+			MaxPerIP uint `toml:"max-per-ip" json:"maxPerIp,omitempty"`
+			DryRun   bool `toml:"dry-run" json:"dryRun,omitempty"`
+		} `toml:"pending-handshakes" json:"pendingHandshakes,omitempty"`
 		AntiReplay struct {
 			Enabled   bool    `toml:"enabled" json:"enabled,omitempty"`
 			MaxSize   string  `toml:"max-size" json:"maxSize,omitempty"`

internal/config/testdata/pending_handshakes.toml

@@ -0,0 +1,6 @@
+secret = "7oe1GqLy6TBc38CV3jx7q09nb29nbGUuY29t"
+bind-to = "0.0.0.0:3128"
+
+[defense.pending-handshakes]
+max-per-ip = 32
+dry-run = true

mtglib/events.go

@@ -157,6 +157,27 @@ func NewEventDomainFronting(streamID string) EventDomainFronting {
 	}
 }
 
+// EventPendingHandshakeLimit is emitted when a connection exceeds the per-IP
+// limit on pending (unauthenticated) handshakes: it is either rejected or, in
+// dry-run mode, observed and admitted.
+type EventPendingHandshakeLimit struct {
+	eventBase
+
+	// Action is PendingHandshakeRejected or PendingHandshakeObserved.
+	Action string
+}
+
+// NewEventPendingHandshakeLimit creates a new EventPendingHandshakeLimit.
+func NewEventPendingHandshakeLimit(streamID, action string) EventPendingHandshakeLimit {
+	return EventPendingHandshakeLimit{
+		eventBase: eventBase{
+			timestamp: time.Now(),
+			streamID:  streamID,
+		},
+		Action: action,
+	}
+}
+
 // NewEventConcurrencyLimited creates a new EventConcurrencyLimited
 // event.
 func NewEventConcurrencyLimited() EventConcurrencyLimited {

mtglib/pending_handshakes.go

@@ -0,0 +1,105 @@
+package mtglib
+
+import (
+	"net"
+	"net/netip"
+	"sync"
+)
+
+// Results of a pending-handshake admission, reported as the action label of
+// the pending_handshake_limit metric.
+const (
+	PendingHandshakeRejected = "rejected"
+	PendingHandshakeObserved = "observed"
+)
+
+// pendingHandshakes limits concurrent unauthenticated handshakes per client
+// IP address.
+//
+// A connection that never completes the handshake holds a worker until the
+// handshake timeout. One address that opens hundreds of such connections can
+// keep the proxy busy and starve legitimate clients. Only connections that
+// have not proven a secret yet are counted: a slot is released as soon as the
+// secret is verified or the connection goes to the fronting domain, so
+// authenticated sessions and bursts of short media connections are not
+// affected. Entries exist only while an address has pending handshakes.
+type pendingHandshakes struct {
+	mu      sync.Mutex
+	pending map[netip.Addr]uint32
+}
+
+func newPendingHandshakes() *pendingHandshakes {
+	return &pendingHandshakes{pending: map[netip.Addr]uint32{}}
+}
+
+// acquire takes a slot for ip. limit == 0 disables the check. It returns an
+// idempotent release function (nil if nothing was taken), the metric action
+// ("" when under the limit) and whether the connection may proceed. In dry-run
+// mode an address over the limit is admitted, still tracked, and reported as
+// observed.
+func (p *pendingHandshakes) acquire(ip net.IP, limit uint32, dryRun bool) (func(), string, bool) {
+	if limit == 0 {
+		return nil, "", true
+	}
+
+	addr, ok := netip.AddrFromSlice(ip)
+	if !ok {
+		return nil, "", true
+	}
+
+	addr = addr.Unmap()
+
+	p.mu.Lock()
+
+	overLimit := p.pending[addr] >= limit
+	if overLimit && !dryRun {
+		p.mu.Unlock()
+
+		return nil, PendingHandshakeRejected, false
+	}
+
+	p.pending[addr]++
+	p.mu.Unlock()
+
+	var once sync.Once
+
+	release := func() {
+		once.Do(func() {
+			p.mu.Lock()
+			defer p.mu.Unlock()
+
+			if p.pending[addr] <= 1 {
+				delete(p.pending, addr)
+			} else {
+				p.pending[addr]--
+			}
+		})
+	}
+
+	if overLimit {
+		return release, PendingHandshakeObserved, true
+	}
+
+	return release, "", true
+}
+
+// pendingFor returns the number of pending handshakes of ip.
+func (p *pendingHandshakes) pendingFor(ip net.IP) uint32 {
+	addr, ok := netip.AddrFromSlice(ip)
+	if !ok {
+		return 0
+	}
+
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	return p.pending[addr.Unmap()]
+}
+
+// trackedIPs returns the number of addresses with pending handshakes.
+func (p *pendingHandshakes) trackedIPs() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	return len(p.pending)
+}