github.com/apernet/quic-go@v0.43.1-0.20240515053213-5e9e635fd9f0/quicvarint/varint.go (about)

     1  package quicvarint
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  )
     7  
     8  // taken from the QUIC draft
     9  const (
    10  	// Min is the minimum value allowed for a QUIC varint.
    11  	Min = 0
    12  
    13  	// Max is the maximum allowed value for a QUIC varint (2^62-1).
    14  	Max = maxVarInt8
    15  
    16  	maxVarInt1 = 63
    17  	maxVarInt2 = 16383
    18  	maxVarInt4 = 1073741823
    19  	maxVarInt8 = 4611686018427387903
    20  )
    21  
    22  // Read reads a number in the QUIC varint format from r.
    23  func Read(r io.ByteReader) (uint64, error) {
    24  	firstByte, err := r.ReadByte()
    25  	if err != nil {
    26  		return 0, err
    27  	}
    28  	// the first two bits of the first byte encode the length
    29  	len := 1 << ((firstByte & 0xc0) >> 6)
    30  	b1 := firstByte & (0xff - 0xc0)
    31  	if len == 1 {
    32  		return uint64(b1), nil
    33  	}
    34  	b2, err := r.ReadByte()
    35  	if err != nil {
    36  		return 0, err
    37  	}
    38  	if len == 2 {
    39  		return uint64(b2) + uint64(b1)<<8, nil
    40  	}
    41  	b3, err := r.ReadByte()
    42  	if err != nil {
    43  		return 0, err
    44  	}
    45  	b4, err := r.ReadByte()
    46  	if err != nil {
    47  		return 0, err
    48  	}
    49  	if len == 4 {
    50  		return uint64(b4) + uint64(b3)<<8 + uint64(b2)<<16 + uint64(b1)<<24, nil
    51  	}
    52  	b5, err := r.ReadByte()
    53  	if err != nil {
    54  		return 0, err
    55  	}
    56  	b6, err := r.ReadByte()
    57  	if err != nil {
    58  		return 0, err
    59  	}
    60  	b7, err := r.ReadByte()
    61  	if err != nil {
    62  		return 0, err
    63  	}
    64  	b8, err := r.ReadByte()
    65  	if err != nil {
    66  		return 0, err
    67  	}
    68  	return uint64(b8) + uint64(b7)<<8 + uint64(b6)<<16 + uint64(b5)<<24 + uint64(b4)<<32 + uint64(b3)<<40 + uint64(b2)<<48 + uint64(b1)<<56, nil
    69  }
    70  
    71  // Append appends i in the QUIC varint format.
    72  func Append(b []byte, i uint64) []byte {
    73  	if i <= maxVarInt1 {
    74  		return append(b, uint8(i))
    75  	}
    76  	if i <= maxVarInt2 {
    77  		return append(b, []byte{uint8(i>>8) | 0x40, uint8(i)}...)
    78  	}
    79  	if i <= maxVarInt4 {
    80  		return append(b, []byte{uint8(i>>24) | 0x80, uint8(i >> 16), uint8(i >> 8), uint8(i)}...)
    81  	}
    82  	if i <= maxVarInt8 {
    83  		return append(b, []byte{
    84  			uint8(i>>56) | 0xc0, uint8(i >> 48), uint8(i >> 40), uint8(i >> 32),
    85  			uint8(i >> 24), uint8(i >> 16), uint8(i >> 8), uint8(i),
    86  		}...)
    87  	}
    88  	panic(fmt.Sprintf("%#x doesn't fit into 62 bits", i))
    89  }
    90  
    91  // AppendWithLen append i in the QUIC varint format with the desired length.
    92  func AppendWithLen(b []byte, i uint64, length int) []byte {
    93  	if length != 1 && length != 2 && length != 4 && length != 8 {
    94  		panic("invalid varint length")
    95  	}
    96  	l := Len(i)
    97  	if l == length {
    98  		return Append(b, i)
    99  	}
   100  	if l > length {
   101  		panic(fmt.Sprintf("cannot encode %d in %d bytes", i, length))
   102  	}
   103  	if length == 2 {
   104  		b = append(b, 0b01000000)
   105  	} else if length == 4 {
   106  		b = append(b, 0b10000000)
   107  	} else if length == 8 {
   108  		b = append(b, 0b11000000)
   109  	}
   110  	for j := 1; j < length-l; j++ {
   111  		b = append(b, 0)
   112  	}
   113  	for j := 0; j < l; j++ {
   114  		b = append(b, uint8(i>>(8*(l-1-j))))
   115  	}
   116  	return b
   117  }
   118  
   119  // Len determines the number of bytes that will be needed to write the number i.
   120  func Len(i uint64) int {
   121  	if i <= maxVarInt1 {
   122  		return 1
   123  	}
   124  	if i <= maxVarInt2 {
   125  		return 2
   126  	}
   127  	if i <= maxVarInt4 {
   128  		return 4
   129  	}
   130  	if i <= maxVarInt8 {
   131  		return 8
   132  	}
   133  	// Don't use a fmt.Sprintf here to format the error message.
   134  	// The function would then exceed the inlining budget.
   135  	panic(struct {
   136  		message string
   137  		num     uint64
   138  	}{"value doesn't fit into 62 bits: ", i})
   139  }