feat: add options to cached addr book

fix test by allowing private ips

cached_addr_book.go

@@ -50,17 +50,35 @@ type peerState struct {
 }
 
 type cachedAddrBook struct {
-	addrBook  peerstore.AddrBook
-	peers     map[peer.ID]*peerState
-	mu        sync.RWMutex // Add mutex for thread safety
-	isProbing bool
+	addrBook        peerstore.AddrBook
+	peers           map[peer.ID]*peerState
+	mu              sync.RWMutex // Add mutex for thread safety
+	isProbing       bool
+	allowPrivateIPs bool // for testing
 }
 
-func newCachedAddrBook() *cachedAddrBook {
-	return &cachedAddrBook{
+type AddrBookOption func(*cachedAddrBook) error
+
+func WithAllowPrivateIPs() AddrBookOption {
+	return func(cab *cachedAddrBook) error {
+		cab.allowPrivateIPs = true
+		return nil
+	}
+}
+
+func newCachedAddrBook(opts ...AddrBookOption) (*cachedAddrBook, error) {
+	cab := &cachedAddrBook{
 		peers:    make(map[peer.ID]*peerState),
 		addrBook: pstoremem.NewAddrBook(),
 	}
+
+	for _, opt := range opts {
+		err := opt(cab)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return cab, nil
 }
 
 func (cab *cachedAddrBook) background(ctx context.Context, host host.Host) {
@@ -149,7 +167,8 @@ func (cab *cachedAddrBook) probePeers(ctx context.Context, host host.Host) {
 	semaphore := make(chan struct{}, MaxConcurrentProbes)
 
 	for i, p := range cab.addrBook.PeersWithAddrs() {
-		if host.Network().Connectedness(p) == network.Connected || host.Network().Connectedness(p) == network.Limited {
+		connectedness := host.Network().Connectedness(p)
+		if connectedness == network.Connected || connectedness == network.Limited {
 			continue // don't probe connected peers
 		}
 
@@ -164,7 +183,10 @@ func (cab *cachedAddrBook) probePeers(ctx context.Context, host host.Host) {
 		}
 
 		addrs := cab.addrBook.Addrs(p)
-		addrs = ma.FilterAddrs(addrs, manet.IsPublicAddr)
+
+		if !cab.allowPrivateIPs {
+			addrs = ma.FilterAddrs(addrs, manet.IsPublicAddr)
+		}
 
 		if len(addrs) == 0 {
 			continue // no addresses to probe

cached_addr_book_test.go

@@ -16,14 +16,16 @@ import (
 
 func TestCachedAddrBook(t *testing.T) {
 	// Create a new cached address book
-	cab := newCachedAddrBook()
+	cab, err := newCachedAddrBook(WithAllowPrivateIPs())
+	require.NoError(t, err)
 	require.NotNil(t, cab)
 	require.NotNil(t, cab.peers)
 	require.NotNil(t, cab.addrBook)
 }
 
 func TestGetCachedAddrs(t *testing.T) {
-	cab := newCachedAddrBook()
+	cab, err := newCachedAddrBook(WithAllowPrivateIPs())
+	require.NoError(t, err)
 
 	// Create a test peer with new PeerID
 	testPeer, err := peer.Decode("12D3KooWCZ67sU8oCvKd82Y6c9NgpqgoZYuZEUcg4upHCjK3n1aj")
@@ -59,7 +61,8 @@ func TestBackground(t *testing.T) {
 	require.NoError(t, err)
 	defer em.Close()
 
-	cab := newCachedAddrBook()
+	cab, err := newCachedAddrBook(WithAllowPrivateIPs())
+	require.NoError(t, err)
 
 	// Create a channel to signal when background processing is ready
 	ready := make(chan struct{})
@@ -135,7 +138,8 @@ func TestProbePeers(t *testing.T) {
 	require.NoError(t, err)
 	defer h.Close()
 
-	cab := newCachedAddrBook()
+	cab, err := newCachedAddrBook(WithAllowPrivateIPs())
+	require.NoError(t, err)
 
 	// Add a test peer with some addresses
 	testPeer, _ := peer.Decode("12D3KooWCZ67sU8oCvKd82Y6c9NgpqgoZYuZEUcg4upHCjK3n1aj")

server.go

@@ -85,7 +85,10 @@ func start(ctx context.Context, cfg *config) error {
 
 	if cfg.cachedAddrBook {
 		fmt.Println("Using cached address book to speed up peer discovery")
-		cachedAddrBook = newCachedAddrBook()
+		cachedAddrBook, err = newCachedAddrBook()
+		if err != nil {
+			return err
+		}
 		go cachedAddrBook.background(ctx, h)
 	}