immediately reject streams for closed sessions (#235)

Author: marten-seemann

session_manager.go

@@ -2,6 +2,7 @@ package webtransport
 
 import (
 	"context"
+	"slices"
 	"sync"
 	"time"
 
@@ -21,11 +22,14 @@ type sessionEntry struct {
 	Session       *Session
 }
 
+const maxRecentlyClosedSessions = 16
+
 type sessionManager struct {
 	timeout time.Duration
 
-	mx       sync.Mutex
-	sessions map[sessionID]sessionEntry
+	mx                     sync.Mutex
+	sessions               map[sessionID]sessionEntry
+	recentlyClosedSessions []sessionID
 }
 
 func newSessionManager(timeout time.Duration) *sessionManager {
@@ -45,6 +49,13 @@ func (m *sessionManager) AddStream(str *quic.Stream, id sessionID) {
 
 	entry, ok := m.sessions[id]
 	if !ok {
+		// Receiving a stream for an unknown session is expected to be rare,
+		// so the performance impact of searching through the slice is negligible.
+		if slices.Contains(m.recentlyClosedSessions, id) {
+			str.CancelRead(WTBufferedStreamRejectedErrorCode)
+			str.CancelWrite(WTBufferedStreamRejectedErrorCode)
+			return
+		}
 		entry = sessionEntry{Unestablished: &unestablishedSession{}}
 		m.sessions[id] = entry
 	}
@@ -67,6 +78,12 @@ func (m *sessionManager) AddUniStream(str *quic.ReceiveStream, id sessionID) {
 
 	entry, ok := m.sessions[id]
 	if !ok {
+		// Receiving a stream for an unknown session is expected to be rare,
+		// so the performance impact of searching through the slice is negligible.
+		if slices.Contains(m.recentlyClosedSessions, id) {
+			str.CancelRead(WTBufferedStreamRejectedErrorCode)
+			return
+		}
 		entry = sessionEntry{Unestablished: &unestablishedSession{}}
 		m.sessions[id] = entry
 	}
@@ -134,12 +151,21 @@ func (m *sessionManager) AddSession(id sessionID, s *Session) {
 	m.sessions[id] = sessionEntry{Session: s}
 
 	context.AfterFunc(s.Context(), func() {
-		m.mx.Lock()
-		defer m.mx.Unlock()
-		delete(m.sessions, id)
+		m.deleteSession(id)
 	})
 }
 
+func (m *sessionManager) deleteSession(id sessionID) {
+	m.mx.Lock()
+	defer m.mx.Unlock()
+
+	delete(m.sessions, id)
+	m.recentlyClosedSessions = append(m.recentlyClosedSessions, id)
+	if len(m.recentlyClosedSessions) > maxRecentlyClosedSessions {
+		m.recentlyClosedSessions = m.recentlyClosedSessions[1:]
+	}
+}
+
 func (m *sessionManager) Close() {
 	m.mx.Lock()
 	defer m.mx.Unlock()

session_manager_test.go

@@ -3,10 +3,14 @@ package webtransport
 import (
 	"context"
 	"io"
+	"maps"
+	"math/rand/v2"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"runtime"
+	"slices"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -15,6 +19,7 @@ import (
 
 	"github.com/quic-go/quic-go"
 	"github.com/quic-go/quic-go/http3"
+	"github.com/quic-go/quic-go/quicvarint"
 	"github.com/quic-go/quic-go/testutils/simnet"
 
 	"github.com/stretchr/testify/require"
@@ -257,3 +262,150 @@ func TestSessionManagerStreamReordering(t *testing.T) {
 		require.Equal(t, []byte("amet"), data)
 	})
 }
+
+func TestSessionManagerSessionClose(t *testing.T) {
+	// synctest works slightly differently on Go 1.24,
+	// so we skip the test
+	if strings.HasPrefix(runtime.Version(), "go1.24") {
+		t.Skip("skipping on Go 1.24 due to synctest issues")
+	}
+
+	synctest.Test(t, func(t *testing.T) {
+		const rtt = 10 * time.Millisecond
+		clientPacketConn, serverPacketConn, closeFn := newSimnetLink(t, rtt)
+		defer closeFn(t)
+		clientConn, serverConn := newConnPair(t, clientPacketConn, serverPacketConn)
+		t.Cleanup(func() {
+			clientConn.CloseWithError(0, "")
+			serverConn.CloseWithError(0, "")
+		})
+
+		tr := http3.Transport{}
+		cc := tr.NewRawClientConn(clientConn)
+
+		server := http3.Server{
+			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusOK)
+				w.(http.Flusher).Flush()
+				<-r.Context().Done()
+			}),
+		}
+		sc, err := server.NewRawServerConn(serverConn)
+		require.NoError(t, err)
+
+		openSession := func() (*http3.RequestStream, error) {
+			reqStr, err := cc.OpenRequestStream(context.Background())
+			if err != nil {
+				return nil, err
+			}
+			if err := reqStr.SendRequestHeader(httptest.NewRequest(http.MethodGet, "/", nil)); err != nil {
+				return nil, err
+			}
+			serverReqStr, err := serverConn.AcceptStream(context.Background())
+			if err != nil {
+				return nil, err
+			}
+			go sc.HandleRequestStream(serverReqStr)
+			if _, err := reqStr.ReadResponse(); err != nil {
+				return nil, err
+			}
+			return reqStr, nil
+		}
+
+		sessMgr := newSessionManager(time.Hour)
+
+		reqStrs := make(map[sessionID]*http3.RequestStream)
+		for range maxRecentlyClosedSessions {
+			reqStr, err := openSession()
+			require.NoError(t, err)
+			sessID := sessionID(rand.Int64N(quicvarint.Max))
+			reqStrs[sessID] = reqStr
+			sess := newSession(context.Background(), sessID, clientConn, reqStr, "")
+			sessMgr.AddSession(sessID, sess)
+			synctest.Wait()
+		}
+		require.Equal(t, maxRecentlyClosedSessions, sessMgr.NumSessions())
+
+		// close a random session
+		sessID := slices.Collect(maps.Keys(reqStrs))[rand.Int64N(int64(len(reqStrs)))]
+		reqStrs[sessID].CancelWrite(0)
+		reqStrs[sessID].CancelRead(0)
+		synctest.Wait()
+		require.Equal(t, maxRecentlyClosedSessions-1, sessMgr.NumSessions())
+		delete(reqStrs, sessID)
+
+		// Consume the HTTP/3 control stream that the server opened during setup.
+		// This is necessary because we're calling AcceptUniStream directly on the QUIC connection,
+		// which returns ALL unidirectional streams, not just WebTransport streams.
+		_, err = clientConn.AcceptUniStream(context.Background())
+		require.NoError(t, err)
+
+		// enqueue streams for the remaining sessions
+		for sessID := range reqStrs {
+			// bidirectional stream
+			serverStr, err := serverConn.OpenStream()
+			require.NoError(t, err)
+			_, err = serverStr.Write([]byte("stream " + strconv.Itoa(int(sessID))))
+			require.NoError(t, err)
+			clientStr, err := clientConn.AcceptStream(context.Background())
+			require.NoError(t, err)
+			sessMgr.AddStream(clientStr, sessID)
+			synctest.Wait()
+			// make sure the stream is not rejected
+			select {
+			case <-clientStr.Context().Done():
+				require.Fail(t, "stream should not be rejected")
+			case <-serverStr.Context().Done():
+				require.Fail(t, "stream should not be rejected")
+			default:
+			}
+
+			// unidirectional stream
+			serverUniStr, err := serverConn.OpenUniStream()
+			require.NoError(t, err)
+			_, err = serverUniStr.Write([]byte("unistream " + strconv.Itoa(int(sessID))))
+			require.NoError(t, err)
+			clientUniStr, err := clientConn.AcceptUniStream(context.Background())
+			require.NoError(t, err)
+			sessMgr.AddUniStream(clientUniStr, sessID)
+			synctest.Wait()
+			// make sure the stream is not rejected
+			select {
+			case <-serverUniStr.Context().Done():
+				require.Fail(t, "unidirectional stream should not be rejected")
+			default:
+			}
+		}
+
+		// test that streams for the closed session are immediately rejected
+		start := time.Now()
+		// bidirectional stream
+		serverStr, err := serverConn.OpenStream()
+		require.NoError(t, err)
+		_, err = serverStr.Write([]byte("stream " + strconv.Itoa(int(sessID))))
+		require.NoError(t, err)
+		clientStrClosed, err := clientConn.AcceptStream(context.Background())
+		require.NoError(t, err)
+		sessMgr.AddStream(clientStrClosed, sessID)
+		synctest.Wait()
+		// make sure the stream is immediately rejected
+		_, err = serverStr.Read([]byte{0})
+		require.ErrorIs(t, err, &quic.StreamError{Remote: true, StreamID: serverStr.StreamID(), ErrorCode: WTBufferedStreamRejectedErrorCode})
+		require.Equal(t, rtt, time.Since(start))
+
+		// unidirectional stream
+		start = time.Now()
+		serverUniStr, err := serverConn.OpenUniStream()
+		require.NoError(t, err)
+		_, err = serverUniStr.Write([]byte("unistream " + strconv.Itoa(int(sessID))))
+		require.NoError(t, err)
+		synctest.Wait()
+		clientUniStr, err := clientConn.AcceptUniStream(context.Background())
+		require.NoError(t, err)
+		sessMgr.AddUniStream(clientUniStr, sessID)
+		synctest.Wait()
+		_, err = clientUniStr.Read([]byte{0})
+		require.ErrorIs(t, err, &quic.StreamError{Remote: false, StreamID: clientUniStr.StreamID(), ErrorCode: WTBufferedStreamRejectedErrorCode})
+		require.Equal(t, rtt/2, time.Since(start))
+	})
+}