From 803e84632f3110ca40bb84f2bfd65bd4e73dbe60 Mon Sep 17 00:00:00 2001
From: sukun <sukunrt@gmail.com>
Date: Mon, 8 Jul 2024 20:25:36 +0530
Subject: [PATCH] config: fix AddrFactory for AutoNAT

---
 config/config.go | 40 +++++++++++++++++++++++-----------------
 libp2p_test.go   |  5 ++++-
 2 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/config/config.go b/config/config.go
index a99da1232d..623cc473fc 100644
--- a/config/config.go
+++ b/config/config.go
@@ -486,26 +486,33 @@ func (cfg *Config) NewNode() (host.Host, error) {
 		)
 	}
 
-	// Note: h.AddrsFactory may be changed by relayFinder, but non-relay version is
-	// used by AutoNAT below.
-	if cfg.EnableAutoRelay {
-		if !cfg.DisableMetrics {
-			mt := autorelay.WithMetricsTracer(
-				autorelay.NewMetricsTracer(autorelay.WithRegisterer(cfg.PrometheusRegisterer)))
-			mtOpts := []autorelay.Option{mt}
-			cfg.AutoRelayOpts = append(mtOpts, cfg.AutoRelayOpts...)
-		}
-		fxopts = append(fxopts,
-			fx.Invoke(func(h *bhost.BasicHost, lifecycle fx.Lifecycle) (*autorelay.AutoRelay, error) {
+	// oldAddrFactory is the AddrFactory before it's modified by autorelay
+	// we need this for checking reachability via autonat
+	oldAddrFactory := func(addrs []ma.Multiaddr) []ma.Multiaddr {
+		return addrs
+	}
+
+	fxopts = append(fxopts,
+		fx.Invoke(func(h *bhost.BasicHost, lifecycle fx.Lifecycle) (*autorelay.AutoRelay, error) {
+			oldAddrFactory = h.AddrsFactory
+			if cfg.EnableAutoRelay {
+				if !cfg.DisableMetrics {
+					mt := autorelay.WithMetricsTracer(
+						autorelay.NewMetricsTracer(autorelay.WithRegisterer(cfg.PrometheusRegisterer)))
+					mtOpts := []autorelay.Option{mt}
+					cfg.AutoRelayOpts = append(mtOpts, cfg.AutoRelayOpts...)
+				}
+
 				ar, err := autorelay.NewAutoRelay(h, cfg.AutoRelayOpts...)
 				if err != nil {
 					return nil, err
 				}
 				lifecycle.Append(fx.StartStopHook(ar.Start, ar.Close))
 				return ar, nil
-			}),
-		)
-	}
+			}
+			return nil, nil
+		}),
+	)
 
 	var bh *bhost.BasicHost
 	fxopts = append(fxopts, fx.Invoke(func(bho *bhost.BasicHost) { bh = bho }))
@@ -523,7 +530,7 @@ func (cfg *Config) NewNode() (host.Host, error) {
 		return nil, err
 	}
 
-	if err := cfg.addAutoNAT(bh); err != nil {
+	if err := cfg.addAutoNAT(bh, oldAddrFactory); err != nil {
 		app.Stop(context.Background())
 		if cfg.Routing != nil {
 			rh.Close()
@@ -539,8 +546,7 @@ func (cfg *Config) NewNode() (host.Host, error) {
 	return &closableBasicHost{App: app, BasicHost: bh}, nil
 }
 
-func (cfg *Config) addAutoNAT(h *bhost.BasicHost) error {
-	addrF := h.AddrsFactory
+func (cfg *Config) addAutoNAT(h *bhost.BasicHost, addrF AddrsFactory) error {
 	autonatOpts := []autonat.Option{
 		autonat.UsingAddresses(func() []ma.Multiaddr {
 			return addrF(h.AllAddrs())
diff --git a/libp2p_test.go b/libp2p_test.go
index ea52e56470..e8339705d8 100644
--- a/libp2p_test.go
+++ b/libp2p_test.go
@@ -377,7 +377,10 @@ func TestRoutedHost(t *testing.T) {
 }
 
 func TestAutoNATService(t *testing.T) {
-	h, err := New(EnableNATService())
+	// h, err := New(EnableNATService())
+	// require.NoError(t, err)
+	// h.Close()
+	h, err := New(EnableNATService(), EnableAutoRelayWithStaticRelays([]peer.AddrInfo{{ID: peer.ID("hello")}}))
 	require.NoError(t, err)
 	h.Close()
 }