Repository: libp2p/go-libp2p
Pull request #2862 (2024-07-04 13:19:28)

webtransport: close underlying h3 connection


Patch:
diff --git a/go.mod b/go.mod
index 2e2edc10..d9038066 100644
--- a/go.mod
+++ b/go.mod
@@ -10,7 +10,7 @@ require (
 	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.3.0
 	github.com/flynn/noise v1.1.0
 	github.com/google/gopacket v1.1.19
-	github.com/gorilla/websocket v1.5.3
+	github.com/gorilla/websocket v1.5.1
 	github.com/hashicorp/golang-lru/arc/v2 v2.0.7
 	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/ipfs/go-cid v0.4.1
diff --git a/go.sum b/go.sum
index 3c44afb2..8cf26679 100644
--- a/go.sum
+++ b/go.sum
@@ -121,8 +121,8 @@ github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
 github.com/googleapis/gax-go v2.0.0+incompatible/go.mod h1:SFVmujtThgffbyetf+mdk2eWhX2bMyUtNHzFKcPA9HY=
 github.com/googleapis/gax-go/v2 v2.0.3/go.mod h1:LLvjysVCY1JZeum8Z6l8qUty8fiNwE08qbEPm1M08qg=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
-github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
-github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/gorilla/websocket v1.5.1 h1:gmztn0JnHVt9JZquRuzLw3g4wouNVzKL15iLr/zn/QY=
+github.com/gorilla/websocket v1.5.1/go.mod h1:x3kM2JMyaluk02fnUJpQuwD2dCS5NDG2ZHL0uE0tcaY=
 github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
 github.com/grpc-ecosystem/grpc-gateway v1.5.0/go.mod h1:RSKVYQBd5MCa4OVpNdGskqpgL2+G+NZTnrVHpWWfpdw=
 github.com/hashicorp/golang-lru/arc/v2 v2.0.7 h1:QxkVTxwColcduO+LP7eJO56r2hFiG8zEbfAAzRv52KQ=
diff --git a/p2p/host/peerstore/pstoreds/protobook.go b/p2p/host/peerstore/pstoreds/protobook.go
index 9ef7d1c9..40fa7d95 100644
--- a/p2p/host/peerstore/pstoreds/protobook.go
+++ b/p2p/host/peerstore/pstoreds/protobook.go
@@ -48,7 +48,7 @@ func NewProtoBook(meta pstore.PeerMetadata, opts ...ProtoBookOption) (*dsProtoBo
 			}
 			return ret
 		}(),
-		maxProtos: 128,
+		maxProtos: 1024,
 	}
 
 	for _, opt := range opts {
diff --git a/p2p/host/peerstore/pstoremem/protobook.go b/p2p/host/peerstore/pstoremem/protobook.go
index b28ffe11..51c4b028 100644
--- a/p2p/host/peerstore/pstoremem/protobook.go
+++ b/p2p/host/peerstore/pstoremem/protobook.go
@@ -26,6 +26,9 @@ type memoryProtoBook struct {
 	segments protoSegments
 
 	maxProtos int
+
+	lk       sync.RWMutex
+	interned map[protocol.ID]protocol.ID
 }
 
 var _ pstore.ProtoBook = (*memoryProtoBook)(nil)
@@ -41,6 +44,7 @@ func WithMaxProtocols(num int) ProtoBookOption {
 
 func NewProtoBook(opts ...ProtoBookOption) (*memoryProtoBook, error) {
 	pb := &memoryProtoBook{
+		interned: make(map[protocol.ID]protocol.ID, 256),
 		segments: func() (ret protoSegments) {
 			for i := range ret {
 				ret[i] = &protoSegment{
@@ -49,7 +53,7 @@ func NewProtoBook(opts ...ProtoBookOption) (*memoryProtoBook, error) {
 			}
 			return ret
 		}(),
-		maxProtos: 128,
+		maxProtos: 1024,
 	}
 
 	for _, opt := range opts {
@@ -60,6 +64,30 @@ func NewProtoBook(opts ...ProtoBookOption) (*memoryProtoBook, error) {
 	return pb, nil
 }
 
+func (pb *memoryProtoBook) internProtocol(proto protocol.ID) protocol.ID {
+	// check if it is interned with the read lock
+	pb.lk.RLock()
+	interned, ok := pb.interned[proto]
+	pb.lk.RUnlock()
+
+	if ok {
+		return interned
+	}
+
+	// intern with the write lock
+	pb.lk.Lock()
+	defer pb.lk.Unlock()
+
+	// check again in case it got interned in between locks
+	interned, ok = pb.interned[proto]
+	if ok {
+		return interned
+	}
+
+	pb.interned[proto] = proto
+	return proto
+}
+
 func (pb *memoryProtoBook) SetProtocols(p peer.ID, protos ...protocol.ID) error {
 	if len(protos) > pb.maxProtos {
 		return errTooManyProtocols
@@ -67,7 +95,7 @@ func (pb *memoryProtoBook) SetProtocols(p peer.ID, protos ...protocol.ID) error
 
 	newprotos := make(map[protocol.ID]struct{}, len(protos))
 	for _, proto := range protos {
-		newprotos[proto] = struct{}{}
+		newprotos[pb.internProtocol(proto)] = struct{}{}
 	}
 
 	s := pb.segments.get(p)
@@ -93,7 +121,7 @@ func (pb *memoryProtoBook) AddProtocols(p peer.ID, protos ...protocol.ID) error
 	}
 
 	for _, proto := range protos {
-		protomap[proto] = struct{}{}
+		protomap[pb.internProtocol(proto)] = struct{}{}
 	}
 	return nil
 }
@@ -123,10 +151,7 @@ func (pb *memoryProtoBook) RemoveProtocols(p peer.ID, protos ...protocol.ID) err
 	}
 
 	for _, proto := range protos {
-		delete(protomap, proto)
-	}
-	if len(protomap) == 0 {
-		delete(s.protocols, p)
+		delete(protomap, pb.internProtocol(proto))
 	}
 	return nil
 }
diff --git a/p2p/transport/webtransport/conn.go b/p2p/transport/webtransport/conn.go
index 0e83b1d1..05251247 100644
--- a/p2p/transport/webtransport/conn.go
+++ b/p2p/transport/webtransport/conn.go
@@ -7,6 +7,7 @@ import (
 	tpt "github.com/libp2p/go-libp2p/core/transport"
 
 	ma "github.com/multiformats/go-multiaddr"
+	"github.com/quic-go/quic-go"
 	"github.com/quic-go/webtransport-go"
 )
 
@@ -31,16 +32,18 @@ type conn struct {
 	session   *webtransport.Session
 
 	scope network.ConnManagementScope
+	qconn quic.Connection
 }
 
 var _ tpt.CapableConn = &conn{}
 
-func newConn(tr *transport, sess *webtransport.Session, sconn *connSecurityMultiaddrs, scope network.ConnManagementScope) *conn {
+func newConn(tr *transport, sess *webtransport.Session, sconn *connSecurityMultiaddrs, scope network.ConnManagementScope, qconn quic.Connection) *conn {
 	return &conn{
 		connSecurityMultiaddrs: sconn,
 		transport:              tr,
 		session:                sess,
 		scope:                  scope,
+		qconn:                  qconn,
 	}
 }
 
@@ -70,7 +73,9 @@ func (c *conn) allowWindowIncrease(size uint64) bool {
 func (c *conn) Close() error {
 	c.scope.Done()
 	c.transport.removeConn(c.session)
-	return c.session.CloseWithError(0, "")
+	err := c.session.CloseWithError(0, "")
+	_ = c.qconn.CloseWithError(1, "")
+	return err
 }
 
 func (c *conn) IsClosed() bool           { return c.session.Context().Err() != nil }
diff --git a/p2p/transport/webtransport/listener.go b/p2p/transport/webtransport/listener.go
index 2a7c3546..ff611fe9 100644
--- a/p2p/transport/webtransport/listener.go
+++ b/p2p/transport/webtransport/listener.go
@@ -15,12 +15,61 @@ import (
 	"github.com/libp2p/go-libp2p/p2p/transport/quicreuse"
 
 	ma "github.com/multiformats/go-multiaddr"
+	"github.com/quic-go/quic-go"
+	"github.com/quic-go/quic-go/http3"
 	"github.com/quic-go/webtransport-go"
 )
 
 const queueLen = 16
 const handshakeTimeout = 10 * time.Second
 
+type connKey struct{}
+
+// negotiatingConn is a wrapper around a quic.Connection that lets us wrap it in
+// our own context for the duration of the upgrade process. Upgrading a quic
+// connection to an h3 connection to a webtransport session.
+type negotiatingConn struct {
+	quic.Connection
+	ctx    context.Context
+	cancel context.CancelFunc
+	// stopClose is a function that stops the connection from being closed when
+	// the context is done. Returns true if the connection close function was
+	// not called.
+	stopClose func() bool
+	err       error
+}
+
+func (c *negotiatingConn) Unwrap() (quic.Connection, error) {
+	defer c.cancel()
+	if c.stopClose != nil {
+		// unwrap the first time
+		if !c.stopClose() {
+			c.err = errTimeout
+		}
+		c.stopClose = nil
+	}
+	if c.err != nil {
+		return nil, c.err
+	}
+	return c.Connection, nil
+}
+
+func wrapConn(ctx context.Context, c quic.Connection, handshakeTimeout time.Duration) *negotiatingConn {
+	ctx, cancel := context.WithTimeout(ctx, handshakeTimeout)
+	stopClose := context.AfterFunc(ctx, func() {
+		log.Debugf("failed to handshake on conn: %s", c.RemoteAddr())
+		c.CloseWithError(1, "")
+	})
+	return &negotiatingConn{
+		Connection: c,
+		ctx:        ctx,
+		cancel:     cancel,
+		stopClose:  stopClose,
+	}
+}
+
+var errTimeout = errors.New("timeout")
+
 type listener struct {
 	transport       *transport
 	isStaticTLSConf bool
@@ -56,6 +105,11 @@ func newListener(reuseListener quicreuse.Listener, t *transport, isStaticTLSConf
 		addr:            reuseListener.Addr(),
 		multiaddr:       localMultiaddr,
 		server: webtransport.Server{
+			H3: http3.Server{
+				ConnContext: func(ctx context.Context, c quic.Connection) context.Context {
+					return context.WithValue(ctx, connKey{}, c)
+				},
+			},
 			CheckOrigin: func(r *http.Request) bool { return true },
 		},
 	}
@@ -71,7 +125,8 @@ func newListener(reuseListener quicreuse.Listener, t *transport, isStaticTLSConf
 				log.Debugw("serving failed", "addr", ln.Addr(), "error", err)
 				return
 			}
-			go ln.server.ServeQUICConn(conn)
+			wrapped := wrapConn(ln.ctx, conn, t.handshakeTimeout)
+			go ln.server.ServeQUICConn(wrapped)
 		}
 	}()
 	return ln, nil
@@ -137,13 +192,32 @@ func (l *listener) httpHandlerWithConnScope(w http.ResponseWriter, r *http.Reque
 		return err
 	}
 
-	conn := newConn(l.transport, sess, sconn, connScope)
+	connVal := r.Context().Value(connKey{})
+	if connVal == nil {
+		log.Errorf("missing conn from context")
+		sess.CloseWithError(1, "")
+		return errors.New("invalid context")
+	}
+	nconn, ok := connVal.(*negotiatingConn)
+	if !ok {
+		log.Errorf("unexpected connection in context. invalid conn type: %T", nconn)
+		sess.CloseWithError(1, "")
+		return errors.New("invalid context")
+	}
+	qconn, err := nconn.Unwrap()
+	if err != nil {
+		log.Debugf("handshake timed out: %s", r.RemoteAddr)
+		sess.CloseWithError(1, "")
+		return err
+	}
+
+	conn := newConn(l.transport, sess, sconn, connScope, qconn)
 	l.transport.addConn(sess, conn)
 	select {
 	case l.queue <- conn:
 	default:
 		log.Debugw("accept queue full, dropping incoming connection", "peer", sconn.RemotePeer(), "addr", r.RemoteAddr, "error", err)
-		sess.CloseWithError(1, "")
+		conn.Close()
 		return errors.New("accept queue full")
 	}
 
diff --git a/p2p/transport/webtransport/transport.go b/p2p/transport/webtransport/transport.go
index 97172703..ef8551d6 100644
--- a/p2p/transport/webtransport/transport.go
+++ b/p2p/transport/webtransport/transport.go
@@ -60,6 +60,13 @@ func WithTLSClientConfig(c *tls.Config) Option {
 	}
 }
 
+func WithHandshakeTimeout(d time.Duration) Option {
+	return func(t *transport) error {
+		t.handshakeTimeout = d
+		return nil
+	}
+}
+
 type transport struct {
 	privKey ic.PrivKey
 	pid     peer.ID
@@ -78,8 +85,9 @@ type transport struct {
 
 	noise *noise.Transport
 
-	connMx sync.Mutex
-	conns  map[quic.ConnectionTracingID]*conn // using quic-go's ConnectionTracingKey as map key
+	connMx           sync.Mutex
+	conns            map[quic.ConnectionTracingID]*conn // using quic-go's ConnectionTracingKey as map key
+	handshakeTimeout time.Duration
 }
 
 var _ tpt.Transport = &transport{}
@@ -99,13 +107,14 @@ func New(key ic.PrivKey, psk pnet.PSK, connManager *quicreuse.ConnManager, gater
 		return nil, err
 	}
 	t := &transport{
-		pid:         id,
-		privKey:     key,
-		rcmgr:       rcmgr,
-		gater:       gater,
-		clock:       clock.New(),
-		connManager: connManager,
-		conns:       map[quic.ConnectionTracingID]*conn{},
+		pid:              id,
+		privKey:          key,
+		rcmgr:            rcmgr,
+		gater:            gater,
+		clock:            clock.New(),
+		connManager:      connManager,
+		conns:            map[quic.ConnectionTracingID]*conn{},
+		handshakeTimeout: handshakeTimeout,
 	}
 	for _, opt := range opts {
 		if err := opt(t); err != nil {
@@ -159,7 +168,7 @@ func (t *transport) dialWithScope(ctx context.Context, raddr ma.Multiaddr, p pee
 	}
 
 	maddr, _ := ma.SplitFunc(raddr, func(c ma.Component) bool { return c.Protocol().Code == ma.P_WEBTRANSPORT })
-	sess, err := t.dial(ctx, maddr, url, sni, certHashes)
+	sess, qconn, err := t.dial(ctx, maddr, url, sni, certHashes)
 	if err != nil {
 		return nil, err
 	}
@@ -172,12 +181,12 @@ func (t *transport) dialWithScope(ctx context.Context, raddr ma.Multiaddr, p pee
 		sess.CloseWithError(errorCodeConnectionGating, "")
 		return nil, fmt.Errorf("secured connection gated")
 	}
-	conn := newConn(t, sess, sconn, scope)
+	conn := newConn(t, sess, sconn, scope, qconn)
 	t.addConn(sess, conn)
 	return conn, nil
 }
 
-func (t *transport) dial(ctx context.Context, addr ma.Multiaddr, url, sni string, certHashes []multihash.DecodedMultihash) (*webtransport.Session, error) {
+func (t *transport) dial(ctx context.Context, addr ma.Multiaddr, url, sni string, certHashes []multihash.DecodedMultihash) (*webtransport.Session, quic.Connection, error) {
 	var tlsConf *tls.Config
 	if t.tlsClientConf != nil {
 		tlsConf = t.tlsClientConf.Clone()
@@ -200,7 +209,7 @@ func (t *transport) dial(ctx context.Context, addr ma.Multiaddr, url, sni string
 	}
 	conn, err := t.connManager.DialQUIC(ctx, addr, tlsConf, t.allowWindowIncrease)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	dialer := webtransport.Dialer{
 		DialAddr: func(ctx context.Context, addr string, tlsCfg *tls.Config, cfg *quic.Config) (quic.EarlyConnection, error) {
@@ -210,12 +219,14 @@ func (t *transport) dial(ctx context.Context, addr ma.Multiaddr, url, sni string
 	}
 	rsp, sess, err := dialer.Dial(ctx, url, nil)
 	if err != nil {
-		return nil, err
+		conn.CloseWithError(1, "")
+		return nil, nil, err
 	}
 	if rsp.StatusCode < 200 || rsp.StatusCode > 299 {
-		return nil, fmt.Errorf("invalid response status code: %d", rsp.StatusCode)
+		conn.CloseWithError(1, "")
+		return nil, nil, fmt.Errorf("invalid response status code: %d", rsp.StatusCode)
 	}
-	return sess, err
+	return sess, conn, err
 }
 
 func (t *transport) upgrade(ctx context.Context, sess *webtransport.Session, p peer.ID, certHashes []multihash.DecodedMultihash) (*connSecurityMultiaddrs, error) {


Test patch:
diff --git a/p2p/protocol/identify/id_test.go b/p2p/protocol/identify/id_test.go
index 904e47ce..a65d64f2 100644
--- a/p2p/protocol/identify/id_test.go
+++ b/p2p/protocol/identify/id_test.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"math/rand"
 	"slices"
 	"sync"
 	"testing"
@@ -731,15 +730,6 @@ func TestLargeIdentifyMessage(t *testing.T) {
 	}
 }
 
-func randString(n int) string {
-	chars := "abcdefghijklmnopqrstuvwxyz"
-	buf := make([]byte, n)
-	for i := 0; i < n; i++ {
-		buf[i] = chars[rand.Intn(len(chars))]
-	}
-	return string(buf)
-}
-
 func TestLargePushMessage(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -748,9 +738,9 @@ func TestLargePushMessage(t *testing.T) {
 	h2 := blhost.NewBlankHost(swarmt.GenSwarm(t))
 
 	// add protocol strings to make the message larger
-	// about 3K of protocol strings
-	for i := 0; i < 100; i++ {
-		r := protocol.ID(fmt.Sprintf("%s-%d", randString(30), i))
+	// about 2K of protocol strings
+	for i := 0; i < 500; i++ {
+		r := protocol.ID(fmt.Sprintf("rand%d", i))
 		h1.SetStreamHandler(r, func(network.Stream) {})
 		h2.SetStreamHandler(r, func(network.Stream) {})
 	}
diff --git a/p2p/transport/webtransport/transport_test.go b/p2p/transport/webtransport/transport_test.go
index f6c850a2..bd414462 100644
--- a/p2p/transport/webtransport/transport_test.go
+++ b/p2p/transport/webtransport/transport_test.go
@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"net/http"
 	"os"
 	"runtime"
 	"sync/atomic"
@@ -827,3 +828,37 @@ func TestServerRotatesCertCorrectlyAfterSteps(t *testing.T) {
 		require.True(t, found, "Failed after hour: %v", i)
 	}
 }
+
+func TestH3ConnClosed(t *testing.T) {
+	_, serverKey := newIdentity(t)
+	tr, err := libp2pwebtransport.New(serverKey, nil, newConnManager(t), nil, nil, libp2pwebtransport.WithHandshakeTimeout(1*time.Second))
+	require.NoError(t, err)
+	defer tr.(io.Closer).Close()
+	ln, err := tr.Listen(ma.StringCast("/ip4/127.0.0.1/udp/0/quic-v1/webtransport"))
+	require.NoError(t, err)
+	defer ln.Close()
+
+	p, err := net.ListenPacket("udp", "127.0.0.1:0")
+	require.NoError(t, err)
+	conn, err := quic.Dial(context.Background(), p, ln.Addr(), &tls.Config{
+		InsecureSkipVerify: true,
+		NextProtos:         []string{http3.NextProtoH3},
+	}, nil)
+	require.NoError(t, err)
+	rt := &http3.SingleDestinationRoundTripper{
+		Connection: conn,
+	}
+	rt.Start()
+	require.Eventually(t, func() bool {
+		c := http.Client{
+			Transport: rt,
+			Timeout:   1 * time.Second,
+		}
+		resp, err := c.Get(fmt.Sprintf("https://%s", ln.Addr().String()))
+		if err != nil {
+			return true
+		}
+		resp.Body.Close()
+		return false
+	}, 10*time.Second, 1*time.Second)
+}
diff --git a/test-plans/go.mod b/test-plans/go.mod
index e19fd638..823570c6 100644
--- a/test-plans/go.mod
+++ b/test-plans/go.mod
@@ -28,7 +28,7 @@ require (
 	github.com/google/gopacket v1.1.19 // indirect
 	github.com/google/pprof v0.0.0-20240207164012-fb44976bdcd5 // indirect
 	github.com/google/uuid v1.4.0 // indirect
-	github.com/gorilla/websocket v1.5.3 // indirect
+	github.com/gorilla/websocket v1.5.1 // indirect
 	github.com/huin/goupnp v1.3.0 // indirect
 	github.com/ipfs/go-cid v0.4.1 // indirect
 	github.com/ipfs/go-log/v2 v2.5.1 // indirect
diff --git a/test-plans/go.sum b/test-plans/go.sum
index 5e9755f1..b6cb55bf 100644
--- a/test-plans/go.sum
+++ b/test-plans/go.sum
@@ -101,8 +101,8 @@ github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
 github.com/googleapis/gax-go v2.0.0+incompatible/go.mod h1:SFVmujtThgffbyetf+mdk2eWhX2bMyUtNHzFKcPA9HY=
 github.com/googleapis/gax-go/v2 v2.0.3/go.mod h1:LLvjysVCY1JZeum8Z6l8qUty8fiNwE08qbEPm1M08qg=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
-github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
-github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/gorilla/websocket v1.5.1 h1:gmztn0JnHVt9JZquRuzLw3g4wouNVzKL15iLr/zn/QY=
+github.com/gorilla/websocket v1.5.1/go.mod h1:x3kM2JMyaluk02fnUJpQuwD2dCS5NDG2ZHL0uE0tcaY=
 github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
 github.com/grpc-ecosystem/grpc-gateway v1.5.0/go.mod h1:RSKVYQBd5MCa4OVpNdGskqpgL2+G+NZTnrVHpWWfpdw=
 github.com/huin/goupnp v1.3.0 h1:UvLUlWDNpoUdYzb2TCn+MuTWtcjXKSza2n6CBdQ0xXc=