From 81371564af0bb62094eb02fef37482304e39855e Mon Sep 17 00:00:00 2001
From: Pavel Zbitskiy <pavel@algorand.com>
Date: Thu, 1 Feb 2024 17:21:49 -0500
Subject: [PATCH] p2p: fix infinite loop in dnsaddr resolution

---
 network/p2p/dnsaddr/resolve.go      |  6 +++++
 network/p2p/dnsaddr/resolve_test.go | 40 +++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/network/p2p/dnsaddr/resolve.go b/network/p2p/dnsaddr/resolve.go
index 56caa84f4b..dbe2b002a1 100644
--- a/network/p2p/dnsaddr/resolve.go
+++ b/network/p2p/dnsaddr/resolve.go
@@ -35,8 +35,14 @@ func Iterate(initial multiaddr.Multiaddr, controller ResolveController, f func(d
 	if resolver == nil {
 		return errors.New("passed controller has no resolvers Iterate")
 	}
+	const maxHops = 100
+	hops := 0
 	var toResolve = []multiaddr.Multiaddr{initial}
 	for resolver != nil && len(toResolve) > 0 {
+		hops++
+		if hops > maxHops {
+			return errors.New("max hops reached while resolving dnsaddr " + initial.String())
+		}
 		curr := toResolve[0]
 		maddrs, resolveErr := resolver.Resolve(context.Background(), curr)
 		if resolveErr != nil {
diff --git a/network/p2p/dnsaddr/resolve_test.go b/network/p2p/dnsaddr/resolve_test.go
index 937e4db183..03190ab5f7 100644
--- a/network/p2p/dnsaddr/resolve_test.go
+++ b/network/p2p/dnsaddr/resolve_test.go
@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"net"
 	"testing"
+	"time"
 
 	"github.com/multiformats/go-multiaddr"
 	madns "github.com/multiformats/go-multiaddr-dns"
@@ -109,3 +110,42 @@ func TestMultiaddrsFromResolverDnsFailure(t *testing.T) {
 	assert.Empty(t, maddrs)
 	assert.ErrorContains(t, err, "always errors")
 }
+
+type mockController struct {
+}
+
+func (c mockController) Resolver() Resolver {
+	return selfResolver{}
+}
+
+func (c mockController) NextResolver() Resolver {
+	return nil
+}
+
+type selfResolver struct {
+}
+
+func (r selfResolver) Resolve(ctx context.Context, maddr multiaddr.Multiaddr) ([]multiaddr.Multiaddr, error) {
+	return []multiaddr.Multiaddr{maddr}, nil
+}
+
+// TestIterate ensures the Iterate() does not hang in infinite loop
+// when resolver returns the same dnsaddr
+func TestIterate(t *testing.T) {
+	partitiontest.PartitionTest(t)
+	t.Parallel()
+
+	dnsAddr := "/dnsaddr/foobar.com"
+	require.True(t, isDnsaddr(multiaddr.StringCast(dnsAddr)))
+	ma, err := multiaddr.NewMultiaddr(dnsAddr)
+	require.NoError(t, err)
+
+	require.Eventually(t, func() bool {
+		Iterate(
+			ma,
+			mockController{},
+			func(dnsaddr multiaddr.Multiaddr, entries []multiaddr.Multiaddr) error { return nil },
+		)
+		return true
+	}, 100*time.Millisecond, 50*time.Millisecond)
+}