peer: close late connections after disconnect

In this commit, we serialize connection association with peer
disconnection. Previously, Disconnect could run before
AssociateConnection, close quit without a socket, and leave a later
connection published with no remaining path that could close it.

We now close connections handed to a peer that is already disconnecting,
and publish the connected flag only after the socket and connection time
are initialized. The tests pin both call orderings and race the
transitions to verify the socket closes exactly once.

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)

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