varint.go

package varint

import (
	"errors"
	"fmt"
	"io"
	"math"
)

const (
	// max bytes that can be received for one integer
	MaxVarintLen32 = 5
	MaxVarintLen64 = 10
)

// PutVarint encodes an int32 into buf and returns the number of bytes written.
// If the buffer is too small, PutVarint will panic.
// Try to allocate a buffer of MaxVarintLen32 bytes to avoid panics.
// Format: ESDDDDDD EDDDDDDD EDD... Extended, Sign, Data
// E: is next byte part of the current integer
// S: Sign of integer
// Data, Integer bits that follow the sign
func PutVarint(buf []byte, x int) int {
	if x < math.MinInt32 || math.MaxInt32 < x {
		panic("ERROR: value to Pack is out of bounds, should be within range [-2147483648:2147483647] (32bit)")
	}

	const intSize = 4

	// stack allocated buffer
	data := [MaxVarintLen32]byte{} // predefined content of zeroes
	index := 0

	data[index] = byte(x>>(intSize*8-7)) & 0b01000000 // set sign bit if i<0
	x = x ^ (x >> (intSize*8 - 1))                    // if(i<0) i = ~i

	data[index] |= byte(x) & 0b00111111 // pack 6bit into data
	x >>= 6                             // discard 6 bits

	if x != 0 {
		data[index] |= 0b10000000 // set extend bit

		for {
			index++
			data[index] = byte(x) & 0b01111111 //  pack 7 bits
			x >>= 7                            // discard 7 bits

			if x != 0 {
				data[index] |= 1 << 7 // set extend bit
			} else {
				break // break if x is 0
			}

		}
	}

	size := index + 1

	if len(buf) < size {
		panic(fmt.Sprintf("varint buffer needs to have at least %d bytes but has %d", len(data), len(buf)))
	}

	return copy(buf, data[:size])
}

// Varint decodes an int from buf and returns that value and the number of bytes read (> 0).
// If an error occurred, the value is 0 and the number of bytes n is <= 0 with the following meaning:
//
//	n == 0: buf too small
//	n  < 0: value larger than 32 bits (overflow)
//	        and -n is the number of bytes read
func Varint(buf []byte) (i int, n int) {

	if len(buf) == 0 {
		return 0, 0
	}

	index := 0
	// handle first byte (most right side)
	sign := int((buf[index] >> 6) & 0b00000001)
	value := int(buf[index] & 0b00111111)

	// no E bit set, return after parsing first byte
	if buf[index] < 0b10000000 {
		value ^= -sign // if(sign) value = ~(value)
		index++
		return value, index
	}

	// handle 2nd - nth byte
	buf = buf[1:]
	const maxAllowedLen = MaxVarintLen32 - 1
	for i, b := range buf {
		index++
		// overflow check
		// 1 sign bit + 6 data bits = 7 bits
		// 7 bits * 4 bytes = 28 bits
		// 7 + 28 = 35 bits, 3 too many
		// last byte can only have 4 bits
		if i == maxAllowedLen-1 && b > 0b00001111 {
			return 0, -(i + 1)
		}

		value |= int(b&0b01111111) << (6 + 7*i)
		if b < 0b10000000 {
			// no extend bit set
			break
		}
	}

	value ^= -sign // if(sign) value = ~(value)
	index++

	return value, index
}

// AppendVarint appends the varint-encoded form of x, as generated by PutVarint, to buf and returns the extended buffer.
func AppendVarint(buf []byte, x int) []byte {
	arr := [MaxVarintLen32]byte{}
	sbuf := arr[:]
	n := PutVarint(sbuf, x)
	sbuf = sbuf[:n]
	return append(buf, sbuf...)
}

// ReadVarint can decode a stream of bytes
func ReadVarint(r io.ByteReader) (int, error) {
	b, err := r.ReadByte()
	if err != nil {
		return 0, err
	}

	index := 0
	// handle first byte (most right side)
	sign := int((b >> 6) & 0b00000001)
	value := int(b & 0b00111111)

	// no E bit set, return after parsing first byte
	if b < 0b10000000 {
		value ^= -sign // if(sign) value = ~(value)
		index++
		return value, nil
	}

	// handle 2nd - nth byte
	const maxAllowedLen = MaxVarintLen32 - 1
	for i := 0; i < MaxVarintLen32; i++ {
		b, err := r.ReadByte()
		if err != nil {
			if errors.Is(err, io.EOF) {
				return value, io.ErrUnexpectedEOF
			}
			return value, nil
		}
		index++

		// overflow check
		// 1 sign bit + 6 data bits = 7 bits
		// 7 bits * 4 bytes = 28 bits
		// 7 + 28 = 35 bits, 3 too many
		// 7 - 3 = 7 - 3 = last byte can only have 4 bits
		if i == maxAllowedLen-1 && b > 0b00001111 {
			return 0, errors.New("overflow due to invalid last byte")
		}

		value |= int(b&0b01111111) << (6 + 7*i)
		if b < 0b10000000 {
			break
		}

	}

	value ^= -sign // if(sign) value = ~(value)
	index++

	return value, nil
}

// PutBigVarint is a variant of varint that can encode 64-bit integers.
// It is out of the scope of the Teeworlds protocol but can be used to check for potential security issues.
func PutBigVarint(buf []byte, x int64) int {
	const intSize = 8

	// stack allocated buffer
	data := [MaxVarintLen64]byte{} // predefined content of zeroes
	index := 0

	data[index] = byte(x>>(intSize*8-7)) & 0b01000000 // set sign bit if i<0
	x = x ^ (x >> (intSize*8 - 1))                    // if(i<0) i = ~i

	data[index] |= byte(x) & 0b00111111 // pack 6bit into data
	x >>= 6                             // discard 6 bits

	if x != 0 {
		data[index] |= 0b10000000 // set extend bit

		for {
			index++
			data[index] = byte(x) & 0b01111111 //  pack 7 bits
			x >>= 7                            // discard 7 bits

			if x != 0 {
				data[index] |= 1 << 7 // set extend bit
			} else {
				break // break if x is 0
			}
		}
	}

	size := index + 1

	if len(buf) < size {
		panic(fmt.Sprintf("varint buffer needs to have at least %d bytes but has %d", len(data), len(buf)))
	}

	return copy(buf, data[:size])
}

// Big decodes an int from buf and returns that value and the number of bytes read (> 0).
// If an error occurred, the value is 0 and the number of bytes n is <= 0 with the following meaning:
//
//	n == 0: buf too small
//	n  < 0: value larger than 64 bits (overflow/underflow)
//	        and -n is the number of bytes read
func BigVarint(buf []byte) (i int64, n int) {

	if len(buf) == 0 {
		return 0, 0
	}

	index := 0
	// handle first byte (most right side)
	sign := int64((buf[index] >> 6) & 0b00000001)
	value := int64(buf[index] & 0b00111111)

	// no E bit set, return after parsing first byte
	if buf[index] < 0b10000000 {
		value ^= -sign // if(sign) value = ~(value)
		index++
		return value, index
	}

	// handle 2nd - nth byte
	buf = buf[1:]
	const maxAllowedLen = MaxVarintLen64 - 1
	for i, b := range buf {
		index++
		// overflow check
		// 1 sign bit + 6 data bits = 7 bits
		// 7 bits * 9 bytes = 63 bits
		// 7 + 63 = 70 bits, 6 too many
		// 7 - 6 = last byte can only have 1 bit
		if i == maxAllowedLen-1 && b > 0b00000001 {
			return 0, -(i + 1)
		}

		value |= int64(b&0b01111111) << (6 + 7*i)
		if b < 0b10000000 {
			// no extend bit set
			break
		}
	}

	value ^= -sign // if(sign) value = ~(value)
	index++

	return value, index
}

// AppendBigVarint appends the varint-encoded form of x, as generated by PutBigVarint, to buf and returns the extended buffer.
// It is out of the scope of the Teeworlds protocol but can be used to check for potential security issues.
func AppendBigVarint(buf []byte, x int64) []byte {
	arr := [MaxVarintLen64]byte{}
	sbuf := arr[:]
	n := PutBigVarint(sbuf, x)
	sbuf = sbuf[:n]
	return append(buf, sbuf...)
}

// ReadVarint can decode a stream of bytes
func ReadBigVarint(r io.ByteReader) (int64, error) {
	b, err := r.ReadByte()
	if err != nil {
		return 0, err
	}

	index := 0
	// handle first byte (most right side)
	sign := int64((b >> 6) & 0b00000001)
	value := int64(b & 0b00111111)

	// no E bit set, return after parsing first byte
	if b < 0b10000000 {
		value ^= -sign // if(sign) value = ~(value)
		index++
		return value, nil
	}

	// handle 2nd - nth byte
	const maxAllowedLen = MaxVarintLen64 - 1
	for i := 0; i < MaxVarintLen64; i++ {
		b, err := r.ReadByte()
		if err != nil {
			if errors.Is(err, io.EOF) {
				return value, io.ErrUnexpectedEOF
			}
			return value, nil
		}
		index++

		// overflow check
		// 1 sign bit + 6 data bits = 7 bits
		// 7 bits * 4 bytes = 28 bits
		// 7 + 28 = 35 bits, 3 too many
		// 7 - 3 = 7 - 3 = last byte can only have 4 bits
		if i == maxAllowedLen-1 && b > 0b00000001 {
			return 0, errors.New("overflow due to invalid last byte")
		}

		value |= int64(b&0b01111111) << (6 + 7*i)
		if b < 0b10000000 {
			break
		}

	}

	value ^= -sign // if(sign) value = ~(value)
	index++

	return value, nil
}