diff --git a/peer/peer.go b/peer/peer.go index ca057b30..cf950a2f 100644 --- a/peer/peer.go +++ b/peer/peer.go @@ -451,7 +451,9 @@ type Peer struct { connected int32 disconnect int32 - conn net.Conn + // connMtx serializes connection association with disconnection. + connMtx sync.Mutex + conn net.Conn // These fields are set at creation time and never modified, so they are // safe to read from concurrently without a mutex. @@ -2013,8 +2015,11 @@ func (p *Peer) Disconnect() { } log.Tracef("Disconnecting %s", p) - if atomic.LoadInt32(&p.connected) != 0 { - p.conn.Close() + p.connMtx.Lock() + conn := p.conn + p.connMtx.Unlock() + if conn != nil { + _ = conn.Close() } close(p.quit) } @@ -2446,16 +2451,27 @@ func (p *Peer) start() error { return nil } -// AssociateConnection associates the given conn to the peer. Calling this -// function when the peer is already connected will have no effect. +// AssociateConnection associates the given conn to the peer. Calling this +// function when the peer is already connected will have no effect. When the +// peer is already disconnecting, the connection is closed instead. func (p *Peer) AssociateConnection(conn net.Conn) { - // Already connected? - if !atomic.CompareAndSwapInt32(&p.connected, 0, 1) { + p.connMtx.Lock() + if atomic.LoadInt32(&p.connected) != 0 { + p.connMtx.Unlock() + return + } + if atomic.LoadInt32(&p.disconnect) != 0 { + p.connMtx.Unlock() + _ = conn.Close() return } p.conn = conn + p.statsMtx.Lock() p.timeConnected = time.Now() + p.statsMtx.Unlock() + atomic.StoreInt32(&p.connected, 1) + p.connMtx.Unlock() if p.cfg.UsingV2Conn { p.V2Transport.UseReadWriter(conn) diff --git a/peer/peer_test.go b/peer/peer_test.go index 2903cb93..e7a66584 100644 --- a/peer/peer_test.go +++ b/peer/peer_test.go @@ -10,6 +10,7 @@ import ( "io" "net" "strconv" + "sync" "sync/atomic" "testing" "time" @@ -48,6 +49,18 @@ type conn struct { proxy bool } +// countingConn records how many times its embedded connection is closed. +type countingConn struct { + net.Conn + closes int32 +} + +// Close closes the embedded connection and records the call. +func (c *countingConn) Close() error { + atomic.AddInt32(&c.closes, 1) + return c.Conn.Close() +} + // LocalAddr returns the local address for the connection. func (c conn) LocalAddr() net.Addr { return &addr{c.lnet, c.laddr} @@ -79,6 +92,81 @@ func (c conn) SetDeadline(t time.Time) error { return nil } func (c conn) SetReadDeadline(t time.Time) error { return nil } func (c conn) SetWriteDeadline(t time.Time) error { return nil } +// TestDisconnectBeforeAssociateConnection verifies a connection handed to an +// already-disconnected peer is closed instead of being published and leaked. +func TestDisconnectBeforeAssociateConnection(t *testing.T) { + local, remote := net.Pipe() + defer local.Close() + defer remote.Close() + + trackedConn := &countingConn{Conn: local} + p := peer.NewInboundPeer(&peer.Config{}) + p.Disconnect() + p.Disconnect() + p.WaitForDisconnect() + + p.AssociateConnection(trackedConn) + + if got := atomic.LoadInt32(&trackedConn.closes); got != 1 { + t.Fatalf("unexpected connection close count: got %d, want 1", got) + } + if p.Connected() { + t.Fatal("disconnected peer accepted a connection") + } +} + +// TestAssociateConnectionDisconnectRace verifies concurrent association and +// disconnection always close the transferred connection exactly once. +func TestAssociateConnectionDisconnectRace(t *testing.T) { + const iterations = 100 + + for i := 0; i < iterations; i++ { + local, remote := net.Pipe() + trackedConn := &countingConn{Conn: local} + p, err := peer.NewOutboundPeer( + &peer.Config{ + NewestBlock: func() (*chainhash.Hash, int32, error) { + return &chainhash.Hash{}, 0, nil + }, + AllowSelfConns: true, + }, + "127.0.0.1:8333", + ) + if err != nil { + t.Fatalf("NewOutboundPeer: unexpected error: %v", err) + } + + start := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + <-start + p.AssociateConnection(trackedConn) + }() + go func() { + defer wg.Done() + <-start + p.Disconnect() + }() + + close(start) + wg.Wait() + p.WaitForDisconnect() + + if got := atomic.LoadInt32(&trackedConn.closes); got != 1 { + t.Fatalf("iteration %d: unexpected connection close count: "+ + "got %d, want 1", i, got) + } + if p.Connected() { + t.Fatalf("iteration %d: peer remained connected", i) + } + + _ = local.Close() + _ = remote.Close() + } +} + // addr mocks a network address type addr struct { net, address string