vppn/peer/packets-util.go

package peer

import (
	"net/netip"
	"sync/atomic"
	"time"
	"unsafe"
)

var traceIDCounter uint64 = uint64(time.Now().Unix()<<30) + 1

func newTraceID() uint64 {
	return atomic.AddUint64(&traceIDCounter, 1)
}

// ----------------------------------------------------------------------------

type binWriter struct {
	b []byte
	i int
}

func newBinWriter(buf []byte) *binWriter {
	buf = buf[:cap(buf)]
	return &binWriter{buf, 0}
}

func (w *binWriter) Bool(b bool) *binWriter {
	if b {
		return w.Byte(1)
	}
	return w.Byte(0)
}

func (w *binWriter) Byte(b byte) *binWriter {
	w.b[w.i] = b
	w.i++
	return w
}

func (w *binWriter) SharedKey(key [32]byte) *binWriter {
	copy(w.b[w.i:w.i+32], key[:])
	w.i += 32
	return w
}

func (w *binWriter) Uint16(x uint16) *binWriter {
	*(*uint16)(unsafe.Pointer(&w.b[w.i])) = x
	w.i += 2
	return w
}

func (w *binWriter) Uint64(x uint64) *binWriter {
	*(*uint64)(unsafe.Pointer(&w.b[w.i])) = x
	w.i += 8
	return w
}

func (w *binWriter) Int64(x int64) *binWriter {
	*(*int64)(unsafe.Pointer(&w.b[w.i])) = x
	w.i += 8
	return w
}

func (w *binWriter) AddrPort(addrPort netip.AddrPort) *binWriter {
	w.Bool(addrPort.IsValid())
	addr := addrPort.Addr().As16()
	copy(w.b[w.i:w.i+16], addr[:])
	w.i += 16
	return w.Uint16(addrPort.Port())
}

func (w *binWriter) AddrPort8(l [8]netip.AddrPort) *binWriter {
	for _, addrPort := range l {
		w.AddrPort(addrPort)
	}
	return w
}

func (w *binWriter) Build() []byte {
	return w.b[:w.i]
}

// ----------------------------------------------------------------------------

type binReader struct {
	b   []byte
	i   int
	err error
}

func newBinReader(buf []byte) *binReader {
	return &binReader{b: buf}
}

func (r *binReader) hasBytes(n int) bool {
	if r.err != nil || (len(r.b)-r.i) < n {
		r.err = errMalformedPacket
		return false
	}
	return true
}

func (r *binReader) Bool(b *bool) *binReader {
	var bb byte
	r.Byte(&bb)
	*b = bb != 0
	return r
}

func (r *binReader) Byte(b *byte) *binReader {
	if !r.hasBytes(1) {
		return r
	}
	*b = r.b[r.i]
	r.i++
	return r
}

func (r *binReader) SharedKey(x *[32]byte) *binReader {
	if !r.hasBytes(32) {
		return r
	}
	*x = ([32]byte)(r.b[r.i : r.i+32])
	r.i += 32
	return r
}

func (r *binReader) Uint16(x *uint16) *binReader {
	if !r.hasBytes(2) {
		return r
	}
	*x = *(*uint16)(unsafe.Pointer(&r.b[r.i]))
	r.i += 2
	return r
}

func (r *binReader) Uint64(x *uint64) *binReader {
	if !r.hasBytes(8) {
		return r
	}
	*x = *(*uint64)(unsafe.Pointer(&r.b[r.i]))
	r.i += 8
	return r
}

func (r *binReader) Int64(x *int64) *binReader {
	if !r.hasBytes(8) {
		return r
	}
	*x = *(*int64)(unsafe.Pointer(&r.b[r.i]))
	r.i += 8
	return r
}

func (r *binReader) AddrPort(x *netip.AddrPort) *binReader {
	if !r.hasBytes(19) {
		return r
	}

	var (
		valid bool
		port  uint16
	)

	r.Bool(&valid)
	addr := netip.AddrFrom16(([16]byte)(r.b[r.i : r.i+16])).Unmap()
	r.i += 16

	r.Uint16(&port)

	if valid {
		*x = netip.AddrPortFrom(addr, port)
	} else {
		*x = netip.AddrPort{}
	}

	return r
}

func (r *binReader) AddrPort8(x *[8]netip.AddrPort) *binReader {
	for i := range x {
		r.AddrPort(&x[i])
	}
	return r
}

func (r *binReader) Error() error {
	return r.err
}