autonat: Clean up after close

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))

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()
-}

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,

p2p/host/autonat/autonat.go

@@ -434,6 +434,9 @@ func (as *AmbientAutoNAT) Close() error {
 		as.service.Disable()
 	}
 	<-as.backgroundRunning
+	if as.service != nil {
+		return as.service.Close()
+	}
 	return nil
 }
 
@@ -444,7 +447,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
 }

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)
 

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()