Limit total number of resolved addresses from DNS response (#63)

resolve.go

@@ -21,6 +21,8 @@ var (
 	DefaultResolver     = &Resolver{def: net.DefaultResolver}
 )
 
+const maxResolvedAddrs = 100
+
 const dnsaddrTXTPrefix = "dnsaddr="
 
 // BasicResolver is a low level interface for DNS resolution
@@ -254,6 +256,10 @@ func (r *Resolver) Resolve(ctx context.Context, maddr ma.Multiaddr) ([]ma.Multia
 		return nil, nil
 	}
 
+	if len(resolved) > maxResolvedAddrs {
+		resolved = resolved[:maxResolvedAddrs]
+	}
+
 	if preDNS != nil {
 		for i, m := range resolved {
 			resolved[i] = preDNS.Encapsulate(m)

resolve_test.go

@@ -3,6 +3,7 @@ package madns
 import (
 	"context"
 	"net"
+	"strconv"
 	"testing"
 
 	ma "github.com/multiformats/go-multiaddr"
@@ -363,3 +364,26 @@ func TestCustomResolver(t *testing.T) {
 		t.Fatal("expected result to be ip5")
 	}
 }
+
+func TestLimitResolver(t *testing.T) {
+	var ipaddrs []net.IPAddr
+	for i := 0; i < 255; i++ {
+		ipaddrs = append(ipaddrs, net.IPAddr{IP: net.ParseIP("1.2.3." + strconv.Itoa(i))})
+	}
+
+	mock := &MockResolver{
+		IP: map[string][]net.IPAddr{
+			"example.com": ipaddrs,
+		},
+		TXT: map[string][]string{},
+	}
+	resolver := &Resolver{def: mock}
+
+	addrs, err := resolver.Resolve(context.Background(), ma.StringCast("/dns4/example.com"))
+	if err != nil {
+		t.Error(err)
+	}
+	if len(addrs) != maxResolvedAddrs {
+		t.Fatalf("expected %d, got %d", maxResolvedAddrs, len(addrs))
+	}
+}