lux/proto/lux_buffer.go

package proto

import (
	"encoding/binary"
	"fmt"
	"net"
	"net/netip"
)

var NO = binary.BigEndian

const MIN_BUFFER_SIZE = 256

type LuxBuffer struct {
	data   []byte
	offset int
	len    int
}

func AllocLuxBuffer(cap int) LuxBuffer {
	return LuxBuffer{
		data:   make([]byte, cap),
		offset: 0,
		len:    0,
	}
}

func NewLuxBuffer() LuxBuffer {
	return AllocLuxBuffer(MIN_BUFFER_SIZE)
}

func FromSlice(bytes []byte) LuxBuffer {
	return LuxBuffer{
		data:   bytes,
		offset: 0,
		len:    len(bytes),
	}
}

func (buf *LuxBuffer) Capacity() int {
	return len(buf.data)
}

func (buf *LuxBuffer) Offset() int {
	return buf.offset
}

func (buf *LuxBuffer) Length() int {
	return buf.len
}

// available capacity to write (before growing)
func (buf *LuxBuffer) Available() int {
	return buf.Capacity() - buf.len
}

// available length to read
func (buf *LuxBuffer) Remaining() int {
	return buf.len - buf.offset
}

func (buf *LuxBuffer) Grow(grow int) {
	ocap, ncap := buf.Capacity(), (buf.Capacity()*2 + grow)
	odata := buf.data

	buf.data = make([]byte, ncap)
	copy(buf.data[:ocap], odata[:])
}

// will return byte buffer, sliced to real length
func (buf *LuxBuffer) AllBytes() []byte {
	return buf.data[:buf.len]
}

// ensure capacity for new write, return slice pointing to a place of a new write
func (buf *LuxBuffer) WriteNext(size int) []byte {
	if buf.Available() < size {
		buf.Grow(size)
	}

	next := buf.data[buf.len : buf.len+size]
	buf.len += size

	return next
}

func (buf *LuxBuffer) WriteBytes(bytes []byte) {
	copy(buf.WriteNext(len(bytes)), bytes)
}

func (buf *LuxBuffer) ReadNext(size int) ([]byte, error) {
	if buf.offset+size > buf.len {
		return nil, fmt.Errorf("ReadNext %d+%d > %d", buf.offset, size, buf.len)
	}
	next := buf.data[buf.offset : buf.offset+size]
	buf.offset += size

	return next, nil
}

func (buf *LuxBuffer) Skip(skip int) {
	buf.offset += skip
}

// explicit copy of read bytes
func (buf *LuxBuffer) CopyBytes(size int) ([]byte, error) {
	bytes, err := buf.ReadNext(size)
	if err != nil {
		return nil, err
	}

	read := make([]byte, size)
	copy(read, bytes)

	return read, nil
}

func (buf *LuxBuffer) ReadUint16() (uint16, error) {
	rd, err := buf.ReadNext(2)
	if err != nil {
		return 0, err
	}
	return NO.Uint16(rd), nil
}

func (buf *LuxBuffer) ReadUint32() (uint32, error) {
	rd, err := buf.ReadNext(4)
	if err != nil {
		return 0, err
	}
	return NO.Uint32(rd), nil
}

func (buf *LuxBuffer) ReadUint64() (uint64, error) {
	rd, err := buf.ReadNext(8)
	if err != nil {
		return 0, err
	}
	return NO.Uint64(rd), nil
}

func (buf *LuxBuffer) WriteUint16(val uint16) {
	NO.PutUint16(buf.WriteNext(2), val)
}

func (buf *LuxBuffer) WriteUint32(val uint32) {
	NO.PutUint32(buf.WriteNext(4), val)
}

func (buf *LuxBuffer) WriteUint64(val uint64) {
	NO.PutUint64(buf.WriteNext(8), val)
}

// variable-length blocks can cause misaligned size,
// so we make method for them, forcing padding

func (buf *LuxBuffer) ReadVarBlock() ([]byte, error) {
	len, err := buf.ReadUint16()
	if err != nil {
		return nil, err
	}
	bytes, err := buf.ReadNext(int(len))
	if err != nil {
		return nil, err
	}

	if len%2 != 0 {
		buf.Skip(1) // skip padding
	}

	return bytes, nil
}

func (buf *LuxBuffer) WriteVarBlock(bytes []byte) {
	len := uint16(len(bytes))

	buf.WriteUint16(len)
	buf.WriteBytes(bytes)
	if len%2 != 0 {
		buf.WriteNext(1) // padding
	}
}

// strings in LUX protocol format will be padded to align of 2
func (buf *LuxBuffer) ReadString() (string, error) {
	block, err := buf.ReadVarBlock()
	if err != nil {
		return "", err
	}
	return string(block), nil
}

func (buf *LuxBuffer) WriteString(val string) {
	buf.WriteVarBlock([]byte(val))
}

func (buf *LuxBuffer) ReadIP() (netip.Addr, error) {
	var addr netip.Addr

	ipOctets, err := buf.ReadVarBlock()
	if err != nil {
		return addr, err
	}

	switch len(ipOctets) {
	case net.IPv4len:
		addr = netip.AddrFrom4([4]byte(ipOctets))
	case net.IPv6len:
		addr = netip.AddrFrom16([16]byte(ipOctets))
	default:
		return addr, fmt.Errorf("wrong ip size %d", len(ipOctets))
	}

	return addr, nil
}

func (buf *LuxBuffer) WriteIP(addr netip.Addr) {
	if addr.Is4() {
		octets := addr.As4()
		buf.WriteVarBlock(octets[:])
	} else if addr.Is6() {
		octets := addr.As16()
		buf.WriteVarBlock(octets[:])
	} else {
		panic("WriteIP addr is neither IPv4 nor IPv6")
	}
}

func (buf *LuxBuffer) ReadIPPort() (netip.AddrPort, error) {
	var addrPort netip.AddrPort

	addr, err := buf.ReadIP()
	if err != nil {
		return addrPort, err
	}

	portNum, err := buf.ReadUint16()
	if err != nil {
		return addrPort, err
	}

	return netip.AddrPortFrom(addr, portNum), nil
}

func (buf *LuxBuffer) WriteIPPort(addrPort netip.AddrPort) {
	buf.WriteIP(addrPort.Addr())
	buf.WriteUint16(addrPort.Port())
}

func LuxProtoIPToAddr(ip net.IP) netip.Addr {
	if octets4 := ip.To4(); octets4 != nil {
		return netip.AddrFrom4([4]byte(octets4))
	} else {
		return netip.AddrFrom16([16]byte(ip.To16()))
	}
}

func LuxProtoAddrToIP(addr netip.Addr) net.IP {
	if addr.Is4() {
		octets4 := addr.As4()
		return octets4[:]
	} else {
		octets16 := addr.As16()
		return octets16[:]
	}
}