enable and use the QUIC Stream Resets with Partial Delivery extension (#239)

* enable and use the QUIC Stream Resets with Partial Delivery extension

* fix stream header sending for concurrent Close / CancelWrite calls

Author: marten-seemann

client.go

@@ -58,9 +58,17 @@ func (d *Dialer) Dial(ctx context.Context, urlStr string, reqHdr http.Header) (*
 
 	quicConf := d.QUICConfig
 	if quicConf == nil {
-		quicConf = &quic.Config{EnableDatagrams: true}
-	} else if !d.QUICConfig.EnableDatagrams {
-		return nil, nil, errors.New("webtransport: DATAGRAM support required, enable it via QUICConfig.EnableDatagrams")
+		quicConf = &quic.Config{
+			EnableDatagrams:                  true,
+			EnableStreamResetPartialDelivery: true,
+		}
+	} else {
+		if !d.QUICConfig.EnableDatagrams {
+			return nil, nil, errors.New("webtransport: DATAGRAM support required, enable it via QUICConfig.EnableDatagrams")
+		}
+		if !d.QUICConfig.EnableStreamResetPartialDelivery {
+			return nil, nil, errors.New("webtransport: stream reset partial delivery required, enable it via QUICConfig.EnableStreamResetPartialDelivery")
+		}
 	}
 
 	tlsConf := d.TLSClientConfig

server.go

@@ -112,12 +112,14 @@ func (s *Server) Serve(conn net.PacketConn) error {
 	if err := s.initialize(); err != nil {
 		return err
 	}
-	quicConf := s.H3.QUICConfig
-	if quicConf == nil {
+	var quicConf *quic.Config
+	if s.H3.QUICConfig != nil {
+		quicConf = s.H3.QUICConfig.Clone()
+	} else {
 		quicConf = &quic.Config{}
 	}
-	quicConf = quicConf.Clone()
 	quicConf.EnableDatagrams = true
+	quicConf.EnableStreamResetPartialDelivery = true
 	ln, err := quic.ListenEarly(conn, s.H3.TLSConfig, quicConf)
 	if err != nil {
 		return err
@@ -142,6 +144,13 @@ func (s *Server) Serve(conn net.PacketConn) error {
 
 // ServeQUICConn serves a single QUIC connection.
 func (s *Server) ServeQUICConn(conn *quic.Conn) error {
+	connState := conn.ConnectionState()
+	if !connState.SupportsDatagrams.Local {
+		return errors.New("webtransport: QUIC DATAGRAM support required, enable it via QUICConfig.EnableDatagrams")
+	}
+	if !connState.SupportsStreamResetPartialDelivery.Local {
+		return errors.New("webtransport: QUIC Stream Resets with Partial Delivery required, enable it via QUICConfig.EnableStreamResetPartialDelivery")
+	}
 	if err := s.initialize(); err != nil {
 		return err
 	}

server_test.go

@@ -82,7 +82,7 @@ func TestServerReorderedUpgradeRequest(t *testing.T) {
 		context.Background(),
 		fmt.Sprintf("localhost:%d", port),
 		&tls.Config{RootCAs: webtransport.CertPool, NextProtos: []string{http3.NextProtoH3}},
-		&quic.Config{EnableDatagrams: true},
+		&quic.Config{EnableDatagrams: true, EnableStreamResetPartialDelivery: true},
 	)
 	require.NoError(t, err)
 	// Open a new stream for a WebTransport session we'll establish later. Stream ID: 0.
@@ -235,3 +235,65 @@ func TestImmediateClose(t *testing.T) {
 	s := webtransport.Server{H3: &http3.Server{}}
 	require.NoError(t, s.Close())
 }
+
+func TestServerConnectionStateChecks(t *testing.T) {
+	tests := []struct {
+		name                     string
+		enableDatagrams          bool
+		enableStreamResetPartial bool
+		wantErr                  string
+	}{
+		{
+			name:                     "missing datagram support",
+			enableDatagrams:          false,
+			enableStreamResetPartial: true,
+			wantErr:                  "webtransport: QUIC DATAGRAM support required",
+		},
+		{
+			name:                     "missing stream reset partial delivery support",
+			enableDatagrams:          true,
+			enableStreamResetPartial: false,
+			wantErr:                  "webtransport: QUIC Stream Resets with Partial Delivery required",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := webtransport.Server{H3: &http3.Server{TLSConfig: webtransport.TLSConf}}
+			defer s.Close()
+
+			serverConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0})
+			require.NoError(t, err)
+			defer serverConn.Close()
+
+			clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0})
+			require.NoError(t, err)
+			defer clientConn.Close()
+
+			ln, err := quic.ListenEarly(serverConn, webtransport.TLSConf, &quic.Config{
+				EnableDatagrams:                  tt.enableDatagrams,
+				EnableStreamResetPartialDelivery: tt.enableStreamResetPartial,
+			})
+			require.NoError(t, err)
+			defer ln.Close()
+
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			defer cancel()
+			_, err = quic.DialEarly(ctx, clientConn, ln.Addr(), &tls.Config{
+				ServerName: "localhost",
+				NextProtos: []string{http3.NextProtoH3},
+				RootCAs:    webtransport.CertPool,
+			}, &quic.Config{
+				EnableDatagrams:                  tt.enableDatagrams,
+				EnableStreamResetPartialDelivery: tt.enableStreamResetPartial,
+			})
+			require.NoError(t, err)
+
+			qconn, err := ln.Accept(ctx)
+			require.NoError(t, err)
+			defer qconn.CloseWithError(0, "")
+
+			require.ErrorContains(t, s.ServeQUICConn(qconn), tt.wantErr)
+		})
+	}
+}

session_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/quic-go/quic-go/http3"
 	"github.com/quic-go/quic-go/quicvarint"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -54,9 +55,10 @@ func newConnPair(t *testing.T, clientConn, serverConn net.PacketConn) (client, s
 		serverConn,
 		TLSConf,
 		&quic.Config{
-			InitialStreamReceiveWindow:     1 << 60,
-			InitialConnectionReceiveWindow: 1 << 60,
-			EnableDatagrams:                true,
+			InitialStreamReceiveWindow:       1 << 60,
+			InitialConnectionReceiveWindow:   1 << 60,
+			EnableDatagrams:                  true,
+			EnableStreamResetPartialDelivery: true,
 		},
 	)
 	require.NoError(t, err)
@@ -72,18 +74,27 @@ func newConnPair(t *testing.T, clientConn, serverConn net.PacketConn) (client, s
 			NextProtos: []string{http3.NextProtoH3},
 			RootCAs:    CertPool,
 		},
-		&quic.Config{EnableDatagrams: true},
+		&quic.Config{
+			EnableDatagrams:                  true,
+			EnableStreamResetPartialDelivery: true,
+		},
 	)
 	require.NoError(t, err)
-	require.True(t, cl.ConnectionState().SupportsDatagrams.Remote)
+	assert.True(t, cl.ConnectionState().SupportsDatagrams.Local)
+	assert.True(t, cl.ConnectionState().SupportsDatagrams.Remote)
+	assert.True(t, cl.ConnectionState().SupportsStreamResetPartialDelivery.Local)
+	assert.True(t, cl.ConnectionState().SupportsStreamResetPartialDelivery.Remote)
 	t.Cleanup(func() { cl.CloseWithError(0, "") })
 
 	conn, err := ln.Accept(ctx)
 	require.NoError(t, err)
 	t.Cleanup(func() { conn.CloseWithError(0, "") })
 	select {
 	case <-conn.HandshakeComplete():
-		require.True(t, conn.ConnectionState().SupportsDatagrams.Remote)
+		assert.True(t, conn.ConnectionState().SupportsDatagrams.Local)
+		assert.True(t, conn.ConnectionState().SupportsDatagrams.Remote)
+		assert.True(t, conn.ConnectionState().SupportsStreamResetPartialDelivery.Local)
+		assert.True(t, conn.ConnectionState().SupportsStreamResetPartialDelivery.Remote)
 	case <-ctx.Done():
 		t.Fatal("timeout")
 	}

stream.go

@@ -18,6 +18,7 @@ type quicSendStream interface {
 	CancelWrite(quic.StreamErrorCode)
 	Context() context.Context
 	SetWriteDeadline(time.Time) error
+	SetReliableBoundary()
 }
 
 var (
@@ -44,6 +45,9 @@ type SendStream struct {
 	// Might be initialized to nil if this sendStream is part of an incoming bidirectional stream.
 	streamHdr   []byte
 	streamHdrMu sync.Mutex
+	// Set to true when a goroutine is spawned to send the header asynchronously.
+	// This only happens if the stream is closed / reset immediately after creation.
+	sendingHdrAsync bool
 
 	onClose func() // to remove the stream from the streamsMap
 
@@ -116,23 +120,24 @@ func (s *SendStream) handleSessionGoneError() error {
 }
 
 func (s *SendStream) write(b []byte) (int, error) {
-	if err := s.maybeSendStreamHeader(); err != nil {
+	s.streamHdrMu.Lock()
+	err := s.maybeSendStreamHeader()
+	s.streamHdrMu.Unlock()
+	if err != nil {
 		return 0, err
 	}
 	return s.str.Write(b)
 }
 
 func (s *SendStream) maybeSendStreamHeader() error {
-	s.streamHdrMu.Lock()
-	defer s.streamHdrMu.Unlock()
-
 	if len(s.streamHdr) == 0 {
 		return nil
 	}
 	n, err := s.str.Write(s.streamHdr)
 	if n > 0 {
 		s.streamHdr = s.streamHdr[n:]
 	}
+	s.str.SetReliableBoundary()
 	if err != nil {
 		return err
 	}
@@ -145,6 +150,32 @@ func (s *SendStream) maybeSendStreamHeader() error {
 // Write will unblock immediately, and future calls to Write will fail.
 // When called multiple times it is a no-op.
 func (s *SendStream) CancelWrite(e StreamErrorCode) {
+	// if a Goroutine is already sending the header, return immediately
+	s.streamHdrMu.Lock()
+	if s.sendingHdrAsync {
+		s.streamHdrMu.Unlock()
+		return
+	}
+
+	if len(s.streamHdr) > 0 {
+		// Sending the stream header might block if we are blocked by flow control.
+		// Send a stream header async so that CancelWrite can return immediately.
+		s.sendingHdrAsync = true
+		streamHdr := s.streamHdr
+		s.streamHdr = nil
+		s.streamHdrMu.Unlock()
+
+		go func() {
+			s.SetWriteDeadline(time.Time{})
+			_, _ = s.str.Write(streamHdr)
+			s.str.SetReliableBoundary()
+			s.str.CancelWrite(webtransportCodeToHTTPCode(e))
+			s.onClose()
+		}()
+		return
+	}
+	s.streamHdrMu.Unlock()
+
 	s.str.CancelWrite(webtransportCodeToHTTPCode(e))
 	s.onClose()
 }
@@ -160,9 +191,32 @@ func (s *SendStream) closeWithSession(err error) {
 // Close closes the write-direction of the stream.
 // Future calls to Write are not permitted after calling Close.
 func (s *SendStream) Close() error {
-	if err := s.maybeSendStreamHeader(); err != nil {
-		return err
+	// if a Goroutine is already sending the header, return immediately
+	s.streamHdrMu.Lock()
+	if s.sendingHdrAsync {
+		s.streamHdrMu.Unlock()
+		return nil
 	}
+
+	if len(s.streamHdr) > 0 {
+		// Sending the stream header might block if we are blocked by flow control.
+		// Send a stream header async so that CancelWrite can return immediately.
+		s.sendingHdrAsync = true
+		streamHdr := s.streamHdr
+		s.streamHdr = nil
+		s.streamHdrMu.Unlock()
+
+		go func() {
+			s.SetWriteDeadline(time.Time{})
+			_, _ = s.str.Write(streamHdr)
+			s.str.SetReliableBoundary()
+			_ = s.str.Close()
+			s.onClose()
+		}()
+		return nil
+	}
+	s.streamHdrMu.Unlock()
+
 	s.onClose()
 	return maybeConvertStreamError(s.str.Close())
 }