From baf63aef3f7f1850b507d169b9c40919bc9160d1 Mon Sep 17 00:00:00 2001
From: Marco Munizaga <git@marcopolo.io>
Date: Thu, 21 Mar 2024 13:52:03 -0700
Subject: [PATCH] Clean up after autonat

---
 config/config.go            | 39 ++++++++++++++++++++++---------------
 leaky_tests/leaky_test.go   |  7 -------
 libp2p_test.go              |  6 ++++++
 p2p/host/autonat/autonat.go |  4 ++--
 p2p/host/autonat/svc.go     |  5 +++++
 p2p/net/swarm/swarm.go      |  5 +++++
 6 files changed, 41 insertions(+), 25 deletions(-)

diff --git a/config/config.go b/config/config.go
index ff6fa765e3..3f0cd85e91 100644
--- a/config/config.go
+++ b/config/config.go
@@ -24,7 +24,6 @@ import (
 	"github.com/libp2p/go-libp2p/p2p/host/autonat"
 	"github.com/libp2p/go-libp2p/p2p/host/autorelay"
 	bhost "github.com/libp2p/go-libp2p/p2p/host/basic"
-	blankhost "github.com/libp2p/go-libp2p/p2p/host/blank"
 	"github.com/libp2p/go-libp2p/p2p/host/eventbus"
 	"github.com/libp2p/go-libp2p/p2p/host/peerstore/pstoremem"
 	rcmgr "github.com/libp2p/go-libp2p/p2p/host/resource-manager"
@@ -466,8 +465,7 @@ func (cfg *Config) addAutoNAT(h *bhost.BasicHost) error {
 		}
 
 		// Pull out the pieces of the config that we _actually_ care about.
-		// Specifically, don't set up things like autorelay, listeners,
-		// identify, etc.
+		// Specifically, don't set up things like listeners, identify, etc.
 		autoNatCfg := Config{
 			Transports:         cfg.Transports,
 			Muxers:             cfg.Muxers,
@@ -486,30 +484,39 @@ func (cfg *Config) addAutoNAT(h *bhost.BasicHost) error {
 			},
 		}
 
-		dialer, err := autoNatCfg.makeSwarm(eventbus.NewBus(), false)
-		if err != nil {
-			return err
-		}
-		dialerHost := blankhost.NewBlankHost(dialer)
 		fxopts, err := autoNatCfg.addTransports()
 		if err != nil {
-			dialerHost.Close()
 			return err
 		}
+		var dialer *swarm.Swarm
+
 		fxopts = append(fxopts,
-			fx.Supply(dialerHost.ID()),
-			fx.Supply(dialer),
+			fx.Provide(eventbus.NewBus),
+			fx.Provide(func(lifecycle fx.Lifecycle, b event.Bus) (*swarm.Swarm, error) {
+				lifecycle.Append(fx.Hook{
+					OnStop: func(context.Context) error {
+						return ps.Close()
+					}})
+				var err error
+				dialer, err = autoNatCfg.makeSwarm(b, false)
+				return dialer, err
+
+			}),
 			fx.Provide(func() crypto.PrivKey { return autonatPrivKey }),
 		)
 		app := fx.New(fxopts...)
 		if err := app.Err(); err != nil {
-			dialerHost.Close()
 			return err
 		}
-		// NOTE: We're dropping the blank host here but that's fine. It
-		// doesn't really _do_ anything and doesn't even need to be
-		// closed (as long as we close the underlying network).
-		autonatOpts = append(autonatOpts, autonat.EnableService(dialerHost.Network()))
+		err = app.Start(context.Background())
+		if err != nil {
+			return err
+		}
+		go func() {
+			<-dialer.Done() // The swarm used for autonat has closed, we can cleanup now
+			app.Stop(context.Background())
+		}()
+		autonatOpts = append(autonatOpts, autonat.EnableService(dialer))
 	}
 	if cfg.AutoNATConfig.ForceReachability != nil {
 		autonatOpts = append(autonatOpts, autonat.WithReachability(*cfg.AutoNATConfig.ForceReachability))
diff --git a/leaky_tests/leaky_test.go b/leaky_tests/leaky_test.go
index fd7d164ac4..172b656149 100644
--- a/leaky_tests/leaky_test.go
+++ b/leaky_tests/leaky_test.go
@@ -5,7 +5,6 @@ import (
 	"testing"
 
 	"github.com/libp2p/go-libp2p"
-	"github.com/stretchr/testify/require"
 )
 
 func TestBadTransportConstructor(t *testing.T) {
@@ -18,9 +17,3 @@ func TestBadTransportConstructor(t *testing.T) {
 		t.Error("expected error to contain debugging info")
 	}
 }
-
-func TestAutoNATService(t *testing.T) {
-	h, err := libp2p.New(libp2p.EnableNATService())
-	require.NoError(t, err)
-	h.Close()
-}
diff --git a/libp2p_test.go b/libp2p_test.go
index fe05b5aaec..7e4ac1b61e 100644
--- a/libp2p_test.go
+++ b/libp2p_test.go
@@ -370,6 +370,12 @@ func TestRoutedHost(t *testing.T) {
 	require.Equal(t, []peer.ID{id}, mockRouter.queried)
 }
 
+func TestAutoNATService(t *testing.T) {
+	h, err := New(EnableNATService())
+	require.NoError(t, err)
+	h.Close()
+}
+
 func TestMain(m *testing.M) {
 	goleak.VerifyTestMain(
 		m,
diff --git a/p2p/host/autonat/autonat.go b/p2p/host/autonat/autonat.go
index fc8c6763b4..804606a037 100644
--- a/p2p/host/autonat/autonat.go
+++ b/p2p/host/autonat/autonat.go
@@ -434,7 +434,7 @@ func (as *AmbientAutoNAT) Close() error {
 		as.service.Disable()
 	}
 	<-as.backgroundRunning
-	return nil
+	return as.service.Close()
 }
 
 // Status returns the AutoNAT observed reachability status.
@@ -444,7 +444,7 @@ func (s *StaticAutoNAT) Status() network.Reachability {
 
 func (s *StaticAutoNAT) Close() error {
 	if s.service != nil {
-		s.service.Disable()
+		return s.service.Close()
 	}
 	return nil
 }
diff --git a/p2p/host/autonat/svc.go b/p2p/host/autonat/svc.go
index cf1dff8e72..d29355014e 100644
--- a/p2p/host/autonat/svc.go
+++ b/p2p/host/autonat/svc.go
@@ -273,6 +273,11 @@ func (as *autoNATService) Disable() {
 	}
 }
 
+func (as *autoNATService) Close() error {
+	as.Disable()
+	return as.config.dialer.Close()
+}
+
 func (as *autoNATService) background(ctx context.Context) {
 	defer close(as.backgroundRunning)
 
diff --git a/p2p/net/swarm/swarm.go b/p2p/net/swarm/swarm.go
index f1a0f590de..0140c3f596 100644
--- a/p2p/net/swarm/swarm.go
+++ b/p2p/net/swarm/swarm.go
@@ -263,6 +263,11 @@ func (s *Swarm) Close() error {
 	return nil
 }
 
+// Done returns a channel that is closed when the swarm is closed.
+func (s *Swarm) Done() <-chan struct{} {
+	return s.ctx.Done()
+}
+
 func (s *Swarm) close() {
 	s.ctxCancel()