From d7ba37217c6f1a0d739666f78ac19a87728a7280 Mon Sep 17 00:00:00 2001
From: Julien Muret <birros@protonmail.com>
Date: Sun, 5 Jun 2022 21:59:32 +0200
Subject: [PATCH] swarm: add ListenClose (#1586)

* feat(swarm): add ListenClose

* fix(swarm): fix ListenClose behavior

* refactor(swarm): move the listener's closing log
---
 p2p/net/swarm/swarm_listen.go | 48 ++++++++++++++++++++++++++++++-----
 p2p/net/swarm/swarm_test.go   | 19 ++++++++++++++
 2 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/p2p/net/swarm/swarm_listen.go b/p2p/net/swarm/swarm_listen.go
index ca54280c08..94fc6cce90 100644
--- a/p2p/net/swarm/swarm_listen.go
+++ b/p2p/net/swarm/swarm_listen.go
@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/libp2p/go-libp2p-core/network"
+	"github.com/libp2p/go-libp2p-core/transport"
 
 	ma "github.com/multiformats/go-multiaddr"
 )
@@ -35,6 +36,27 @@ func (s *Swarm) Listen(addrs ...ma.Multiaddr) error {
 	return nil
 }
 
+// ListenClose stop and delete listeners for all of the given addresses.
+func (s *Swarm) ListenClose(addrs ...ma.Multiaddr) {
+	var listenersToClose []transport.Listener
+
+	s.listeners.Lock()
+	for l := range s.listeners.m {
+		if !containsMultiaddr(addrs, l.Multiaddr()) {
+			continue
+		}
+
+		delete(s.listeners.m, l)
+		listenersToClose = append(listenersToClose, l)
+	}
+	s.listeners.cacheEOL = time.Time{}
+	s.listeners.Unlock()
+
+	for _, l := range listenersToClose {
+		l.Close()
+	}
+}
+
 // AddListenAddr tells the swarm to listen on a single address. Unlike Listen,
 // this method does not attempt to filter out bad addresses.
 func (s *Swarm) AddListenAddr(a ma.Multiaddr) error {
@@ -78,12 +100,19 @@ func (s *Swarm) AddListenAddr(a ma.Multiaddr) error {
 
 	go func() {
 		defer func() {
-			list.Close()
 			s.listeners.Lock()
-			delete(s.listeners.m, list)
-			s.listeners.cacheEOL = time.Time{}
+			_, ok := s.listeners.m[list]
+			if ok {
+				delete(s.listeners.m, list)
+				s.listeners.cacheEOL = time.Time{}
+			}
 			s.listeners.Unlock()
 
+			if ok {
+				list.Close()
+				log.Errorf("swarm listener unintentionally closed")
+			}
+
 			// signal to our notifiees on listen close.
 			s.notifyAll(func(n network.Notifiee) {
 				n.ListenClose(s, maddr)
@@ -93,10 +122,6 @@ func (s *Swarm) AddListenAddr(a ma.Multiaddr) error {
 		for {
 			c, err := list.Accept()
 			if err != nil {
-				if s.ctx.Err() == nil {
-					// only log if the swarm is still running.
-					log.Errorf("swarm listener accept error: %s", err)
-				}
 				return
 			}
 
@@ -119,3 +144,12 @@ func (s *Swarm) AddListenAddr(a ma.Multiaddr) error {
 	}()
 	return nil
 }
+
+func containsMultiaddr(addrs []ma.Multiaddr, addr ma.Multiaddr) bool {
+	for _, a := range addrs {
+		if addr == a {
+			return true
+		}
+	}
+	return false
+}
diff --git a/p2p/net/swarm/swarm_test.go b/p2p/net/swarm/swarm_test.go
index 2218b821b4..a5e193f03d 100644
--- a/p2p/net/swarm/swarm_test.go
+++ b/p2p/net/swarm/swarm_test.go
@@ -542,3 +542,22 @@ func TestResourceManagerAcceptStream(t *testing.T) {
 	}
 	require.EqualError(t, err, "stream reset")
 }
+
+func TestListenCloseCount(t *testing.T) {
+	s := GenSwarm(t, OptDialOnly)
+	addrsToListen := []ma.Multiaddr{
+		ma.StringCast("/ip4/0.0.0.0/tcp/0"),
+		ma.StringCast("/ip4/0.0.0.0/udp/0/quic"),
+	}
+
+	if err := s.Listen(addrsToListen...); err != nil {
+		t.Fatal(err)
+	}
+	listenedAddrs := s.ListenAddresses()
+	require.Equal(t, 2, len(listenedAddrs))
+
+	s.ListenClose(listenedAddrs...)
+
+	remainingAddrs := s.ListenAddresses()
+	require.Equal(t, 0, len(remainingAddrs))
+}