Refactor Swarm.resolveAddrs

Refactors how DNS Address resolution works.

config/config.go

@@ -40,7 +40,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 
 	ma "github.com/multiformats/go-multiaddr"
-	madns "github.com/multiformats/go-multiaddr-dns"
 	manet "github.com/multiformats/go-multiaddr/net"
 	"github.com/quic-go/quic-go"
 	"go.uber.org/fx"
@@ -114,7 +113,7 @@ type Config struct {
 	Peerstore  peerstore.Peerstore
 	Reporter   metrics.Reporter
 
-	MultiaddrResolver *madns.Resolver
+	MultiaddrResolver swarm.MultiaddrDNSResolver
 
 	DisablePing bool
 
@@ -286,7 +285,6 @@ func (cfg *Config) addTransports() ([]fx.Option, error) {
 		fx.Provide(func() connmgr.ConnectionGater { return cfg.ConnectionGater }),
 		fx.Provide(func() pnet.PSK { return cfg.PSK }),
 		fx.Provide(func() network.ResourceManager { return cfg.ResourceManager }),
-		fx.Provide(func() *madns.Resolver { return cfg.MultiaddrResolver }),
 		fx.Provide(func(cm *quicreuse.ConnManager, sw *swarm.Swarm) libp2pwebrtc.ListenUDPFn {
 			hasQuicAddrPortFor := func(network string, laddr *net.UDPAddr) bool {
 				quicAddrPorts := map[string]struct{}{}

core/peer/addrinfo.go

@@ -61,6 +61,24 @@ func SplitAddr(m ma.Multiaddr) (transport ma.Multiaddr, id ID) {
 	return transport, id
 }
 
+// IDFromP2PAddr extracts the peer ID from a p2p Multiaddr
+func IDFromP2PAddr(m ma.Multiaddr) (ID, error) {
+	if m == nil {
+		return "", ErrInvalidAddr
+	}
+	var lastComponent ma.Component
+	ma.ForEach(m, func(c ma.Component) bool {
+		lastComponent = c
+		return true
+	})
+	if lastComponent.Protocol().Code != ma.P_P2P {
+		return "", ErrInvalidAddr
+	}
+
+	id := ID(lastComponent.RawValue()) // already validated by the multiaddr library.
+	return id, nil
+}
+
 // AddrInfoFromString builds an AddrInfo from the string representation of a Multiaddr
 func AddrInfoFromString(s string) (*AddrInfo, error) {
 	a, err := ma.NewMultiaddr(s)

core/peer/addrinfo_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	. "github.com/libp2p/go-libp2p/core/peer"
+	"github.com/stretchr/testify/require"
 
 	ma "github.com/multiformats/go-multiaddr"
 )
@@ -50,6 +51,19 @@ func TestSplitAddr(t *testing.T) {
 	}
 }
 
+func TestIDFromP2PAddr(t *testing.T) {
+	id, err := IDFromP2PAddr(maddrFull)
+	require.NoError(t, err)
+	require.Equal(t, testID, id)
+
+	id, err = IDFromP2PAddr(maddrPeer)
+	require.NoError(t, err)
+	require.Equal(t, testID, id)
+
+	_, err = IDFromP2PAddr(maddrTpt)
+	require.ErrorIs(t, err, ErrInvalidAddr)
+}
+
 func TestAddrInfoFromP2pAddr(t *testing.T) {
 	ai, err := AddrInfoFromP2pAddr(maddrFull)
 	if err != nil {

defaults.go

@@ -21,7 +21,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 
 	"github.com/multiformats/go-multiaddr"
-	madns "github.com/multiformats/go-multiaddr-dns"
 )
 
 // DefaultSecurity is the default security option.
@@ -128,11 +127,6 @@ var DefaultConnectionManager = func(cfg *Config) error {
 	return cfg.Apply(ConnectionManager(mgr))
 }
 
-// DefaultMultiaddrResolver creates a default connection manager
-var DefaultMultiaddrResolver = func(cfg *Config) error {
-	return cfg.Apply(MultiaddrResolver(madns.DefaultResolver))
-}
-
 // DefaultPrometheusRegisterer configures libp2p to use the default registerer
 var DefaultPrometheusRegisterer = func(cfg *Config) error {
 	return cfg.Apply(PrometheusRegisterer(prometheus.DefaultRegisterer))
@@ -198,10 +192,6 @@ var defaults = []struct {
 		fallback: func(cfg *Config) bool { return cfg.ConnManager == nil },
 		opt:      DefaultConnectionManager,
 	},
-	{
-		fallback: func(cfg *Config) bool { return cfg.MultiaddrResolver == nil },
-		opt:      DefaultMultiaddrResolver,
-	},
 	{
 		fallback: func(cfg *Config) bool { return !cfg.DisableMetrics && cfg.PrometheusRegisterer == nil },
 		opt:      DefaultPrometheusRegisterer,

options.go

@@ -31,7 +31,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 
 	ma "github.com/multiformats/go-multiaddr"
-	madns "github.com/multiformats/go-multiaddr-dns"
 	"go.uber.org/fx"
 )
 
@@ -495,7 +494,7 @@ func UserAgent(userAgent string) Option {
 }
 
 // MultiaddrResolver sets the libp2p dns resolver
-func MultiaddrResolver(rslv *madns.Resolver) Option {
+func MultiaddrResolver(rslv swarm.MultiaddrDNSResolver) Option {
 	return func(cfg *Config) error {
 		cfg.MultiaddrResolver = rslv
 		return nil

p2p/net/swarm/dial_test.go

@@ -52,7 +52,7 @@ func TestBasicDialPeerWithResolver(t *testing.T) {
 	resolver, err := madns.NewResolver(madns.WithDomainResolver("example.com", &mockResolver))
 	require.NoError(t, err)
 
-	swarms := makeSwarms(t, 2, swarmt.WithSwarmOpts(swarm.WithMultiaddrResolver(resolver)))
+	swarms := makeSwarms(t, 2, swarmt.WithSwarmOpts(swarm.WithMultiaddrResolver(swarm.ResolverFromMaDNS{resolver})))
 	defer closeSwarms(swarms)
 	s1 := swarms[0]
 	s2 := swarms[1]

p2p/net/swarm/resolve_test.go

@@ -0,0 +1,86 @@
+package swarm
+
+import (
+	"context"
+	"net"
+	"strconv"
+	"testing"
+
+	"github.com/multiformats/go-multiaddr"
+	madns "github.com/multiformats/go-multiaddr-dns"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSwarmResolver(t *testing.T) {
+	mockResolver := madns.MockResolver{IP: make(map[string][]net.IPAddr)}
+	ipaddr, err := net.ResolveIPAddr("ip4", "127.0.0.1")
+	require.NoError(t, err)
+	mockResolver.IP["example.com"] = []net.IPAddr{*ipaddr}
+	mockResolver.TXT = map[string][]string{
+		"_dnsaddr.example.com": {"dnsaddr=/ip4/127.0.0.1"},
+	}
+	madnsResolver, err := madns.NewResolver(madns.WithDomainResolver("example.com", &mockResolver))
+	require.NoError(t, err)
+	swarmResolver := ResolverFromMaDNS{madnsResolver}
+
+	ctx := context.Background()
+	res, err := swarmResolver.ResolveDNSComponent(ctx, multiaddr.StringCast("/dns/example.com"), 10)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(res))
+	require.Equal(t, "/ip4/127.0.0.1", res[0].String())
+
+	res, err = swarmResolver.ResolveDNSAddr(ctx, "", multiaddr.StringCast("/dnsaddr/example.com"), 1, 10)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(res))
+	require.Equal(t, "/ip4/127.0.0.1", res[0].String())
+
+	t.Run("Test Limits", func(t *testing.T) {
+		var ipaddrs []net.IPAddr
+		var manyDNSAddrs []string
+		for i := 0; i < 255; i++ {
+			ip := "1.2.3." + strconv.Itoa(i)
+			ipaddrs = append(ipaddrs, net.IPAddr{IP: net.ParseIP(ip)})
+			manyDNSAddrs = append(manyDNSAddrs, "dnsaddr=/ip4/"+ip)
+		}
+
+		mockResolver.IP = map[string][]net.IPAddr{
+			"example.com": ipaddrs,
+		}
+		mockResolver.TXT = map[string][]string{
+			"_dnsaddr.example.com": manyDNSAddrs,
+		}
+
+		res, err := swarmResolver.ResolveDNSComponent(ctx, multiaddr.StringCast("/dns/example.com"), 10)
+		require.NoError(t, err)
+		require.Equal(t, 10, len(res))
+		for i := 0; i < 10; i++ {
+			require.Equal(t, "/ip4/1.2.3."+strconv.Itoa(i), res[i].String())
+		}
+
+		res, err = swarmResolver.ResolveDNSAddr(ctx, "", multiaddr.StringCast("/dnsaddr/example.com"), 1, 10)
+		require.NoError(t, err)
+		require.Equal(t, 10, len(res))
+		for i := 0; i < 10; i++ {
+			require.Equal(t, "/ip4/1.2.3."+strconv.Itoa(i), res[i].String())
+		}
+	})
+
+	t.Run("Test Recursive Limits", func(t *testing.T) {
+		recursiveDNSAddr := make(map[string][]string)
+		for i := 0; i < 255; i++ {
+			recursiveDNSAddr["_dnsaddr."+strconv.Itoa(i)+".example.com"] = []string{"dnsaddr=/dnsaddr/" + strconv.Itoa(i+1) + ".example.com"}
+		}
+		recursiveDNSAddr["_dnsaddr.255.example.com"] = []string{"dnsaddr=/ip4/127.0.0.1"}
+		mockResolver.TXT = recursiveDNSAddr
+
+		res, err = swarmResolver.ResolveDNSAddr(ctx, "", multiaddr.StringCast("/dnsaddr/0.example.com"), 256, 10)
+		require.NoError(t, err)
+		require.Equal(t, 1, len(res))
+		require.Equal(t, "/ip4/127.0.0.1", res[0].String())
+
+		res, err = swarmResolver.ResolveDNSAddr(ctx, "", multiaddr.StringCast("/dnsaddr/0.example.com"), 255, 10)
+		require.NoError(t, err)
+		require.Equal(t, 1, len(res))
+		require.Equal(t, "/dnsaddr/255.example.com", res[0].String())
+	})
+}

p2p/net/swarm/swarm.go

@@ -60,9 +60,9 @@ func WithConnectionGater(gater connmgr.ConnectionGater) Option {
 }
 
 // WithMultiaddrResolver sets a custom multiaddress resolver
-func WithMultiaddrResolver(maResolver *madns.Resolver) Option {
+func WithMultiaddrResolver(resolver MultiaddrDNSResolver) Option {
 	return func(s *Swarm) error {
-		s.maResolver = maResolver
+		s.multiaddrResolver = resolver
 		return nil
 	}
 }
@@ -145,6 +145,14 @@ func WithReadOnlyBlackHoleDetector() Option {
 	}
 }
 
+type MultiaddrDNSResolver interface {
+	// ResolveDNSAddr resolves the first /dnsaddr component in a multiaddr.
+	// Recurisvely resolves DNSADDRs up to the recursion limit
+	ResolveDNSAddr(ctx context.Context, expectedPeerID peer.ID, maddr ma.Multiaddr, recursionLimit, outputLimit int) ([]ma.Multiaddr, error)
+	// ResolveDNSComponent resolves the first /{dns,dns4,dns6} component in a multiaddr.
+	ResolveDNSComponent(ctx context.Context, maddr ma.Multiaddr, outputLimit int) ([]ma.Multiaddr, error)
+}
+
 // Swarm is a connection muxer, allowing connections to other peers to
 // be opened and closed, while still using the same Chan for all
 // communication. The Chan sends/receives Messages, which note the
@@ -196,7 +204,7 @@ type Swarm struct {
 		m map[int]transport.Transport
 	}
 
-	maResolver *madns.Resolver
+	multiaddrResolver MultiaddrDNSResolver
 
 	// stream handlers
 	streamh atomic.Pointer[network.StreamHandler]
@@ -231,15 +239,15 @@ func NewSwarm(local peer.ID, peers peerstore.Peerstore, eventBus event.Bus, opts
 	}
 	ctx, cancel := context.WithCancel(context.Background())
 	s := &Swarm{
-		local:            local,
-		peers:            peers,
-		emitter:          emitter,
-		ctx:              ctx,
-		ctxCancel:        cancel,
-		dialTimeout:      defaultDialTimeout,
-		dialTimeoutLocal: defaultDialTimeoutLocal,
-		maResolver:       madns.DefaultResolver,
-		dialRanker:       DefaultDialRanker,
+		local:             local,
+		peers:             peers,
+		emitter:           emitter,
+		ctx:               ctx,
+		ctxCancel:         cancel,
+		dialTimeout:       defaultDialTimeout,
+		dialTimeoutLocal:  defaultDialTimeoutLocal,
+		multiaddrResolver: ResolverFromMaDNS{madns.DefaultResolver},
+		dialRanker:        DefaultDialRanker,
 
 		// A black hole is a binary property. On a network if UDP dials are blocked or there is
 		// no IPv6 connectivity, all dials will fail. So a low success rate of 5 out 100 dials
@@ -624,7 +632,6 @@ func isBetterConn(a, b *Conn) bool {
 
 // bestConnToPeer returns the best connection to peer.
 func (s *Swarm) bestConnToPeer(p peer.ID) *Conn {
-
 	// TODO: Prefer some transports over others.
 	// For now, prefers direct connections over Relayed connections.
 	// For tie-breaking, select the newest non-closed connection with the most streams.
@@ -813,8 +820,10 @@ func (s *Swarm) ResourceManager() network.ResourceManager {
 }
 
 // Swarm is a Network.
-var _ network.Network = (*Swarm)(nil)
-var _ transport.TransportNetwork = (*Swarm)(nil)
+var (
+	_ network.Network            = (*Swarm)(nil)
+	_ transport.TransportNetwork = (*Swarm)(nil)
+)
 
 type connWithMetrics struct {
 	transport.CapableConn
@@ -846,3 +855,95 @@ func (c connWithMetrics) Stat() network.ConnStats {
 }
 
 var _ network.ConnStat = connWithMetrics{}
+
+type ResolverFromMaDNS struct {
+	*madns.Resolver
+}
+
+var _ MultiaddrDNSResolver = ResolverFromMaDNS{}
+
+func startsWithDNSADDR(m ma.Multiaddr) bool {
+	if m == nil {
+		return false
+	}
+
+	startsWithDNSADDR := false
+	// Using ForEach to avoid allocating
+	ma.ForEach(m, func(c ma.Component) bool {
+		startsWithDNSADDR = c.Protocol().Code == ma.P_DNSADDR
+		return false
+	})
+	return startsWithDNSADDR
+}
+
+// ResolveDNSAddr implements MultiaddrDNSResolver
+func (r ResolverFromMaDNS) ResolveDNSAddr(ctx context.Context, expectedPeerID peer.ID, maddr ma.Multiaddr, recursionLimit int, outputLimit int) ([]ma.Multiaddr, error) {
+	if outputLimit <= 0 {
+		return nil, nil
+	}
+	if recursionLimit <= 0 {
+		return []ma.Multiaddr{maddr}, nil
+	}
+	var resolved, toResolve []ma.Multiaddr
+	addrs, err := r.Resolve(ctx, maddr)
+	if err != nil {
+		return nil, err
+	}
+	if len(addrs) > outputLimit {
+		addrs = addrs[:outputLimit]
+	}
+
+	for _, addr := range addrs {
+		if startsWithDNSADDR(addr) {
+			toResolve = append(toResolve, addr)
+		} else {
+			resolved = append(resolved, addr)
+		}
+	}
+
+	for _, addr := range toResolve {
+		resolvedAddrs, err := r.ResolveDNSAddr(ctx, expectedPeerID, addr, recursionLimit-1, outputLimit-len(resolved))
+		if err != nil {
+			log.Warnf("failed to resolve dnsaddr %v %s: ", addr, err)
+			// Dropping this address
+			continue
+		}
+		resolved = append(resolved, resolvedAddrs...)
+	}
+
+	if len(resolved) > outputLimit {
+		resolved = resolved[:outputLimit]
+	}
+
+	// If the address contains a peer id, make sure it matches our expectedPeerID
+	if expectedPeerID != "" {
+		removeMismatchPeerID := func(a ma.Multiaddr) bool {
+			id, err := peer.IDFromP2PAddr(a)
+			if err == peer.ErrInvalidAddr {
+				// This multiaddr didn't contain a peer id, assume it's for this peer.
+				// Handshake will fail later if it's not.
+				return false
+			} else if err != nil {
+				// This multiaddr is invalid, drop it.
+				return true
+			}
+
+			return id != expectedPeerID
+		}
+		resolved = slices.DeleteFunc(resolved, removeMismatchPeerID)
+	}
+
+	return resolved, nil
+}
+
+// ResolveDNSComponent implements MultiaddrDNSResolver
+func (r ResolverFromMaDNS) ResolveDNSComponent(ctx context.Context, maddr ma.Multiaddr, outputLimit int) ([]ma.Multiaddr, error) {
+	addrs, err := r.Resolve(ctx, maddr)
+	if err != nil {
+		return nil, err
+	}
+	if len(addrs) > outputLimit {
+		addrs = addrs[:outputLimit]
+	}
+	return addrs, nil
+}