add vendor data

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

@@ -0,0 +1,478 @@
+package webtransport
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/quic-go/quic-go"
+)
+
+type quicSendStream interface {
+	io.WriteCloser
+	CancelWrite(quic.StreamErrorCode)
+	Context() context.Context
+	SetWriteDeadline(time.Time) error
+	SetReliableBoundary()
+}
+
+var (
+	_ quicSendStream = &quic.SendStream{}
+	_ quicSendStream = &quic.Stream{}
+)
+
+type quicReceiveStream interface {
+	io.Reader
+	CancelRead(quic.StreamErrorCode)
+	SetReadDeadline(time.Time) error
+}
+
+var (
+	_ quicReceiveStream = &quic.ReceiveStream{}
+	_ quicReceiveStream = &quic.Stream{}
+)
+
+// A SendStream is a unidirectional WebTransport send stream.
+type SendStream struct {
+	str quicSendStream
+	// WebTransport stream header.
+	// Set by the constructor, set to nil once sent out.
+	// 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
+
+	closeOnce sync.Once
+	closed    chan struct{}
+	closeErr  error
+
+	deadlineMu       sync.Mutex
+	writeDeadline    time.Time
+	deadlineNotifyCh chan struct{} // receives a value when deadline changes
+}
+
+func newSendStream(str quicSendStream, hdr []byte, onClose func()) *SendStream {
+	return &SendStream{
+		str:       str,
+		closed:    make(chan struct{}),
+		streamHdr: hdr,
+		onClose:   onClose,
+	}
+}
+
+// Write writes data to the stream.
+// Write can be made to time out using [SendStream.SetWriteDeadline].
+// If the stream was canceled, the error is a [StreamError].
+func (s *SendStream) Write(b []byte) (int, error) {
+	n, err := s.write(b)
+	if err != nil && !isTimeoutError(err) {
+		s.onClose()
+	}
+	var strErr *quic.StreamError
+	if errors.As(err, &strErr) && strErr.ErrorCode == WTSessionGoneErrorCode {
+		return n, s.handleSessionGoneError()
+	}
+	return n, maybeConvertStreamError(err)
+}
+
+// handleSessionGoneError waits for the session to be closed after receiving a WTSessionGoneErrorCode.
+// If the peer is initiating the session close, we might need to wait for the CONNECT stream to be closed.
+// While a malicious peer might withhold the session close, this is not an interesting attack vector:
+// 1. a WebTransport stream consumes very little memory, and
+// 2. the number of concurrent WebTransport sessions is limited.
+func (s *SendStream) handleSessionGoneError() error {
+	s.deadlineMu.Lock()
+	if s.deadlineNotifyCh == nil {
+		s.deadlineNotifyCh = make(chan struct{}, 1)
+	}
+	s.deadlineMu.Unlock()
+
+	for {
+		s.deadlineMu.Lock()
+		deadline := s.writeDeadline
+		s.deadlineMu.Unlock()
+
+		var timerCh <-chan time.Time
+		if !deadline.IsZero() {
+			if d := time.Until(deadline); d > 0 {
+				timerCh = time.After(d)
+			} else {
+				return os.ErrDeadlineExceeded
+			}
+		}
+		select {
+		case <-s.closed:
+			return s.closeErr
+		case <-timerCh:
+			return os.ErrDeadlineExceeded
+		case <-s.deadlineNotifyCh:
+		}
+	}
+}
+
+func (s *SendStream) write(b []byte) (int, error) {
+	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 {
+	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
+	}
+	s.streamHdr = nil
+	return nil
+}
+
+// CancelWrite aborts sending on this stream.
+// Data already written, but not yet delivered to the peer is not guaranteed to be delivered reliably.
+// 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()
+}
+
+func (s *SendStream) closeWithSession(err error) {
+	s.closeOnce.Do(func() {
+		s.closeErr = err
+		s.str.CancelWrite(WTSessionGoneErrorCode)
+		close(s.closed)
+	})
+}
+
+// Close closes the write-direction of the stream.
+// Future calls to Write are not permitted after calling Close.
+func (s *SendStream) Close() error {
+	// 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())
+}
+
+// The Context is canceled as soon as the write-side of the stream is closed.
+// This happens when Close() or CancelWrite() is called, or when the peer
+// cancels the read-side of their stream.
+// The cancellation cause is set to the error that caused the stream to
+// close, or `context.Canceled` in case the stream is closed without error.
+func (s *SendStream) Context() context.Context {
+	return s.str.Context()
+}
+
+// SetWriteDeadline sets the deadline for future Write calls
+// and any currently-blocked Write call.
+// Even if write times out, it may return n > 0, indicating that
+// some data was successfully written.
+// A zero value for t means Write will not time out.
+func (s *SendStream) SetWriteDeadline(t time.Time) error {
+	s.deadlineMu.Lock()
+	s.writeDeadline = t
+	if s.deadlineNotifyCh != nil {
+		select {
+		case s.deadlineNotifyCh <- struct{}{}:
+		default:
+		}
+	}
+	s.deadlineMu.Unlock()
+
+	return maybeConvertStreamError(s.str.SetWriteDeadline(t))
+}
+
+// A ReceiveStream is a unidirectional WebTransport receive stream.
+type ReceiveStream struct {
+	str quicReceiveStream
+
+	onClose func() // to remove the stream from the streamsMap
+
+	closeOnce sync.Once
+	closed    chan struct{}
+	closeErr  error
+
+	deadlineMu       sync.Mutex
+	readDeadline     time.Time
+	deadlineNotifyCh chan struct{} // receives a value when deadline changes
+}
+
+func newReceiveStream(str quicReceiveStream, onClose func()) *ReceiveStream {
+	return &ReceiveStream{
+		str:     str,
+		closed:  make(chan struct{}),
+		onClose: onClose,
+	}
+}
+
+// Read reads data from the stream.
+// Read can be made to time out using [ReceiveStream.SetReadDeadline].
+// If the stream was canceled, the error is a [StreamError].
+func (s *ReceiveStream) Read(b []byte) (int, error) {
+	n, err := s.str.Read(b)
+	if err != nil && !isTimeoutError(err) {
+		s.onClose()
+	}
+	var strErr *quic.StreamError
+	if errors.As(err, &strErr) && strErr.ErrorCode == WTSessionGoneErrorCode {
+		return n, s.handleSessionGoneError()
+	}
+	return n, maybeConvertStreamError(err)
+}
+
+// handleSessionGoneError waits for the session to be closed after receiving a WTSessionGoneErrorCode.
+// If the peer is initiating the session close, we might need to wait for the CONNECT stream to be closed.
+// While a malicious peer might withhold the session close, this is not an interesting attack vector:
+// 1. a WebTransport stream consumes very little memory, and
+// 2. the number of concurrent WebTransport sessions is limited.
+func (s *ReceiveStream) handleSessionGoneError() error {
+	s.deadlineMu.Lock()
+	if s.deadlineNotifyCh == nil {
+		s.deadlineNotifyCh = make(chan struct{}, 1)
+	}
+	s.deadlineMu.Unlock()
+
+	for {
+		s.deadlineMu.Lock()
+		deadline := s.readDeadline
+		s.deadlineMu.Unlock()
+
+		var timerCh <-chan time.Time
+		if !deadline.IsZero() {
+			if d := time.Until(deadline); d > 0 {
+				timerCh = time.After(d)
+			} else {
+				return os.ErrDeadlineExceeded
+			}
+		}
+		select {
+		case <-s.closed:
+			return s.closeErr
+		case <-timerCh:
+			return os.ErrDeadlineExceeded
+		case <-s.deadlineNotifyCh:
+		}
+	}
+}
+
+// CancelRead aborts receiving on this stream.
+// It instructs the peer to stop transmitting stream data.
+// Read will unblock immediately, and future Read calls will fail.
+// When called multiple times it is a no-op.
+func (s *ReceiveStream) CancelRead(e StreamErrorCode) {
+	s.str.CancelRead(webtransportCodeToHTTPCode(e))
+	s.onClose()
+}
+
+func (s *ReceiveStream) closeWithSession(err error) {
+	s.closeOnce.Do(func() {
+		s.closeErr = err
+		s.str.CancelRead(WTSessionGoneErrorCode)
+		close(s.closed)
+	})
+}
+
+// SetReadDeadline sets the deadline for future Read calls and
+// any currently-blocked Read call.
+// A zero value for t means Read will not time out.
+func (s *ReceiveStream) SetReadDeadline(t time.Time) error {
+	s.deadlineMu.Lock()
+	s.readDeadline = t
+	if s.deadlineNotifyCh != nil {
+		select {
+		case s.deadlineNotifyCh <- struct{}{}:
+		default:
+		}
+	}
+	s.deadlineMu.Unlock()
+
+	return maybeConvertStreamError(s.str.SetReadDeadline(t))
+}
+
+// Stream is a bidirectional WebTransport stream.
+type Stream struct {
+	sendStr *SendStream
+	recvStr *ReceiveStream
+
+	mx                             sync.Mutex
+	sendSideClosed, recvSideClosed bool
+	onClose                        func()
+}
+
+func newStream(str *quic.Stream, hdr []byte, onClose func()) *Stream {
+	s := &Stream{onClose: onClose}
+	s.sendStr = newSendStream(str, hdr, func() { s.registerClose(true) })
+	s.recvStr = newReceiveStream(str, func() { s.registerClose(false) })
+	return s
+}
+
+// Write writes data to the stream.
+// Write can be made to time out using [Stream.SetWriteDeadline] or [Stream.SetDeadline].
+// If the stream was canceled, the error is a [StreamError].
+func (s *Stream) Write(b []byte) (int, error) {
+	return s.sendStr.Write(b)
+}
+
+// Read reads data from the stream.
+// Read can be made to time out using [Stream.SetReadDeadline] and [Stream.SetDeadline].
+// If the stream was canceled, the error is a [StreamError].
+func (s *Stream) Read(b []byte) (int, error) {
+	return s.recvStr.Read(b)
+}
+
+// CancelWrite aborts sending on this stream.
+// See [SendStream.CancelWrite] for more details.
+func (s *Stream) CancelWrite(e StreamErrorCode) {
+	s.sendStr.CancelWrite(e)
+}
+
+// CancelRead aborts receiving on this stream.
+// See [ReceiveStream.CancelRead] for more details.
+func (s *Stream) CancelRead(e StreamErrorCode) {
+	s.recvStr.CancelRead(e)
+}
+
+// Close closes the send-direction of the stream.
+// It does not close the receive-direction of the stream.
+func (s *Stream) Close() error {
+	return s.sendStr.Close()
+}
+
+func (s *Stream) registerClose(isSendSide bool) {
+	s.mx.Lock()
+	if isSendSide {
+		s.sendSideClosed = true
+	} else {
+		s.recvSideClosed = true
+	}
+	isClosed := s.sendSideClosed && s.recvSideClosed
+	s.mx.Unlock()
+
+	if isClosed {
+		s.onClose()
+	}
+}
+
+func (s *Stream) closeWithSession(err error) {
+	s.sendStr.closeWithSession(err)
+	s.recvStr.closeWithSession(err)
+}
+
+// The Context is canceled as soon as the write-side of the stream is closed.
+// See [SendStream.Context] for more details.
+func (s *Stream) Context() context.Context {
+	return s.sendStr.Context()
+}
+
+// SetWriteDeadline sets the deadline for future Write calls.
+// See [SendStream.SetWriteDeadline] for more details.
+func (s *Stream) SetWriteDeadline(t time.Time) error {
+	return s.sendStr.SetWriteDeadline(t)
+}
+
+// SetReadDeadline sets the deadline for future Read calls.
+// See [ReceiveStream.SetReadDeadline] for more details.
+func (s *Stream) SetReadDeadline(t time.Time) error {
+	return s.recvStr.SetReadDeadline(t)
+}
+
+// SetDeadline sets the read and write deadlines associated with the stream.
+// It is equivalent to calling both SetReadDeadline and SetWriteDeadline.
+func (s *Stream) SetDeadline(t time.Time) error {
+	err1 := s.SetWriteDeadline(t)
+	err2 := s.SetReadDeadline(t)
+	return errors.Join(err1, err2)
+}
+
+func maybeConvertStreamError(err error) error {
+	if err == nil {
+		return nil
+	}
+	var streamErr *quic.StreamError
+	if errors.As(err, &streamErr) {
+		errorCode, cerr := httpCodeToWebtransportCode(streamErr.ErrorCode)
+		if cerr != nil {
+			return fmt.Errorf("stream reset, but failed to convert stream error %d: %w", streamErr.ErrorCode, cerr)
+		}
+		return &StreamError{
+			ErrorCode: errorCode,
+			Remote:    streamErr.Remote,
+		}
+	}
+	return err
+}
+
+func isTimeoutError(err error) bool {
+	nerr, ok := err.(net.Error)
+	if !ok {
+		return false
+	}
+	return nerr.Timeout()
+}