github.com/metacubex/quic-go@v0.44.1-0.20240520163451-20b689a59136/http3/capsule.go (about)

     1  package http3
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/metacubex/quic-go/quicvarint"
     7  )
     8  
     9  // CapsuleType is the type of the capsule.
    10  type CapsuleType uint64
    11  
    12  type exactReader struct {
    13  	R *io.LimitedReader
    14  }
    15  
    16  func (r *exactReader) Read(b []byte) (int, error) {
    17  	n, err := r.R.Read(b)
    18  	if r.R.N > 0 {
    19  		return n, io.ErrUnexpectedEOF
    20  	}
    21  	return n, err
    22  }
    23  
    24  type countingByteReader struct {
    25  	io.ByteReader
    26  	Read int
    27  }
    28  
    29  func (r *countingByteReader) ReadByte() (byte, error) {
    30  	b, err := r.ByteReader.ReadByte()
    31  	if err == nil {
    32  		r.Read++
    33  	}
    34  	return b, err
    35  }
    36  
    37  // ParseCapsule parses the header of a Capsule.
    38  // It returns an io.LimitedReader that can be used to read the Capsule value.
    39  // The Capsule value must be read entirely (i.e. until the io.EOF) before using r again.
    40  func ParseCapsule(r quicvarint.Reader) (CapsuleType, io.Reader, error) {
    41  	cbr := countingByteReader{ByteReader: r}
    42  	ct, err := quicvarint.Read(&cbr)
    43  	if err != nil {
    44  		// If an io.EOF is returned without consuming any bytes, return it unmodified.
    45  		// Otherwise, return an io.ErrUnexpectedEOF.
    46  		if err == io.EOF && cbr.Read > 0 {
    47  			return 0, nil, io.ErrUnexpectedEOF
    48  		}
    49  		return 0, nil, err
    50  	}
    51  	l, err := quicvarint.Read(r)
    52  	if err != nil {
    53  		if err == io.EOF {
    54  			return 0, nil, io.ErrUnexpectedEOF
    55  		}
    56  		return 0, nil, err
    57  	}
    58  	return CapsuleType(ct), &exactReader{R: io.LimitReader(r, int64(l)).(*io.LimitedReader)}, nil
    59  }
    60  
    61  // WriteCapsule writes a capsule
    62  func WriteCapsule(w quicvarint.Writer, ct CapsuleType, value []byte) error {
    63  	b := make([]byte, 0, 16)
    64  	b = quicvarint.Append(b, uint64(ct))
    65  	b = quicvarint.Append(b, uint64(len(value)))
    66  	if _, err := w.Write(b); err != nil {
    67  		return err
    68  	}
    69  	_, err := w.Write(value)
    70  	return err
    71  }