From 114cd59f610c149523d93dad4c0fffc16c128bde Mon Sep 17 00:00:00 2001 From: Andrey Osipuk <45951478+AndreyOsipuk@users.noreply.github.com> Date: Tue, 6 Oct 2026 16:09:18 +0300 Subject: [PATCH 1/2] feat: opt-in per-IP limit on pending handshakes A connection holds a worker until it authenticates, goes to the fronting domain or hits the handshake timeout. One address that opens hundreds of connections and never completes the handshake can keep the whole concurrency budget busy and starve everyone else. [defense.pending-handshakes] max-per-ip limits concurrent unauthenticated handshakes from one IP. The slot is released as soon as the secret is verified or the connection is fronted, so authenticated sessions and bursts of short media connections (authenticated within milliseconds) are not counted. dry-run observes the limit without enforcing it. Disabled by default. Event PendingHandshakeLimit, metric pending_handshake_limit{action}. --- events/event_stream.go | 2 + events/init.go | 4 + events/init_test.go | 4 + events/multi_observer.go | 12 ++ events/noop.go | 21 +-- example.config.toml | 15 ++ internal/cli/run_proxy.go | 3 + internal/config/config.go | 4 + internal/config/config_test.go | 15 ++ internal/config/parse.go | 4 + .../config/testdata/pending_handshakes.toml | 6 + mtglib/events.go | 21 +++ mtglib/pending_handshakes.go | 105 ++++++++++++ mtglib/pending_handshakes_internal_test.go | 160 ++++++++++++++++++ mtglib/proxy.go | 33 ++++ mtglib/proxy_opts.go | 14 ++ mtglib/proxy_test.go | 119 +++++++++++++ mtglib/stream_context.go | 11 ++ stats/init.go | 9 + stats/prometheus.go | 11 ++ stats/prometheus_test.go | 13 ++ stats/statsd.go | 4 + 22 files changed, 580 insertions(+), 10 deletions(-) create mode 100644 internal/config/testdata/pending_handshakes.toml create mode 100644 mtglib/pending_handshakes.go create mode 100644 mtglib/pending_handshakes_internal_test.go diff --git a/events/event_stream.go b/events/event_stream.go index 516f77e7b..b04ce02c5 100644 --- a/events/event_stream.go +++ b/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: diff --git a/events/init.go b/events/init.go index 6d541161c..ba64eb348 100644 --- a/events/init.go +++ b/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) diff --git a/events/init_test.go b/events/init_test.go index 3fc4ff0a4..2f6eb16de 100644 --- a/events/init_test.go +++ b/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) } diff --git a/events/multi_observer.go b/events/multi_observer.go index f59d63e5b..e9a23a794 100644 --- a/events/multi_observer.go +++ b/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{} diff --git a/events/noop.go b/events/noop.go index 68ee0734e..7ed2f5faf 100644 --- a/events/noop.go +++ b/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 { diff --git a/example.config.toml b/example.config.toml index 8dec1a8d4..b973e2c06 100644 --- a/example.config.toml +++ b/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. # diff --git a/internal/cli/run_proxy.go b/internal/cli/run_proxy.go index f1ce8af8b..8ffaafde6 100644 --- a/internal/cli/run_proxy.go +++ b/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) diff --git a/internal/config/config.go b/internal/config/config.go index e9a8e9625..de7191011 100644 --- a/internal/config/config.go +++ b/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 diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 25c98233e..61a86f962 100644 --- a/internal/config/config_test.go +++ b/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{}) diff --git a/internal/config/parse.go b/internal/config/parse.go index 97eb19103..613ac251d 100644 --- a/internal/config/parse.go +++ b/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"` diff --git a/internal/config/testdata/pending_handshakes.toml b/internal/config/testdata/pending_handshakes.toml new file mode 100644 index 000000000..1ee9759a0 --- /dev/null +++ b/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 diff --git a/mtglib/events.go b/mtglib/events.go index 0a1aefd19..1d808ee8f 100644 --- a/mtglib/events.go +++ b/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 { diff --git a/mtglib/pending_handshakes.go b/mtglib/pending_handshakes.go new file mode 100644 index 000000000..89ddbb7fd --- /dev/null +++ b/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) +} diff --git a/mtglib/pending_handshakes_internal_test.go b/mtglib/pending_handshakes_internal_test.go new file mode 100644 index 000000000..62f2ce766 --- /dev/null +++ b/mtglib/pending_handshakes_internal_test.go @@ -0,0 +1,160 @@ +package mtglib + +import ( + "net" + "sync" + "sync/atomic" + "testing" +) + +var ( + testIPA = net.ParseIP("198.51.100.1") + testIPB = net.ParseIP("198.51.100.2") +) + +func TestPendingHandshakesDisabled(t *testing.T) { + t.Parallel() + + p := newPendingHandshakes() + + for range 100 { + release, action, ok := p.acquire(testIPA, 0, false) + if !ok || release != nil || action != "" { + t.Fatal("limit 0 must admit without tracking") + } + } + + if p.trackedIPs() != 0 { + t.Fatal("nothing must be tracked with limit 0") + } +} + +func TestPendingHandshakesRejectOverLimitAndRelease(t *testing.T) { + t.Parallel() + + p := newPendingHandshakes() + + var releases []func() + + for range 3 { + release, action, ok := p.acquire(testIPA, 3, false) + if !ok || action != "" { + t.Fatal("expected admission under the limit") + } + + releases = append(releases, release) + } + + if _, action, ok := p.acquire(testIPA, 3, false); ok || action != PendingHandshakeRejected { + t.Fatal("expected rejection over the limit") + } + + if p.pendingFor(testIPA) != 3 { + t.Fatalf("a rejection must not take a slot, pending %d", p.pendingFor(testIPA)) + } + + releases[0]() + releases[0]() // release is idempotent + + if p.pendingFor(testIPA) != 2 { + t.Fatalf("expected 2 pending after one release, got %d", p.pendingFor(testIPA)) + } + + if _, _, ok := p.acquire(testIPA, 3, false); !ok { + t.Fatal("a released slot must be available again") + } +} + +func TestPendingHandshakesAddressesAreIndependent(t *testing.T) { + t.Parallel() + + p := newPendingHandshakes() + + if _, _, ok := p.acquire(testIPA, 1, false); !ok { + t.Fatal("first address must be admitted") + } + + if _, _, ok := p.acquire(testIPA, 1, false); ok { + t.Fatal("first address must be over the limit") + } + + if _, _, ok := p.acquire(testIPB, 1, false); !ok { + t.Fatal("another address must not be affected") + } + + // IPv4-mapped IPv6 is the same address. + if _, _, ok := p.acquire(net.ParseIP("::ffff:198.51.100.1"), 1, false); ok { + t.Fatal("IPv4-mapped form must count as the same address") + } +} + +func TestPendingHandshakesDryRun(t *testing.T) { + t.Parallel() + + p := newPendingHandshakes() + + r1, action, ok := p.acquire(testIPA, 1, true) + if !ok || action != "" { + t.Fatal("first must be admitted silently") + } + + r2, action, ok := p.acquire(testIPA, 1, true) + if !ok || action != PendingHandshakeObserved { + t.Fatalf("over the limit in dry run must be admitted and observed, got %q %v", action, ok) + } + + if p.pendingFor(testIPA) != 2 { + t.Fatal("dry run must keep counting") + } + + r1() + r2() + + if p.trackedIPs() != 0 { + t.Fatal("entries must be removed when idle") + } +} + +func TestPendingHandshakesConcurrentNeverExceedsLimit(t *testing.T) { + t.Parallel() + + p := newPendingHandshakes() + + const limit = 8 + + var peak atomic.Uint32 + + var wg sync.WaitGroup + + for range 16 { + wg.Go(func() { + for range 2000 { + release, _, ok := p.acquire(testIPA, limit, false) + if !ok { + continue + } + + now := p.pendingFor(testIPA) + + for { + old := peak.Load() + if now <= old || peak.CompareAndSwap(old, now) { + break + } + } + + release() + } + }) + } + + wg.Wait() + + if peak.Load() > limit { + t.Fatalf("peak %d exceeds limit %d", peak.Load(), limit) + } + + if p.trackedIPs() != 0 { + t.Fatal("entries must be removed when idle") + } +} diff --git a/mtglib/proxy.go b/mtglib/proxy.go index 1ea6f2e63..97c6986db 100644 --- a/mtglib/proxy.go +++ b/mtglib/proxy.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "math" "net" "strconv" "sync" @@ -26,6 +27,9 @@ type Proxy struct { streamWaitGroup sync.WaitGroup allowFallbackOnUnknownDC bool + pendingHandshakes *pendingHandshakes + pendingHandshakesPerIP uint32 + pendingHandshakesDryRun bool tolerateTimeSkewness time.Duration idleTimeout time.Duration handshakeTimeout time.Duration @@ -86,10 +90,33 @@ func (p *Proxy) ServeConn(conn essentials.Conn) { ctx.logger.Info("Stream has been finished") }() + // Per-IP limit on pending (unauthenticated) handshakes. The slot is + // released as soon as the secret is verified or the connection goes to the + // fronting domain, so authenticated sessions never hold it. + release, action, admitted := p.pendingHandshakes.acquire( + ctx.ClientIP(), p.pendingHandshakesPerIP, p.pendingHandshakesDryRun) + if action != "" { + p.eventStream.Send(ctx, NewEventPendingHandshakeLimit(ctx.streamID, action)) + } + + if !admitted { + ctx.logger.Info("too many pending handshakes from this ip") + + return + } + + if release != nil { + defer release() + + ctx.releasePendingHandshake = release + } + if !p.doFakeTLSHandshake(ctx) { return } + ctx.finishPendingHandshake() + clientConn, err := p.doppelGanger.NewConn(ctx.clientConn) if err != nil { ctx.logger.InfoError("cannot wrap into doppelganger connection", err) @@ -295,6 +322,9 @@ func (p *Proxy) doTelegramCall(ctx *streamContext) error { } func (p *Proxy) doDomainFronting(ctx *streamContext, conn *connRewind) { + // The handshake is over: a fronted connection must not hold a pending slot. + ctx.finishPendingHandshake() + p.eventStream.Send(p.ctx, NewEventDomainFronting(ctx.streamID)) conn.Rewind() @@ -364,6 +394,9 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { idleTimeout: opts.getIdleTimeout(), handshakeTimeout: opts.getHandshakeTimeout(), allowFallbackOnUnknownDC: opts.AllowFallbackOnUnknownDC, + pendingHandshakes: newPendingHandshakes(), + pendingHandshakesPerIP: uint32(min(opts.PendingHandshakesPerIP, math.MaxUint32)), //nolint: gosec + pendingHandshakesDryRun: opts.PendingHandshakesDryRun, telegram: tg, doppelGanger: doppel.NewGanger( ctx, diff --git a/mtglib/proxy_opts.go b/mtglib/proxy_opts.go index 41ffa84e0..159c5f543 100644 --- a/mtglib/proxy_opts.go +++ b/mtglib/proxy_opts.go @@ -177,6 +177,20 @@ type ProxyOpts struct { // DoppelGangerDRS defines if TLS Dynamic Record Sizing is active. DoppelGangerDRS bool + + // PendingHandshakesPerIP limits concurrent unauthenticated handshakes from + // one client IP. A slot is released as soon as the secret is verified or + // the connection goes to the fronting domain, so authenticated sessions are + // never counted. Excess connections are closed. 0 disables the limit. + // + // This is an optional setting. + PendingHandshakesPerIP uint + + // PendingHandshakesDryRun observes PendingHandshakesPerIP without + // enforcing it: excess connections are admitted and only reported. + // + // This is an optional setting. + PendingHandshakesDryRun bool } func (p ProxyOpts) valid() error { diff --git a/mtglib/proxy_test.go b/mtglib/proxy_test.go index 278ce4f27..8395fa5d5 100644 --- a/mtglib/proxy_test.go +++ b/mtglib/proxy_test.go @@ -3,6 +3,7 @@ package mtglib_test import ( "crypto/tls" "encoding/json" + "errors" "fmt" "io" "net" @@ -198,3 +199,121 @@ func TestProxy(t *testing.T) { t.Parallel() suite.Run(t, &ProxyTestSuite{}) } + +// A second silent connection from the same IP is closed right away when the +// pending-handshake limit is 1; once the first one goes away, a new connection +// is admitted (it stays open waiting for its handshake). +func TestProxyPendingHandshakeLimit(t *testing.T) { + t.Parallel() + + dialer, err := network.NewDefaultDialer(0, 0) + if err != nil { + t.Fatal(err) + } + + ntw, err := network.NewNetwork(dialer, "mtgtest", "1.1.1.1", 0) + if err != nil { + t.Fatal(err) + } + + allowlist, _ := ipblocklist.NewFireholFromFiles( + logger.NewNoopLogger(), + 1, + []files.File{files.NewMem([]*net.IPNet{cidranger.AllIPv4, cidranger.AllIPv6})}, + nil, + ) + + go allowlist.Run(time.Second) + + proxy, err := mtglib.NewProxy(mtglib.ProxyOpts{ + Secret: mtglib.GenerateSecret("httpbin.org"), + Network: ntw, + AntiReplayCache: antireplay.NewNoop(), + IPBlocklist: ipblocklist.NewNoop(), + IPAllowlist: allowlist, + EventStream: events.NewNoopStream(), + Logger: logger.NewNoopLogger(), + UseTestDCs: true, + PendingHandshakesPerIP: 1, + }) + if err != nil { + t.Fatal(err) + } + + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + + go proxy.Serve(listener) //nolint: errcheck + + defer func() { + listener.Close() //nolint: errcheck + proxy.Shutdown() + }() + + closedQuickly := func(conn net.Conn) bool { + conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond)) //nolint: errcheck + + _, err := conn.Read(make([]byte, 1)) + + var netErr net.Error + + return err != nil && !(errors.As(err, &netErr) && netErr.Timeout()) + } + + dial := func() net.Conn { + conn, err := net.Dial("tcp", listener.Addr().String()) + if err != nil { + t.Fatal(err) + } + + return conn + } + + // Open silent connections one by one until one is closed: with the limit + // of 1 exactly one must stay pending. No fixed sleeps - the proxy may pick + // a connection up with a delay under load. + var admitted []net.Conn + + rejected := false + + for range 10 { + conn := dial() + if closedQuickly(conn) { + conn.Close() //nolint: errcheck + + rejected = true + + break + } + + admitted = append(admitted, conn) + } + + if !rejected { + t.Fatal("connections over the pending-handshake limit must be closed") + } + + if len(admitted) != 1 { + t.Fatalf("exactly one pending handshake must be admitted with limit 1, got %d", len(admitted)) + } + + // When the pending connection goes away, its slot is released. + admitted[0].Close() //nolint: errcheck + + for attempt := range 10 { + conn := dial() + ok := !closedQuickly(conn) + + conn.Close() //nolint: errcheck + + if ok { + return + } + + if attempt == 9 { + t.Fatal("after the pending connection is gone, a new one must be admitted") + } + } +} diff --git a/mtglib/stream_context.go b/mtglib/stream_context.go index 3858282ba..0095af22a 100644 --- a/mtglib/stream_context.go +++ b/mtglib/stream_context.go @@ -18,6 +18,17 @@ type streamContext struct { streamID string dc int logger Logger + + // releasePendingHandshake frees the per-IP pending-handshake slot, if any. + releasePendingHandshake func() +} + +// finishPendingHandshake releases the pending-handshake slot of this stream. +// It is safe to call several times. +func (s *streamContext) finishPendingHandshake() { + if s.releasePendingHandshake != nil { + s.releasePendingHandshake() + } } func (s *streamContext) Deadline() (time.Time, bool) { diff --git a/stats/init.go b/stats/init.go index a2236e77f..6da2c79a1 100644 --- a/stats/init.go +++ b/stats/init.go @@ -76,6 +76,12 @@ const ( // Type: counter MetricConcurrencyLimited = "concurrency_limited" + // MetricPendingHandshakeLimit defines a metric for connections over the + // per-IP limit on pending handshakes. + // + // Type: counter, tags: TagAction. + MetricPendingHandshakeLimit = "pending_handshake_limit" + // MetricIPBlocklisted defines a metric for a count of events, when // client was blocked because her IP address was found in blocklists. // @@ -111,6 +117,9 @@ const ( // TagDC defines a name of the 'dc' tag. TagDC = "dc" + // TagAction defines a name of the 'action' tag of MetricPendingHandshakeLimit. + TagAction = "action" + // TagDirection defines a name of the 'direction' tag. TagDirection = "direction" diff --git a/stats/prometheus.go b/stats/prometheus.go index 920332c4d..f4231579e 100644 --- a/stats/prometheus.go +++ b/stats/prometheus.go @@ -123,6 +123,10 @@ func (p prometheusProcessor) EventReplayAttack(_ mtglib.EventReplayAttack) { p.factory.metricReplayAttacks.Inc() } +func (p prometheusProcessor) EventPendingHandshakeLimit(evt mtglib.EventPendingHandshakeLimit) { + p.factory.metricPendingHandshakeLimit.WithLabelValues(evt.Action).Inc() +} + func (p prometheusProcessor) EventIPListSize(evt mtglib.EventIPListSize) { tag := TagIPListBlock if !evt.IsBlockList { @@ -155,6 +159,7 @@ type PrometheusFactory struct { metricTelegramTraffic *prometheus.CounterVec metricDomainFrontingTraffic *prometheus.CounterVec metricIPBlocklisted *prometheus.CounterVec + metricPendingHandshakeLimit *prometheus.CounterVec metricDomainFronting prometheus.Counter metricConcurrencyLimited prometheus.Counter @@ -232,6 +237,11 @@ func NewPrometheus(metricPrefix, httpPath string) *PrometheusFactory { //nolint: Name: MetricIPBlocklisted, Help: "A number of rejected sessions due to ip blocklisting.", }, []string{TagIPList}), + metricPendingHandshakeLimit: prometheus.NewCounterVec(prometheus.CounterOpts{ + Namespace: metricPrefix, + Name: MetricPendingHandshakeLimit, + Help: "Connections over the per-IP limit on pending handshakes: rejected or observed (dry run).", + }, []string{TagAction}), metricDomainFronting: prometheus.NewCounter(prometheus.CounterOpts{ Namespace: metricPrefix, @@ -258,6 +268,7 @@ func NewPrometheus(metricPrefix, httpPath string) *PrometheusFactory { //nolint: registry.MustRegister(factory.metricTelegramTraffic) registry.MustRegister(factory.metricDomainFrontingTraffic) registry.MustRegister(factory.metricIPBlocklisted) + registry.MustRegister(factory.metricPendingHandshakeLimit) registry.MustRegister(factory.metricDomainFronting) registry.MustRegister(factory.metricConcurrencyLimited) diff --git a/stats/prometheus_test.go b/stats/prometheus_test.go index 7448366d9..c0dcbbd68 100644 --- a/stats/prometheus_test.go +++ b/stats/prometheus_test.go @@ -148,6 +148,19 @@ func (suite *PrometheusTestSuite) TestEventConcurrencyLimited() { suite.Contains(data, `mtg_concurrency_limited 1`) } +func (suite *PrometheusTestSuite) TestEventPendingHandshakeLimit() { + suite.prometheus.EventPendingHandshakeLimit(mtglib.NewEventPendingHandshakeLimit("s1", mtglib.PendingHandshakeRejected)) + suite.prometheus.EventPendingHandshakeLimit(mtglib.NewEventPendingHandshakeLimit("s2", mtglib.PendingHandshakeRejected)) + suite.prometheus.EventPendingHandshakeLimit(mtglib.NewEventPendingHandshakeLimit("s3", mtglib.PendingHandshakeObserved)) + + time.Sleep(100 * time.Millisecond) + + data, err := suite.Get() + suite.NoError(err) + suite.Contains(data, `mtg_pending_handshake_limit{action="rejected"} 2`) + suite.Contains(data, `mtg_pending_handshake_limit{action="observed"} 1`) +} + func (suite *PrometheusTestSuite) TestEventIPBlocklisted() { suite.prometheus.EventIPBlocklisted( mtglib.NewEventIPBlocklisted(net.ParseIP("2001:db8::68"))) diff --git a/stats/statsd.go b/stats/statsd.go index 64b51172b..1344f66d7 100644 --- a/stats/statsd.go +++ b/stats/statsd.go @@ -126,6 +126,10 @@ func (s statsdProcessor) EventReplayAttack(_ mtglib.EventReplayAttack) { s.client.Incr(MetricReplayAttacks, 1) } +func (s statsdProcessor) EventPendingHandshakeLimit(evt mtglib.EventPendingHandshakeLimit) { + s.client.Incr(MetricPendingHandshakeLimit, 1, statsd.StringTag(TagAction, evt.Action)) +} + func (s statsdProcessor) EventIPListSize(evt mtglib.EventIPListSize) { tag := TagIPListBlock if !evt.IsBlockList { From 1a5c9a8b271e484a683bd77d4c88e0f4afd990f3 Mon Sep 17 00:00:00 2001 From: Andrey Osipuk <45951478+AndreyOsipuk@users.noreply.github.com> Date: Tue, 6 Oct 2026 23:57:11 +0300 Subject: [PATCH 2/2] test: wait for the allowlist before probing the pending-handshake limit The allowlist loads asynchronously and rejects every connection until it is ready, which on a slow CI runner looked like a rejection by the limit (0 admitted). Also apply De Morgan's law as staticcheck suggests. --- mtglib/proxy_test.go | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/mtglib/proxy_test.go b/mtglib/proxy_test.go index 8395fa5d5..2b53a4a79 100644 --- a/mtglib/proxy_test.go +++ b/mtglib/proxy_test.go @@ -225,6 +225,12 @@ func TestProxyPendingHandshakeLimit(t *testing.T) { go allowlist.Run(time.Second) + // The allowlist loads asynchronously; until then every connection is + // rejected by it, which would look like a rejection by the limit. + for !allowlist.Contains(net.ParseIP("127.0.0.1")) { + time.Sleep(10 * time.Millisecond) + } + proxy, err := mtglib.NewProxy(mtglib.ProxyOpts{ Secret: mtglib.GenerateSecret("httpbin.org"), Network: ntw, @@ -259,7 +265,7 @@ func TestProxyPendingHandshakeLimit(t *testing.T) { var netErr net.Error - return err != nil && !(errors.As(err, &netErr) && netErr.Timeout()) + return err != nil && (!errors.As(err, &netErr) || !netErr.Timeout()) } dial := func() net.Conn {