fix infinite discovery address leak

zeta-chain · Nov 15, 2024 · ab13fd6 · ab13fd6
1 parent 12f139f
commit ab13fd6
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/p2p/discovery.go b/p2p/discovery.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"io"
-	"slices"
 	"sync"
 	"time"
 
@@ -85,7 +84,7 @@ func (pd *PeerDiscovery) addPeer(pinfo peer.AddrInfo) {
 	oldPinfo, ok := pd.knownPeers[pinfo.ID]
 	if ok {
 		for _, addr := range pinfo.Addrs {
-			if !slices.Contains(oldPinfo.Addrs, addr) {
+			if !multiaddr.Contains(oldPinfo.Addrs, addr) {
 				oldPinfo.Addrs = append(oldPinfo.Addrs, addr)
 			}
 		}

diff --git a/p2p/discovery_test.go b/p2p/discovery_test.go
@@ -83,4 +83,9 @@ func TestDiscovery(t *testing.T) {
 	assert.Equal(t, 4, len(comm2.host.Peerstore().Peers()))
 	assert.Equal(t, 4, len(comm3.host.Peerstore().Peers()))
 	assert.Equal(t, 4, len(comm4.host.Peerstore().Peers()))
+
+	assert.Equal(t, 3, len(comm.discovery.knownPeers))
+	for peer, knownPeers := range comm.discovery.knownPeers {
+		assert.LessOrEqual(t, len(knownPeers.Addrs), 4, "%s has more than 4 addresses (%d)?", peer.String(), len(knownPeers.Addrs))
+	}
 }