github.com/MerlinKodo/quic-go@v0.39.2/http3/http_stream.go (about)

     1  package http3
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  
     7  	"github.com/MerlinKodo/quic-go"
     8  	"github.com/MerlinKodo/quic-go/internal/utils"
     9  )
    10  
    11  // A Stream is a HTTP/3 stream.
    12  // When writing to and reading from the stream, data is framed in HTTP/3 DATA frames.
    13  type Stream quic.Stream
    14  
    15  // The stream conforms to the quic.Stream interface, but instead of writing to and reading directly
    16  // from the QUIC stream, it writes to and reads from the HTTP stream.
    17  type stream struct {
    18  	quic.Stream
    19  
    20  	buf []byte
    21  
    22  	onFrameError          func()
    23  	bytesRemainingInFrame uint64
    24  }
    25  
    26  var _ Stream = &stream{}
    27  
    28  func newStream(str quic.Stream, onFrameError func()) *stream {
    29  	return &stream{
    30  		Stream:       str,
    31  		onFrameError: onFrameError,
    32  		buf:          make([]byte, 0, 16),
    33  	}
    34  }
    35  
    36  func (s *stream) Read(b []byte) (int, error) {
    37  	if s.bytesRemainingInFrame == 0 {
    38  	parseLoop:
    39  		for {
    40  			frame, err := parseNextFrame(s.Stream, nil)
    41  			if err != nil {
    42  				return 0, err
    43  			}
    44  			switch f := frame.(type) {
    45  			case *headersFrame:
    46  				// skip HEADERS frames
    47  				continue
    48  			case *dataFrame:
    49  				s.bytesRemainingInFrame = f.Length
    50  				break parseLoop
    51  			default:
    52  				s.onFrameError()
    53  				// parseNextFrame skips over unknown frame types
    54  				// Therefore, this condition is only entered when we parsed another known frame type.
    55  				return 0, fmt.Errorf("peer sent an unexpected frame: %T", f)
    56  			}
    57  		}
    58  	}
    59  
    60  	var n int
    61  	var err error
    62  	if s.bytesRemainingInFrame < uint64(len(b)) {
    63  		n, err = s.Stream.Read(b[:s.bytesRemainingInFrame])
    64  	} else {
    65  		n, err = s.Stream.Read(b)
    66  	}
    67  	s.bytesRemainingInFrame -= uint64(n)
    68  	return n, err
    69  }
    70  
    71  func (s *stream) hasMoreData() bool {
    72  	return s.bytesRemainingInFrame > 0
    73  }
    74  
    75  func (s *stream) Write(b []byte) (int, error) {
    76  	s.buf = s.buf[:0]
    77  	s.buf = (&dataFrame{Length: uint64(len(b))}).Append(s.buf)
    78  	if _, err := s.Stream.Write(s.buf); err != nil {
    79  		return 0, err
    80  	}
    81  	return s.Stream.Write(b)
    82  }
    83  
    84  var errTooMuchData = errors.New("peer sent too much data")
    85  
    86  type lengthLimitedStream struct {
    87  	*stream
    88  	contentLength int64
    89  	read          int64
    90  	resetStream   bool
    91  }
    92  
    93  var _ Stream = &lengthLimitedStream{}
    94  
    95  func newLengthLimitedStream(str *stream, contentLength int64) *lengthLimitedStream {
    96  	return &lengthLimitedStream{
    97  		stream:        str,
    98  		contentLength: contentLength,
    99  	}
   100  }
   101  
   102  func (s *lengthLimitedStream) checkContentLengthViolation() error {
   103  	if s.read > s.contentLength || s.read == s.contentLength && s.hasMoreData() {
   104  		if !s.resetStream {
   105  			s.CancelRead(quic.StreamErrorCode(ErrCodeMessageError))
   106  			s.CancelWrite(quic.StreamErrorCode(ErrCodeMessageError))
   107  			s.resetStream = true
   108  		}
   109  		return errTooMuchData
   110  	}
   111  	return nil
   112  }
   113  
   114  func (s *lengthLimitedStream) Read(b []byte) (int, error) {
   115  	if err := s.checkContentLengthViolation(); err != nil {
   116  		return 0, err
   117  	}
   118  	n, err := s.stream.Read(b[:utils.Min(int64(len(b)), s.contentLength-s.read)])
   119  	s.read += int64(n)
   120  	if err := s.checkContentLengthViolation(); err != nil {
   121  		return n, err
   122  	}
   123  	return n, err
   124  }