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=