package identify import ( "context" "fmt" "testing" "time" "github.com/libp2p/go-libp2p/core/network" "github.com/libp2p/go-libp2p/core/peer" "github.com/libp2p/go-libp2p/core/peerstore" recordPb "github.com/libp2p/go-libp2p/core/record/pb" blhost "github.com/libp2p/go-libp2p/p2p/host/blank" swarmt "github.com/libp2p/go-libp2p/p2p/net/swarm/testing" "google.golang.org/protobuf/proto" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestFastDisconnect(t *testing.T) { // This test checks to see if we correctly abort sending an identify // response if the peer disconnects before we handle the request. ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() target := blhost.NewBlankHost(swarmt.GenSwarm(t)) defer target.Close() ids, err := NewIDService(target) require.NoError(t, err) defer ids.Close() ids.Start() sync := make(chan struct{}) target.SetStreamHandler(ID, func(s network.Stream) { // Wait till the stream is set up on both sides. select { case <-sync: case <-ctx.Done(): return } // Kill the connection, and make sure we're completely disconnected. assert.Eventually(t, func() bool { for _, conn := range target.Network().ConnsToPeer(s.Conn().RemotePeer()) { conn.Close() } return target.Network().Connectedness(s.Conn().RemotePeer()) != network.Connected }, 2*time.Second, time.Millisecond, ) // Now try to handle the response. // This should not block indefinitely, or panic, or anything like that. // // However, if we have a bug, that _could_ happen. ids.handleIdentifyRequest(s) // Ok, allow the outer test to continue. select { case <-sync: case <-ctx.Done(): return } }) source := blhost.NewBlankHost(swarmt.GenSwarm(t)) defer source.Close() // only connect to the first address, to make sure we only end up with one connection require.NoError(t, source.Connect(ctx, peer.AddrInfo{ID: target.ID(), Addrs: target.Addrs()})) s, err := source.NewStream(ctx, target.ID(), ID) require.NoError(t, err) select { case sync <- struct{}{}: case <-ctx.Done(): t.Fatal(ctx.Err()) } s.Reset() select { case sync <- struct{}{}: case <-ctx.Done(): t.Fatal(ctx.Err()) } // double-check to make sure we didn't actually timeout somewhere. require.NoError(t, ctx.Err()) } func TestWrongSignedPeerRecord(t *testing.T) { h1 := blhost.NewBlankHost(swarmt.GenSwarm(t)) defer h1.Close() ids, err := NewIDService(h1) require.NoError(t, err) ids.Start() defer ids.Close() h2 := blhost.NewBlankHost(swarmt.GenSwarm(t)) defer h2.Close() ids2, err := NewIDService(h2) require.NoError(t, err) ids2.Start() defer ids2.Close() h3 := blhost.NewBlankHost(swarmt.GenSwarm(t)) defer h2.Close() ids3, err := NewIDService(h3) require.NoError(t, err) ids3.Start() defer ids3.Close() h2.Connect(context.Background(), peer.AddrInfo{ID: h1.ID(), Addrs: h1.Addrs()}) s, err := h2.NewStream(context.Background(), h1.ID(), IDPush) require.NoError(t, err) err = ids3.sendIdentifyResp(s, true) // This should fail because the peer record is signed by h3, not h2 require.NoError(t, err) time.Sleep(time.Second) require.Empty(t, h1.Peerstore().Addrs(h3.ID()), "h1 should not know about h3 since it was relayed over h2") } func TestInvalidSignedPeerRecord(t *testing.T) { h1 := blhost.NewBlankHost(swarmt.GenSwarm(t)) defer h1.Close() ids, err := NewIDService(h1) require.NoError(t, err) ids.Start() defer ids.Close() h2 := blhost.NewBlankHost(swarmt.GenSwarm(t)) defer h2.Close() ids2, err := NewIDService(h2) require.NoError(t, err) // We don't want to start the identify service, we'll manage the messages h2 // sends manually so we can tweak it // ids2.Start() h2.Connect(context.Background(), peer.AddrInfo{ID: h1.ID(), Addrs: h1.Addrs()}) require.Empty(t, h1.Peerstore().Addrs(h2.ID())) s, err := h2.NewStream(context.Background(), h1.ID(), IDPush) require.NoError(t, err) ids2.updateSnapshot() ids2.currentSnapshot.Lock() snapshot := ids2.currentSnapshot.snapshot ids2.currentSnapshot.Unlock() mes := ids2.createBaseIdentifyResponse(s.Conn(), &snapshot) fmt.Println("Signed record is", snapshot.record) marshalled, err := snapshot.record.Marshal() require.NoError(t, err) var envPb recordPb.Envelope err = proto.Unmarshal(marshalled, &envPb) require.NoError(t, err) envPb.Signature = []byte("invalid") mes.SignedPeerRecord, err = proto.Marshal(&envPb) require.NoError(t, err) err = ids2.writeChunkedIdentifyMsg(s, mes) require.NoError(t, err) fmt.Println("Done sending msg") s.Close() // Wait a bit for h1 to process the message time.Sleep(1 * time.Second) cab, ok := h1.Peerstore().(peerstore.CertifiedAddrBook) require.True(t, ok) require.Nil(t, cab.GetPeerRecord(h2.ID())) }