Repository: libp2p/go-libp2p
Pull request #3243 (2025-03-19 17:25:07)

webrtc: fix memory leak with udpmux.muxedConnection context
This context wasn't being cancelled on all code paths. In particular, contexts for connections that didn't complete negotiation were not being cancelled.

The change arranges for either udpmux.muxedConnection.Close or  RemoveConnByUfrag to call the other.

Fixes: #3223


Patch:
diff --git a/p2p/http/libp2phttp.go b/p2p/http/libp2phttp.go
index ae73c2b5..26f4825d 100644
--- a/p2p/http/libp2phttp.go
+++ b/p2p/http/libp2phttp.go
@@ -45,10 +45,6 @@ const LegacyWellKnownProtocols = "/.well-known/libp2p"
 const peerMetadataLimit = 8 << 10 // 8KB
 const peerMetadataLRUSize = 256   // How many different peer's metadata to keep in our LRU cache
 
-// DefaultNewStreamTimeout is the default value for new stream establishing timeout.
-// It is the same value as basic_host.DefaultNegotiationTimeout
-var DefaultNewStreamTimeout = 10 * time.Second
-
 type clientPeerIDContextKey struct{}
 type serverPeerIDContextKey struct{}
 
@@ -500,16 +496,7 @@ func (rt *streamRoundTripper) RoundTrip(r *http.Request) (*http.Response, error)
 		})
 	}
 
-	// If r.Context() timeout is greater than DefaultNewStreamTimeout
-	// use DefaultNewStreamTimeout for new stream negotiation.
-	newStreamCtx := r.Context()
-	if deadline, ok := newStreamCtx.Deadline(); !ok || deadline.After(time.Now().Add(DefaultNewStreamTimeout)) {
-		var cancel context.CancelFunc
-		newStreamCtx, cancel = context.WithTimeout(context.Background(), DefaultNewStreamTimeout)
-		defer cancel()
-	}
-
-	s, err := rt.h.NewStream(newStreamCtx, rt.server, ProtocolIDForMultistreamSelect)
+	s, err := rt.h.NewStream(r.Context(), rt.server, ProtocolIDForMultistreamSelect)
 	if err != nil {
 		return nil, err
 	}
diff --git a/p2p/transport/webrtc/udpmux/mux.go b/p2p/transport/webrtc/udpmux/mux.go
index f01facae..76d68d8e 100644
--- a/p2p/transport/webrtc/udpmux/mux.go
+++ b/p2p/transport/webrtc/udpmux/mux.go
@@ -271,12 +271,13 @@ func (mux *UDPMux) RemoveConnByUfrag(ufrag string) {
 
 	for _, isIPv6 := range [...]bool{true, false} {
 		key := ufragConnKey{ufrag: ufrag, isIPv6: isIPv6}
-		if _, ok := mux.ufragMap[key]; ok {
+		if conn, ok := mux.ufragMap[key]; ok {
 			delete(mux.ufragMap, key)
 			for _, addr := range mux.ufragAddrMap[key] {
 				delete(mux.addrMap, addr.String())
 			}
 			delete(mux.ufragAddrMap, key)
+			conn.close()
 		}
 	}
 }
@@ -293,7 +294,7 @@ func (mux *UDPMux) getOrCreateConn(ufrag string, isIPv6 bool, _ *UDPMux, addr ne
 		return false, conn
 	}
 
-	conn := newMuxedConnection(mux, func() { mux.RemoveConnByUfrag(ufrag) })
+	conn := newMuxedConnection(mux, ufrag)
 	mux.ufragMap[key] = conn
 	mux.addrMap[addr.String()] = conn
 	mux.ufragAddrMap[key] = append(mux.ufragAddrMap[key], addr)
diff --git a/p2p/transport/webrtc/udpmux/muxed_connection.go b/p2p/transport/webrtc/udpmux/muxed_connection.go
index 2af5d332..84d30c84 100644
--- a/p2p/transport/webrtc/udpmux/muxed_connection.go
+++ b/p2p/transport/webrtc/udpmux/muxed_connection.go
@@ -23,31 +23,29 @@ const queueLen = 128
 // from which this connection (indexed by ufrag) received
 // data.
 type muxedConnection struct {
-	ctx     context.Context
-	cancel  context.CancelFunc
-	onClose func()
-	queue   chan packet
-	mux     *UDPMux
+	ctx    context.Context
+	cancel context.CancelFunc
+	queue  chan packet
+	mux    *UDPMux
+	ufrag  string
 }
 
 var _ net.PacketConn = &muxedConnection{}
 
-func newMuxedConnection(mux *UDPMux, onClose func()) *muxedConnection {
+func newMuxedConnection(mux *UDPMux, ufrag string) *muxedConnection {
 	ctx, cancel := context.WithCancel(mux.ctx)
 	return &muxedConnection{
-		ctx:     ctx,
-		cancel:  cancel,
-		queue:   make(chan packet, queueLen),
-		onClose: onClose,
-		mux:     mux,
+		ctx:    ctx,
+		cancel: cancel,
+		queue:  make(chan packet, queueLen),
+		mux:    mux,
+		ufrag:  ufrag,
 	}
 }
 
 func (c *muxedConnection) Push(buf []byte, addr net.Addr) error {
-	select {
-	case <-c.ctx.Done():
+	if c.ctx.Err() != nil {
 		return errors.New("closed")
-	default:
 	}
 	select {
 	case c.queue <- packet{buf: buf, addr: addr}:
@@ -76,12 +74,21 @@ func (c *muxedConnection) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 }
 
 func (c *muxedConnection) Close() error {
-	select {
-	case <-c.ctx.Done():
+	if c.ctx.Err() != nil {
 		return nil
-	default:
 	}
-	c.onClose()
+	// mux calls close to actually close the connection
+	//
+	// Removing the connection from the mux or closing the connection
+	// must trigger the other.
+	// Doing this here ensures we don't need to call both RemoveConnByUfrag
+	// and close on all code paths.
+	c.mux.RemoveConnByUfrag(c.ufrag)
+	return nil
+}
+
+// closes the connection. Must only be called by the mux.
+func (c *muxedConnection) close() {
 	c.cancel()
 	// drain the packet queue
 	for {
@@ -89,7 +96,7 @@ func (c *muxedConnection) Close() error {
 		case p := <-c.queue:
 			pool.Put(p.buf)
 		default:
-			return nil
+			return
 		}
 	}
 }


Test patch:
diff --git a/p2p/transport/webrtc/udpmux/mux_test.go b/p2p/transport/webrtc/udpmux/mux_test.go
index 298e5c92..b75f3e83 100644
--- a/p2p/transport/webrtc/udpmux/mux_test.go
+++ b/p2p/transport/webrtc/udpmux/mux_test.go
@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/pion/stun"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -246,3 +247,28 @@ func TestMuxedConnection(t *testing.T) {
 	}
 	require.Empty(t, addrUfragMap)
 }
+
+func TestRemovingUfragClosesConn(t *testing.T) {
+	c := newPacketConn(t)
+	m := NewUDPMux(c)
+	m.Start()
+	defer m.Close()
+	remoteAddr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 1234}
+	conn, err := m.GetConn("a", remoteAddr)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	connClosed := make(chan bool)
+	go func() {
+		_, _, err := conn.ReadFrom(make([]byte, 100))
+		assert.ErrorIs(t, err, context.Canceled)
+		close(connClosed)
+	}()
+	require.NoError(t, err)
+	m.RemoveConnByUfrag("a")
+	select {
+	case <-connClosed:
+	case <-time.After(1 * time.Second):
+		t.Fatalf("expected the connection to be closed")
+	}
+}