From ad8d94e62cc36df5659c468254ffdd482a86995e Mon Sep 17 00:00:00 2001 From: Boris Nagaev Date: Mon, 13 Jul 2026 19:16:16 -0500 Subject: [PATCH] connmgr: bound accepted inbound connections In this commit, we add an optional inbound connection limit at the listener boundary. Each accepted socket owns a slot until Close releases it exactly once, which makes connection lifetime, rather than accept rate, the accounting unit. The default remains unlimited for existing callers. The new tests cover the configured bound, concurrent admission, callback ordering, and slot reuse. --- connmgr/connmanager.go | 79 ++++++++++++- connmgr/connmanager_test.go | 214 ++++++++++++++++++++++++++++++++++++ go.mod | 3 +- go.sum | 2 + 4 files changed, 291 insertions(+), 7 deletions(-) diff --git a/connmgr/connmanager.go b/connmgr/connmanager.go index 97fcd6c9..58259ee6 100644 --- a/connmgr/connmanager.go +++ b/connmgr/connmanager.go @@ -11,6 +11,8 @@ import ( "sync" "sync/atomic" "time" + + "golang.org/x/time/rate" ) // maxFailedAttempts is the maximum number of successive failed connection @@ -120,6 +122,11 @@ type Config struct { // connections in that case. OnAccept func(net.Conn) + // MaxInbound limits the number of accepted inbound connections that may + // be active at once. A nil value preserves the historical unlimited + // behavior, while a pointer to zero disables inbound connections. + MaxInbound *uint32 + // TargetOutbound is the number of outbound network connections to // maintain. Defaults to 8. TargetOutbound uint32 @@ -191,17 +198,66 @@ type handleFailed struct { // ConnManager provides a manager to handle network connections. type ConnManager struct { // The following variables must only be used atomically. - connReqCount uint64 - start int32 - stop int32 + connReqCount uint64 + inboundRejected uint64 + start int32 + stop int32 cfg Config wg sync.WaitGroup failedAttempts uint64 requests chan interface{} + inboundSlots chan struct{} + inboundLog rate.Sometimes quit chan struct{} } +// limitedInboundConn releases an inbound connection slot exactly once when +// the connection is closed. +type limitedInboundConn struct { + net.Conn + + releaseOnce sync.Once + release func() +} + +// Close closes the underlying connection and releases its inbound slot. +func (c *limitedInboundConn) Close() error { + err := c.Conn.Close() + c.releaseOnce.Do(c.release) + + return err +} + +// limitInbound attempts to reserve an inbound slot for conn. The returned +// connection releases the slot when it is closed. +func (cm *ConnManager) limitInbound(conn net.Conn) (net.Conn, bool) { + if cm.inboundSlots == nil { + return conn, true + } + + select { + case cm.inboundSlots <- struct{}{}: + return &limitedInboundConn{ + Conn: conn, + release: func() { + <-cm.inboundSlots + }, + }, true + + default: + rejected := atomic.AddUint64(&cm.inboundRejected, 1) + cm.inboundLog.Do(func() { + log.Warnf("Inbound connection limit reached: "+ + "rejected=%d remote=%s", rejected, + conn.RemoteAddr()) + }) + + _ = conn.Close() + return nil, false + } +} + // handleFailedConn handles a connection failed due to a disconnect or any // other failure. If permanent, it retries the connection after the configured // retry duration. Otherwise, if required, it makes a new connection request. @@ -527,6 +583,12 @@ func (cm *ConnManager) listenHandler(listener net.Listener) { } continue } + + conn, ok := cm.limitInbound(conn) + if !ok { + continue + } + go cm.cfg.OnAccept(conn) } @@ -597,9 +659,14 @@ func New(cfg *Config) (*ConnManager, error) { cfg.TargetOutbound = defaultTargetOutbound } cm := ConnManager{ - cfg: *cfg, // Copy so caller can't mutate - requests: make(chan interface{}), - quit: make(chan struct{}), + cfg: *cfg, // Copy so caller can't mutate. + requests: make(chan interface{}), + inboundLog: rate.Sometimes{First: 3, Interval: 30 * time.Second}, + quit: make(chan struct{}), } + if cfg.MaxInbound != nil { + cm.inboundSlots = make(chan struct{}, int(*cfg.MaxInbound)) + } + return &cm, nil } diff --git a/connmgr/connmanager_test.go b/connmgr/connmanager_test.go index 94cb65ff..49c5f843 100644 --- a/connmgr/connmanager_test.go +++ b/connmgr/connmanager_test.go @@ -9,6 +9,7 @@ import ( "fmt" "io" "net" + "sync" "sync/atomic" "testing" "time" @@ -83,6 +84,160 @@ func TestNewConfig(t *testing.T) { } } +// TestInboundLimit verifies that accepted inbound connections hold slots for +// their entire lifetime and release them exactly once on close. +func TestInboundLimit(t *testing.T) { + maxInbound := uint32(2) + + cmgr, err := New(&Config{ + Dial: mockDialer, + MaxInbound: &maxInbound, + }) + if err != nil { + t.Fatalf("New error: %v", err) + } + + newConn := func(port int) net.Conn { + return &mockConn{rAddr: &net.TCPAddr{ + IP: net.ParseIP("192.0.2.1"), + Port: port, + }} + } + + first, ok := cmgr.limitInbound(newConn(1001)) + if !ok { + t.Fatal("first inbound connection was rejected") + } + second, ok := cmgr.limitInbound(newConn(1002)) + if !ok { + t.Fatal("second inbound connection was rejected") + } + if _, ok := cmgr.limitInbound(newConn(1003)); ok { + t.Fatal("connection above the inbound limit was accepted") + } + + if got := len(cmgr.inboundSlots); got != int(maxInbound) { + t.Fatalf("unexpected occupied slots: got %d, want %d", got, + maxInbound) + } + + if err := first.Close(); err != nil { + t.Fatalf("first close failed: %v", err) + } + if err := first.Close(); err != nil { + t.Fatalf("second close failed: %v", err) + } + if got := len(cmgr.inboundSlots); got != 1 { + t.Fatalf("double close released the slot more than once: got %d", got) + } + + replacement, ok := cmgr.limitInbound(newConn(1004)) + if !ok { + t.Fatal("replacement connection was rejected after slot release") + } + + _ = second.Close() + _ = replacement.Close() + if got := len(cmgr.inboundSlots); got != 0 { + t.Fatalf("inbound slots leaked: got %d", got) + } +} + +// TestInboundLimitDisabled verifies that an explicit zero limit rejects every +// accepted inbound connection. +func TestInboundLimitDisabled(t *testing.T) { + maxInbound := uint32(0) + cmgr, err := New(&Config{ + Dial: mockDialer, + MaxInbound: &maxInbound, + }) + if err != nil { + t.Fatalf("New error: %v", err) + } + + conn := &mockConn{rAddr: &net.TCPAddr{ + IP: net.ParseIP("192.0.2.1"), Port: 1001, + }} + if _, ok := cmgr.limitInbound(conn); ok { + t.Fatal("inbound connection accepted with a zero limit") + } + if got := atomic.LoadUint64(&cmgr.inboundRejected); got != 1 { + t.Fatalf("unexpected rejection count: got %d, want 1", got) + } +} + +// TestInboundLimitConcurrent verifies that concurrent admission never exceeds +// the configured accepted-connection bound. +func TestInboundLimitConcurrent(t *testing.T) { + const attempts = 100 + maxInbound := uint32(8) + + cmgr, err := New(&Config{ + Dial: mockDialer, + MaxInbound: &maxInbound, + }) + if err != nil { + t.Fatalf("New error: %v", err) + } + + var ( + wg sync.WaitGroup + live int32 + maxLive int32 + ) + release := make(chan struct{}) + + for i := 0; i < attempts; i++ { + wg.Add(1) + go func(port int) { + defer wg.Done() + + conn, ok := cmgr.limitInbound(&mockConn{ + rAddr: &net.TCPAddr{ + IP: net.ParseIP("198.51.100.1"), + Port: port, + }, + }) + if !ok { + return + } + + current := atomic.AddInt32(&live, 1) + for { + observed := atomic.LoadInt32(&maxLive) + if current <= observed || atomic.CompareAndSwapInt32( + &maxLive, observed, current, + ) { + break + } + } + + <-release + atomic.AddInt32(&live, -1) + _ = conn.Close() + }(10000 + i) + } + + deadline := time.After(time.Second) + for len(cmgr.inboundSlots) < int(maxInbound) { + select { + case <-deadline: + t.Fatal("timed out waiting for inbound slots to fill") + case <-time.After(time.Millisecond): + } + } + close(release) + wg.Wait() + + if maxLive > int32(maxInbound) { + t.Fatalf("live inbound connections exceeded limit: got %d, want <= %d", + maxLive, maxInbound) + } + if got := len(cmgr.inboundSlots); got != 0 { + t.Fatalf("inbound slots leaked: got %d", got) + } +} + // TestStartStop tests that the connection manager starts and stops as // expected. func TestStartStop(t *testing.T) { @@ -664,3 +819,62 @@ out: cmgr.Stop() cmgr.Wait() } + +// TestListenerInboundLimit verifies excess sockets are rejected before the +// accept callback is started and capacity returns when an admitted socket +// closes. +func TestListenerInboundLimit(t *testing.T) { + maxInbound := uint32(1) + listener := newMockListener("127.0.0.1:8333") + accepted := make(chan net.Conn, 2) + + cmgr, err := New(&Config{ + Listeners: []net.Listener{listener}, + OnAccept: func(conn net.Conn) { accepted <- conn }, + MaxInbound: &maxInbound, + Dial: mockDialer, + }) + if err != nil { + t.Fatalf("New error: %v", err) + } + cmgr.Start() + defer func() { + cmgr.Stop() + cmgr.Wait() + }() + + go listener.Connect("127.0.0.1", 10001) + var first net.Conn + select { + case first = <-accepted: + case <-time.After(50 * time.Millisecond): + t.Fatal("first inbound callback did not run") + } + + go listener.Connect("127.0.0.1", 10002) + deadline := time.After(time.Second) + for atomic.LoadUint64(&cmgr.inboundRejected) != 1 { + select { + case <-deadline: + t.Fatal("connection above the limit was not rejected") + + default: + time.Sleep(time.Millisecond) + } + } + select { + case <-accepted: + t.Fatal("connection above the limit reached the accept callback") + + default: + } + + _ = first.Close() + go listener.Connect("127.0.0.1", 10003) + select { + case replacement := <-accepted: + _ = replacement.Close() + case <-time.After(50 * time.Millisecond): + t.Fatal("released inbound slot was not reused") + } +} diff --git a/go.mod b/go.mod index 577d70a8..78624c26 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/btcsuite/btcd -go 1.25 +go 1.25.0 require ( github.com/btcsuite/btcd/address/v2 v2.0.0 @@ -24,6 +24,7 @@ require ( github.com/syndtr/goleveldb v1.0.1-0.20210819022825-2ae1ddf74ef7 golang.org/x/crypto v0.40.0 golang.org/x/sys v0.35.0 + golang.org/x/time v0.15.0 pgregory.net/rapid v1.2.0 ) diff --git a/go.sum b/go.sum index dd7efb50..a7ea81ce 100644 --- a/go.sum +++ b/go.sum @@ -103,6 +103,8 @@ golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.27.0 h1:4fGWRpyh641NLlecmyl4LOe6yDdfaYNrGb2zdfo4JV4= golang.org/x/text v0.27.0/go.mod h1:1D28KMCvyooCX9hBiosv5Tz/+YLxj0j7XhWjpSUF7CU= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=