identify: Fix IdentifyWait when Connected events happen out of order

core/network/conn.go

@@ -33,6 +33,10 @@ type Conn interface {
 
 	// GetStreams returns all open streams over this conn.
 	GetStreams() []Stream
+
+	// IsClosed returns whether a connection is fully closed, so it can
+	// be garbage collected.
+	IsClosed() bool
 }
 
 // ConnectionState holds information about the connection.

p2p/host/basic/basic_host_test.go

@@ -40,9 +40,11 @@ func TestHostSimple(t *testing.T) {
 	h1, err := NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)
 	defer h1.Close()
+	h1.Start()
 	h2, err := NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)
 	defer h2.Close()
+	h2.Start()
 
 	h2pi := h2.Peerstore().PeerInfo(h2.ID())
 	require.NoError(t, h1.Connect(ctx, h2pi))
@@ -447,8 +449,10 @@ func TestNewDialOld(t *testing.T) {
 func TestNewStreamResolve(t *testing.T) {
 	h1, err := NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)
+	h1.Start()
 	h2, err := NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)
+	h2.Start()
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()

p2p/host/routed/routed_test.go

@@ -28,6 +28,7 @@ func TestRoutedHostConnectToObsoleteAddresses(t *testing.T) {
 	h1, err := basic.NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)
 	defer h1.Close()
+	h1.Start()
 
 	h2, err := basic.NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)

p2p/net/connmgr/connmgr_test.go

@@ -802,6 +802,7 @@ func (m mockConn) LocalMultiaddr() ma.Multiaddr                          { panic
 func (m mockConn) RemoteMultiaddr() ma.Multiaddr                         { panic("implement me") }
 func (m mockConn) Stat() network.ConnStats                               { return m.stats }
 func (m mockConn) ID() string                                            { panic("implement me") }
+func (m mockConn) IsClosed() bool                                        { panic("implement me") }
 func (m mockConn) NewStream(ctx context.Context) (network.Stream, error) { panic("implement me") }
 func (m mockConn) GetStreams() []network.Stream                          { panic("implement me") }
 func (m mockConn) Scope() network.ConnScope                              { panic("implement me") }

p2p/net/mock/mock_conn.go

@@ -41,6 +41,8 @@ type conn struct {
 
 	closeOnce sync.Once
 
+	isClosed atomic.Bool
+
 	sync.RWMutex
 }
 
@@ -67,12 +69,17 @@ func newConn(ln, rn *peernet, l *link, dir network.Direction) *conn {
 	return c
 }
 
+func (c *conn) IsClosed() bool {
+	return c.isClosed.Load()
+}
+
 func (c *conn) ID() string {
 	return strconv.FormatInt(c.id, 10)
 }
 
 func (c *conn) Close() error {
 	c.closeOnce.Do(func() {
+		c.isClosed.Store(true)
 		go c.rconn.Close()
 		c.teardown()
 	})

p2p/net/swarm/swarm_conn.go

@@ -43,6 +43,10 @@ type Conn struct {
 
 var _ network.Conn = &Conn{}
 
+func (c *Conn) IsClosed() bool {
+	return c.conn.IsClosed()
+}
+
 func (c *Conn) ID() string {
 	// format: <first 10 chars of peer id>-<global conn ordinal>
 	return fmt.Sprintf("%s-%d", c.RemotePeer().Pretty()[0:10], c.id)

p2p/protocol/identify/id.go

@@ -349,11 +349,17 @@ func (ids *idService) IdentifyWait(c network.Conn) <-chan struct{} {
 	defer ids.connsMu.Unlock()
 
 	e, found := ids.conns[c]
-	if !found { // No entry found. Connection was most likely closed (and removed from this map) recently.
-		log.Debugw("connection not found in identify service", "peer", c.RemotePeer())
-		ch := make(chan struct{})
-		close(ch)
-		return ch
+	if !found {
+		// No entry found. We may have gotten an out of order notification. Check it we should have this conn (because we're still connected)
+		// We hold the ids.connsMu lock so this is safe since a disconnect event will be processed later if we are connected.
+		if c.IsClosed() {
+			log.Debugw("connection not found in identify service", "peer", c.RemotePeer())
+			ch := make(chan struct{})
+			close(ch)
+			return ch
+		} else {
+			ids.addConnWithLock(c)
+		}
 	}
 
 	if e.IdentifyWaitChan != nil {
@@ -863,6 +869,15 @@ func (ids *idService) consumeObservedAddress(observed []byte, c network.Conn) {
 	ids.observedAddrs.Record(c, maddr)
 }
 
+// addConnWithLock assuems caller holds the connsMu lock
+func (ids *idService) addConnWithLock(c network.Conn) {
+	_, found := ids.conns[c]
+	if !found {
+		<-ids.setupCompleted
+		ids.conns[c] = entry{}
+	}
+}
+
 func signedPeerRecordFromMessage(msg *pb.Identify) (*record.Envelope, error) {
 	if msg.SignedPeerRecord == nil || len(msg.SignedPeerRecord) == 0 {
 		return nil, nil
@@ -883,10 +898,8 @@ func (nn *netNotifiee) Connected(_ network.Network, c network.Conn) {
 	// The swarm implementation guarantees this.
 	ids := nn.IDService()
 
-	<-ids.setupCompleted
-
 	ids.connsMu.Lock()
-	ids.conns[c] = entry{}
+	ids.addConnWithLock(c)
 	ids.connsMu.Unlock()
 
 	nn.IDService().IdentifyWait(c)

p2p/protocol/identify/id_test.go

@@ -2,6 +2,7 @@ package identify_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 	"testing"
@@ -17,6 +18,7 @@ import (
 	"github.com/libp2p/go-libp2p/core/protocol"
 	"github.com/libp2p/go-libp2p/core/record"
 	coretest "github.com/libp2p/go-libp2p/core/test"
+	basichost "github.com/libp2p/go-libp2p/p2p/host/basic"
 	blhost "github.com/libp2p/go-libp2p/p2p/host/blank"
 	"github.com/libp2p/go-libp2p/p2p/host/eventbus"
 	"github.com/libp2p/go-libp2p/p2p/host/peerstore/pstoremem"
@@ -894,6 +896,46 @@ func TestIncomingIDStreamsTimeout(t *testing.T) {
 	}
 }
 
+func TestOutOfOrderConnectedNotifs(t *testing.T) {
+	h1, err := libp2p.New(libp2p.NoListenAddrs)
+	require.NoError(t, err)
+	h2, err := libp2p.New(libp2p.ListenAddrs(ma.StringCast("/ip4/127.0.0.1/udp/0/quic-v1")))
+	require.NoError(t, err)
+
+	doneCh := make(chan struct{})
+	errCh := make(chan error)
+
+	// This callback may be called before identify's Connnected callback completes. If it does, the IdentifyWait should still finish successfully.
+	h1.Network().Notify(&network.NotifyBundle{
+		ConnectedF: func(n network.Network, c network.Conn) {
+			bh1 := h1.(*basichost.BasicHost)
+			idChan := bh1.IDService().IdentifyWait(c)
+			go func() {
+				<-idChan
+				protos, err := bh1.Peerstore().GetProtocols(h2.ID())
+				if err != nil {
+					errCh <- err
+				}
+				if len(protos) == 0 {
+					errCh <- errors.New("no protocols found. Identify did not complete")
+				}
+
+				close(doneCh)
+			}()
+		},
+	})
+
+	h1.Connect(context.Background(), peer.AddrInfo{ID: h2.ID(), Addrs: h2.Addrs()})
+
+	select {
+	case <-doneCh:
+	case err := <-errCh:
+		t.Fatalf("err: %v", err)
+	case <-time.After(3 * time.Second):
+		t.Fatalf("identify wait never completed")
+	}
+}
+
 func waitForAddrInStream(t *testing.T, s <-chan ma.Multiaddr, expected ma.Multiaddr, timeout time.Duration, failMsg string) {
 	t.Helper()
 	for {

p2p/protocol/ping/ping_test.go

@@ -20,9 +20,11 @@ func TestPing(t *testing.T) {
 	h1, err := bhost.NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)
 	defer h1.Close()
+	h1.Start()
 	h2, err := bhost.NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)
 	defer h2.Close()
+	h2.Start()
 
 	err = h1.Connect(ctx, peer.AddrInfo{
 		ID:    h2.ID(),

p2p/test/backpressure/backpressure_test.go

@@ -25,8 +25,10 @@ func TestStBackpressureStreamWrite(t *testing.T) {
 
 	h1, err := bhost.NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)
+	h1.Start()
 	h2, err := bhost.NewHost(swarmt.GenSwarm(t), nil)
 	require.NoError(t, err)
+	h2.Start()
 
 	// setup sender handler on 2
 	h2.SetStreamHandler(protocol.TestingID, func(s network.Stream) {

p2p/test/reconnects/reconnect_test.go

@@ -37,6 +37,7 @@ func TestReconnect5(t *testing.T) {
 			h, err := bhost.NewHost(swarmt.GenSwarm(t, swarmOpt), nil)
 			require.NoError(t, err)
 			defer h.Close()
+			h.Start()
 			hosts = append(hosts, h)
 			h.SetStreamHandler(protocol.TestingID, EchoStreamHandler)
 		}