add vendor data

vendor/github.com/quic-go/webtransport-go/session_manager.go

@@ -0,0 +1,179 @@
+package webtransport
+
+import (
+	"context"
+	"slices"
+	"sync"
+	"time"
+
+	"github.com/quic-go/quic-go"
+)
+
+type unestablishedSession struct {
+	Streams    []*quic.Stream
+	UniStreams []*quic.ReceiveStream
+
+	Timer *time.Timer
+}
+
+type sessionEntry struct {
+	// at any point in time, only one of these will be non-nil
+	Unestablished *unestablishedSession
+	Session       *Session
+}
+
+const maxRecentlyClosedSessions = 16
+
+type sessionManager struct {
+	timeout time.Duration
+
+	mx                     sync.Mutex
+	sessions               map[sessionID]sessionEntry
+	recentlyClosedSessions []sessionID
+}
+
+func newSessionManager(timeout time.Duration) *sessionManager {
+	return &sessionManager{
+		timeout:  timeout,
+		sessions: make(map[sessionID]sessionEntry),
+	}
+}
+
+// AddStream adds a new bidirectional stream to a WebTransport session.
+// If the WebTransport session has not yet been established,
+// the stream is buffered until the session is established.
+// If that takes longer than timeout, the stream is reset.
+func (m *sessionManager) AddStream(str *quic.Stream, id sessionID) {
+	m.mx.Lock()
+	defer m.mx.Unlock()
+
+	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
+	}
+	if entry.Session != nil {
+		entry.Session.addIncomingStream(str)
+		return
+	}
+
+	entry.Unestablished.Streams = append(entry.Unestablished.Streams, str)
+	m.resetTimer(id)
+}
+
+// AddUniStream adds a new unidirectional stream to a WebTransport session.
+// If the WebTransport session has not yet been established,
+// the stream is buffered until the session is established.
+// If that takes longer than timeout, the stream is reset.
+func (m *sessionManager) AddUniStream(str *quic.ReceiveStream, id sessionID) {
+	m.mx.Lock()
+	defer m.mx.Unlock()
+
+	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
+	}
+	if entry.Session != nil {
+		entry.Session.addIncomingUniStream(str)
+		return
+	}
+
+	entry.Unestablished.UniStreams = append(entry.Unestablished.UniStreams, str)
+	m.resetTimer(id)
+}
+
+func (m *sessionManager) resetTimer(id sessionID) {
+	entry := m.sessions[id]
+	if entry.Unestablished.Timer != nil {
+		entry.Unestablished.Timer.Reset(m.timeout)
+		return
+	}
+	entry.Unestablished.Timer = time.AfterFunc(m.timeout, func() { m.onTimer(id) })
+}
+
+func (m *sessionManager) onTimer(id sessionID) {
+	m.mx.Lock()
+	defer m.mx.Unlock()
+
+	sessionEntry, ok := m.sessions[id]
+	if !ok { // session already closed
+		return
+	}
+	if sessionEntry.Session != nil { // session already established
+		return
+	}
+	for _, str := range sessionEntry.Unestablished.Streams {
+		str.CancelRead(WTBufferedStreamRejectedErrorCode)
+		str.CancelWrite(WTBufferedStreamRejectedErrorCode)
+	}
+	for _, uniStr := range sessionEntry.Unestablished.UniStreams {
+		uniStr.CancelRead(WTBufferedStreamRejectedErrorCode)
+	}
+	delete(m.sessions, id)
+}
+
+// AddSession adds a new WebTransport session.
+func (m *sessionManager) AddSession(id sessionID, s *Session) {
+	m.mx.Lock()
+	defer m.mx.Unlock()
+
+	entry, ok := m.sessions[id]
+
+	if ok && entry.Unestablished != nil {
+		// We might already have an entry of this session.
+		// This can happen when we receive streams for this WebTransport session before we complete
+		// the Extended CONNECT request.
+		for _, str := range entry.Unestablished.Streams {
+			s.addIncomingStream(str)
+		}
+		for _, uniStr := range entry.Unestablished.UniStreams {
+			s.addIncomingUniStream(uniStr)
+		}
+		if entry.Unestablished.Timer != nil {
+			entry.Unestablished.Timer.Stop()
+		}
+		entry.Unestablished = nil
+	}
+	m.sessions[id] = sessionEntry{Session: s}
+
+	context.AfterFunc(s.Context(), func() {
+		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()
+
+	for _, entry := range m.sessions {
+		if entry.Unestablished != nil && entry.Unestablished.Timer != nil {
+			entry.Unestablished.Timer.Stop()
+		}
+	}
+	clear(m.sessions)
+}