cached_addr_book.go

package main

import (
	"context"
	"io"
	"sync"
	"sync/atomic"
	"time"

	lru "github.com/hashicorp/golang-lru/v2"
	"github.com/ipfs/boxo/routing/http/types"
	"github.com/libp2p/go-libp2p-kad-dht/amino"
	"github.com/libp2p/go-libp2p/core/event"
	"github.com/libp2p/go-libp2p/core/host"
	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/peer"
	"github.com/libp2p/go-libp2p/core/peerstore"
	"github.com/libp2p/go-libp2p/core/record"
	"github.com/libp2p/go-libp2p/p2p/host/peerstore/pstoremem"
	"github.com/libp2p/go-libp2p/p2p/protocol/circuitv2/relay"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr/net"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

const (
	Subsystem = "cached_addr_book"

	// The default TTL to keep recently connected peers' multiaddrs for
	DefaultRecentlyConnectedAddrTTL = amino.DefaultProvideValidity

	// Connected peers don't expire until they disconnect
	ConnectedAddrTTL = peerstore.ConnectedAddrTTL

	// How long to wait since last connection before probing a peer again
	PeerProbeThreshold = time.Hour

	// How often to run the probe peers loop
	ProbeInterval = time.Minute * 15

	// How many concurrent probes to run at once
	MaxConcurrentProbes = 20

	// How long to wait for a connect in a probe to complete.
	// The worst case is a peer behind a relay, so we use the relay connect timeout.
	ConnectTimeout = relay.ConnectTimeout

	// How many peers to cache in the peer state cache
	// 1_000_000 is 10x the default number of signed peer records cached by the memory address book.
	PeerCacheSize = 1_000_000

	// Maximum backoff duration for probing a peer. After this duration, we will stop
	// trying to connect to the peer and remove it from the cache.
	MaxBackoffDuration = amino.DefaultProvideValidity

	probeResult        = "result"
	probeResultOnline  = "online"
	probeResultOffline = "offline"
)

// DefaultRelayAddrTTL bounds how long someguy serves a cached /p2p-circuit
// (relay) address. A relay reservation lasts at most the relay's reservation
// TTL (relay.DefaultResources().ReservationTTL) and is dropped the instant the
// reserving peer disconnects from the relay, so a relay address is far more
// perishable than a direct one. Caching it for the full
// DefaultRecentlyConnectedAddrTTL would keep handing clients relay paths that
// died hours ago. The probe loop re-extends this TTL for peers that are still
// reachable, so live relay-only peers survive while dead relays age out
// quickly. It is set to twice the reservation TTL so the probe loop has room to
// re-confirm a live peer before its entry expires.
var DefaultRelayAddrTTL = 2 * relay.DefaultResources().ReservationTTL

var (
	probeDurationHistogram = promauto.NewHistogram(prometheus.HistogramOpts{
		Name:      "probe_duration_seconds",
		Namespace: name,
		Subsystem: Subsystem,
		Help:      "Duration of peer probing operations in seconds",
		// Buckets probe durations from 5s to 15 minutes
		Buckets: []float64{5, 10, 30, 60, 120, 300, 600, 900},
	})

	probedPeersCounter = promauto.NewCounterVec(prometheus.CounterOpts{
		Name:      "probed_peers",
		Subsystem: Subsystem,
		Namespace: name,
		Help:      "Number of peers probed",
	},
		[]string{probeResult},
	)

	peerStateSize = promauto.NewGauge(prometheus.GaugeOpts{
		Name:      "peer_state_size",
		Subsystem: Subsystem,
		Namespace: name,
		Help:      "Number of peers object currently in the peer state",
	})
)

type peerState struct {
	lastConnTime       time.Time // last time we successfully connected to this peer
	lastFailedConnTime time.Time // last time we failed to find or connect to this peer
	connectFailures    uint      // number of times we've failed to connect to this peer
}

type cachedAddrBook struct {
	addrBook             peerstore.AddrBook             // someguy's own address book: durable, probed, written here
	hostPeerstore        peerstore.AddrBook             // libp2p host peerstore, DHT-populated, read-only fallback
	peerCache            *lru.Cache[peer.ID, peerState] // LRU cache with additional metadata about peer
	probingEnabled       bool
	isProbing            atomic.Bool
	allowPrivateIPs      bool // for testing
	recentlyConnectedTTL time.Duration
	relayAddrTTL         time.Duration
}

type AddrBookOption func(*cachedAddrBook) error

// WithHostPeerstore lets GetCachedAddrs fall back to the libp2p host peerstore,
// which go-libp2p-kad-dht populates with provider addresses during
// FindProviders (under a short TempAddrTTL). This catches peers seen very
// recently as providers that have not yet been copied into someguy's own
// longer-lived address book.
func WithHostPeerstore(ps peerstore.AddrBook) AddrBookOption {
	return func(cab *cachedAddrBook) error {
		cab.hostPeerstore = ps
		return nil
	}
}

func WithAllowPrivateIPs() AddrBookOption {
	return func(cab *cachedAddrBook) error {
		cab.allowPrivateIPs = true
		return nil
	}
}

func WithRecentlyConnectedTTL(ttl time.Duration) AddrBookOption {
	return func(cab *cachedAddrBook) error {
		cab.recentlyConnectedTTL = ttl
		return nil
	}
}

// WithRelayAddrTTL overrides the TTL used for /p2p-circuit (relay) addresses.
// See DefaultRelayAddrTTL for why these are kept shorter than direct addresses.
func WithRelayAddrTTL(ttl time.Duration) AddrBookOption {
	return func(cab *cachedAddrBook) error {
		cab.relayAddrTTL = ttl
		return nil
	}
}

func WithActiveProbing(enabled bool) AddrBookOption {
	return func(cab *cachedAddrBook) error {
		cab.probingEnabled = enabled
		return nil
	}
}

func newCachedAddrBook(opts ...AddrBookOption) (*cachedAddrBook, error) {
	peerCache, err := lru.New[peer.ID, peerState](PeerCacheSize)
	if err != nil {
		return nil, err
	}

	cab := &cachedAddrBook{
		peerCache:            peerCache,
		addrBook:             pstoremem.NewAddrBook(),
		recentlyConnectedTTL: DefaultRecentlyConnectedAddrTTL, // Set default value
		relayAddrTTL:         DefaultRelayAddrTTL,             // Set default value
	}

	for _, opt := range opts {
		err := opt(cab)
		if err != nil {
			return nil, err
		}
	}
	logger.Infof("Using TTL of %s for recently connected peers", cab.recentlyConnectedTTL)
	logger.Infof("Using TTL of %s for relay (/p2p-circuit) addresses", cab.relayAddrTTL)
	logger.Infof("Probing enabled: %t", cab.probingEnabled)
	return cab, nil
}

func (cab *cachedAddrBook) background(ctx context.Context, host host.Host) {
	sub, err := host.EventBus().Subscribe([]any{
		&event.EvtPeerIdentificationCompleted{},
		&event.EvtPeerConnectednessChanged{},
	})
	if err != nil {
		logger.Errorf("failed to subscribe to peer identification events: %v", err)
		return
	}
	defer sub.Close()

	probeTicker := time.NewTicker(ProbeInterval)
	defer probeTicker.Stop()

	for {
		select {
		case <-ctx.Done():
			cabCloser, ok := cab.addrBook.(io.Closer)
			if ok {
				errClose := cabCloser.Close()
				if errClose != nil {
					logger.Warnf("failed to close addr book: %v", errClose)
				}
			}
			return
		case ev := <-sub.Out():
			switch ev := ev.(type) {
			case event.EvtPeerIdentificationCompleted:
				pState, exists := cab.peerCache.Peek(ev.Peer)
				if !exists {
					pState = peerState{}
				}
				pState.lastConnTime = time.Now()
				pState.lastFailedConnTime = time.Time{} // reset failed connection time
				pState.connectFailures = 0              // reset connect failures on successful connection
				cab.peerCache.Add(ev.Peer, pState)
				peerStateSize.Set(float64(cab.peerCache.Len())) // update metric

				ttl := cab.getTTL(host.Network().Connectedness(ev.Peer))

				// A completed identify reports the peer's current advertised
				// addresses, which supersede the set accumulated from provider
				// records, DHT gossip, and earlier identifies. Replace the
				// stored set instead of unioning so stale certhashes, dead
				// relay circuits, and rotated NAT ports do not pile up.
				//
				// Drop the remote addresses of inbound connections: that is the
				// peer's ephemeral source port, which nobody can dial back to,
				// so caching it would reintroduce exactly the junk this prune
				// removes. Outbound (and direction-unknown) remotes are kept.
				var connAddrs []ma.Multiaddr
				for _, c := range host.Network().ConnsToPeer(ev.Peer) {
					if c.Stat().Direction != network.DirInbound {
						connAddrs = append(connAddrs, c.RemoteMultiaddr())
					}
				}
				cab.replacePeerAddrs(ev.Peer, ev.SignedPeerRecord, ev.ListenAddrs, connAddrs, ttl)
			case event.EvtPeerConnectednessChanged:
				// On disconnect, move the peer's addresses off the connected TTL,
				// then cap its relay addresses so a now-idle relay path is not
				// served for the full recentlyConnectedTTL.
				if !hasValidConnectedness(ev.Connectedness) {
					cab.addrBook.UpdateAddrs(ev.Peer, ConnectedAddrTTL, cab.recentlyConnectedTTL)
					cab.capRelayAddrTTL(ev.Peer)
				}
			}
		case <-probeTicker.C:
			if !cab.probingEnabled {
				logger.Debug("Probing disabled, skipping")
				continue
			}
			if cab.isProbing.Load() {
				logger.Debug("Skipping peer probe, still running")
				continue
			}
			logger.Debug("Starting to probe peers")
			cab.isProbing.Store(true)
			go cab.probePeers(ctx, host)
		}
	}
}

// replacePeerAddrs replaces p's stored addresses with the authoritative set
// from a completed identify: the signed peer record, else the identify listen
// addresses, plus any live-connection address so an active session is kept.
//
// Clearing first drops addresses absent from the current set (stale certhashes,
// dead relay circuits, rotated NAT ports) instead of letting them linger to TTL.
//
// libp2p/go-libp2p#3487 does the same prune inside ConsumePeerRecord (not yet
// in the pinned version); re-adding the same set here stays correct once it
// lands, so a dependency bump will not regress this.
func (cab *cachedAddrBook) replacePeerAddrs(p peer.ID, signed *record.Envelope, listenAddrs, connAddrs []ma.Multiaddr, ttl time.Duration) {
	// Nothing authoritative to apply. Return before clearing so an identify that
	// carried no usable addresses never wipes a peer's existing cached set.
	if signed == nil && len(listenAddrs) == 0 && len(connAddrs) == 0 {
		return
	}

	// Drop the accumulated set so addresses absent from the current advertised
	// set are removed instead of unioned.
	cab.addrBook.ClearAddrs(p)

	accepted := false
	if signed != nil {
		if certBook, ok := peerstore.GetCertifiedAddrBook(cab.addrBook); ok {
			ok, err := certBook.ConsumePeerRecord(signed, ttl)
			if err != nil {
				logger.Warnf("failed to consume signed peer record: %v", err)
			}
			accepted = ok
		}
	}
	if !accepted {
		// No signed record, no certified addr book, or the record was rejected
		// (e.g. a sequence-number check in some go-libp2p version). Fall back to
		// the identify listen addresses so the clear never leaves the peer with
		// zero addresses.
		cab.addrBook.AddAddrs(p, listenAddrs, ttl)
	}

	// Cap relay (/p2p-circuit) addresses at the shorter relayAddrTTL. The
	// advertised set is freshly verified, but a relay reservation can lapse long
	// before ttl, so a relay path should not inherit the full TTL. Run this
	// before re-adding connAddrs so a live relay session keeps the connected TTL.
	cab.capRelayAddrTTL(p)

	// Preserve live-connection addresses at the connected TTL even when absent
	// from the advertised set, so an active session is never dropped.
	if len(connAddrs) > 0 {
		cab.addrBook.AddAddrs(p, connAddrs, ConnectedAddrTTL)
	}
}

// Loops over all peers with addresses and probes them if they haven't been probed recently
func (cab *cachedAddrBook) probePeers(ctx context.Context, host host.Host) {
	defer cab.isProbing.Store(false)

	start := time.Now()
	defer func() {
		duration := time.Since(start).Seconds()
		probeDurationHistogram.Observe(duration)
		logger.Debugf("Finished probing peers in %s", duration)
	}()

	var wg sync.WaitGroup
	// semaphore channel to limit the number of concurrent probes
	semaphore := make(chan struct{}, MaxConcurrentProbes)

	for i, p := range cab.addrBook.PeersWithAddrs() {
		if hasValidConnectedness(host.Network().Connectedness(p)) {
			continue // don't probe connected peers
		}

		if !cab.ShouldProbePeer(p) {
			continue
		}

		addrs := cab.addrBook.Addrs(p)

		if !cab.allowPrivateIPs {
			addrs = ma.FilterAddrs(addrs, manet.IsPublicAddr)
		}

		if len(addrs) == 0 {
			continue // no addresses to probe
		}

		wg.Add(1)
		semaphore <- struct{}{}
		go func() {
			defer func() {
				<-semaphore // Release semaphore
				wg.Done()
			}()
			ctx, cancel := context.WithTimeout(ctx, ConnectTimeout)
			defer cancel()
			logger.Debugf("Probe %d: PeerID: %s, Addrs: %v", i+1, p, addrs)
			// if connect succeeds and identify runs, the background loop will take care of updating the peer state and cache
			err := host.Connect(ctx, peer.AddrInfo{
				ID:    p,
				Addrs: addrs,
			})
			if err != nil {
				logger.Debugf("failed to connect to peer %s: %v", p, err)
				cab.RecordFailedConnection(p)
				probedPeersCounter.WithLabelValues(probeResultOffline).Inc()
			} else {
				probedPeersCounter.WithLabelValues(probeResultOnline).Inc()
			}
		}()
	}
	wg.Wait()
}

// Returns the cached addresses for a peer, incrementing the return count
func (cab *cachedAddrBook) GetCachedAddrs(p peer.ID) []types.Multiaddr {
	cachedAddrs := cab.addrBook.Addrs(p)

	// Fall back to the host peerstore, which the DHT fills with provider
	// addresses during FindProviders (short TempAddrTTL). Lets peer routing
	// serve a peer seen as a provider moments ago but absent from peer routing.
	if len(cachedAddrs) == 0 && cab.hostPeerstore != nil {
		cachedAddrs = cab.hostPeerstore.Addrs(p)
	}

	if len(cachedAddrs) == 0 {
		return nil
	}

	result := make([]types.Multiaddr, 0, len(cachedAddrs)) // convert to local Multiaddr type 🙃
	for _, addr := range cachedAddrs {
		result = append(result, types.Multiaddr{Multiaddr: addr})
	}
	return result
}

// CacheAddrs stores addresses observed for a peer outside of a direct
// connection (e.g. embedded in a provider record returned by FindProviders) so
// that later peer-routing lookups can serve them from the same peerbook.
// Private addresses are dropped unless explicitly allowed. These addresses are
// unverified, so direct addresses are stored with the recently-connected TTL
// and relay (/p2p-circuit) addresses with the shorter relayAddrTTL; the probe
// loop confirms or evicts them.
func (cab *cachedAddrBook) CacheAddrs(p peer.ID, addrs []types.Multiaddr) {
	if len(addrs) == 0 {
		return
	}

	maddrs := make([]ma.Multiaddr, 0, len(addrs))
	for _, addr := range addrs {
		if !cab.allowPrivateIPs && !manet.IsPublicAddr(addr.Multiaddr) {
			continue
		}
		maddrs = append(maddrs, addr.Multiaddr)
	}

	if len(maddrs) == 0 {
		return
	}

	// Relay (/p2p-circuit) addresses are far more perishable than direct ones,
	// so cache them under the shorter relayAddrTTL. AddAddrs only ever extends a
	// TTL, so adding the relay set here never shortens one already held by a live
	// connection.
	direct, relayAddrs := splitRelayAddrs(maddrs)
	cab.addrBook.AddAddrs(p, direct, cab.recentlyConnectedTTL)
	cab.addrBook.AddAddrs(p, relayAddrs, cab.relayAddrTTL)
}

// Update the peer cache with information about a failed connection
// This should be called when a connection attempt to a peer fails
func (cab *cachedAddrBook) RecordFailedConnection(p peer.ID) {
	pState, exists := cab.peerCache.Peek(p)
	if !exists {
		pState = peerState{}
	}
	now := time.Now()
	// once probing of offline peer reached MaxBackoffDuration and still failed,
	// we opportunistically remove the dead peer from cache to save time on probing it further
	if exists && pState.connectFailures > 1 && now.Sub(pState.lastFailedConnTime) > MaxBackoffDuration {
		cab.peerCache.Remove(p)
		peerStateSize.Set(float64(cab.peerCache.Len())) // update metric
		// remove the peer from the addr book. Otherwise it will be probed again in the probe loop
		cab.addrBook.ClearAddrs(p)
		return
	}
	pState.lastFailedConnTime = now
	pState.connectFailures++
	cab.peerCache.Add(p, pState)
}

// Returns true if we should probe a peer (either by dialing known addresses or by dispatching a FindPeer)
// based on the last failed connection time and connection failures
func (cab *cachedAddrBook) ShouldProbePeer(p peer.ID) bool {
	pState, exists := cab.peerCache.Peek(p)
	if !exists {
		return true // default to probing if the peer is not in the cache
	}

	var backoffDuration time.Duration
	if pState.connectFailures > 0 {
		// Calculate backoff only if we have failures
		// this is effectively 2^(connectFailures - 1) * PeerProbeThreshold
		// A single failure results in a 1 hour backoff and each additional failure doubles the backoff
		backoffDuration = PeerProbeThreshold * time.Duration(1<<(pState.connectFailures-1))
		backoffDuration = min(backoffDuration, MaxBackoffDuration) // clamp to max backoff duration
	} else {
		backoffDuration = PeerProbeThreshold
	}

	// Only dispatch if we've waited long enough based on the backoff
	return time.Since(pState.lastFailedConnTime) > backoffDuration
}

func hasValidConnectedness(connectedness network.Connectedness) bool {
	return connectedness == network.Connected || connectedness == network.Limited
}

func (cab *cachedAddrBook) getTTL(connectedness network.Connectedness) time.Duration {
	if hasValidConnectedness(connectedness) {
		return ConnectedAddrTTL
	}
	return cab.recentlyConnectedTTL
}

// isRelayAddr reports whether a is a circuit-relay (/p2p-circuit) address.
func isRelayAddr(a ma.Multiaddr) bool {
	_, err := a.ValueForProtocol(ma.P_CIRCUIT)
	return err == nil
}

// splitRelayAddrs partitions addrs into direct addresses and circuit-relay
// (/p2p-circuit) addresses, preserving order within each group.
func splitRelayAddrs(addrs []ma.Multiaddr) (direct, relay []ma.Multiaddr) {
	for _, a := range addrs {
		if isRelayAddr(a) {
			relay = append(relay, a)
		} else {
			direct = append(direct, a)
		}
	}
	return direct, relay
}

// capRelayAddrTTL lowers the TTL of p's stored relay (/p2p-circuit) addresses to
// relayAddrTTL. It uses SetAddrs, which sets an exact TTL and so, unlike
// AddAddrs, can shorten an entry; a caller that must keep a live relay session
// re-adds it at the connected TTL afterward.
func (cab *cachedAddrBook) capRelayAddrTTL(p peer.ID) {
	if _, relayAddrs := splitRelayAddrs(cab.addrBook.Addrs(p)); len(relayAddrs) > 0 {
		cab.addrBook.SetAddrs(p, relayAddrs, cab.relayAddrTTL)
	}
}