DebTik/internal/bwtest/protocol.go

package bwtest

import (
	"context"
	"encoding/binary"
	"errors"
	"hash/fnv"
	"math"
	"time"
)

const (
	MagicBytes uint32 = 0x42575453 // "BWTS"
	ProtoVer   uint8  = 1

	PacketTypeData   uint8 = 1
	PacketTypeReport uint8 = 2
	PacketTypeStop   uint8 = 3

	HeaderSize = 36
)

var (
	ErrInvalidMagic = errors.New("invalid bandwidth test packet magic")
	ErrShortPacket  = errors.New("packet shorter than minimum header size")
)

// PacketHeader represents the 36-byte binary header on every UDP test datagram
type PacketHeader struct {
	Magic         uint32
	Version       uint8
	PacketType    uint8
	StreamID      uint8
	Flags         uint8
	SessionHash   uint64
	Sequence      uint64
	TimestampNano int64
	PayloadLength uint32
}

// HashSessionID generates a fast 64-bit FNV-1a hash of the session ID string
func HashSessionID(sessionID string) uint64 {
	h := fnv.New64a()
	_, _ = h.Write([]byte(sessionID))
	return h.Sum64()
}

// EncodeHeader serializes the packet header into the provided byte slice (must be >= 36 bytes)
func EncodeHeader(b []byte, hdr PacketHeader) {
	binary.BigEndian.PutUint32(b[0:4], MagicBytes)
	b[4] = ProtoVer
	b[5] = hdr.PacketType
	b[6] = hdr.StreamID
	b[7] = hdr.Flags
	binary.BigEndian.PutUint64(b[8:16], hdr.SessionHash)
	binary.BigEndian.PutUint64(b[16:24], hdr.Sequence)
	binary.BigEndian.PutUint64(b[24:32], uint64(hdr.TimestampNano))
	binary.BigEndian.PutUint32(b[32:36], hdr.PayloadLength)
}

// DecodeHeader deserializes the binary packet header from byte slice
func DecodeHeader(b []byte) (PacketHeader, error) {
	if len(b) < HeaderSize {
		return PacketHeader{}, ErrShortPacket
	}
	magic := binary.BigEndian.Uint32(b[0:4])
	if magic != MagicBytes {
		return PacketHeader{}, ErrInvalidMagic
	}
	return PacketHeader{
		Magic:         magic,
		Version:       b[4],
		PacketType:    b[5],
		StreamID:      b[6],
		Flags:         b[7],
		SessionHash:   binary.BigEndian.Uint64(b[8:16]),
		Sequence:      binary.BigEndian.Uint64(b[16:24]),
		TimestampNano: int64(binary.BigEndian.Uint64(b[24:32])),
		PayloadLength: binary.BigEndian.Uint32(b[32:36]),
	}, nil
}

// JitterCalculator computes RFC 3550 statistical interarrival jitter
type JitterCalculator struct {
	hasPrev     bool
	prevTransit int64
	jitterFloat float64 // in nanoseconds
}

// Update incorporates a newly received packet's sender timestamp and arrival time
func (jc *JitterCalculator) Update(senderTsNano, arrivalTsNano int64) float64 {
	transit := arrivalTsNano - senderTsNano
	if !jc.hasPrev {
		jc.prevTransit = transit
		jc.hasPrev = true
		return 0
	}
	diff := math.Abs(float64(transit - jc.prevTransit))
	jc.prevTransit = transit
	// RFC 3550: J(i) = J(i-1) + (|D(i-1,i)| - J(i-1))/16
	jc.jitterFloat += (diff - jc.jitterFloat) / 16.0
	return jc.jitterFloat / 1e6 // Return in milliseconds
}

// RateLimiter implements a high-precision token bucket for pacing outbound test traffic
type RateLimiter struct {
	bps           int64
	tokens        float64
	maxTokens     float64
	lastCheckTime time.Time
}

// NewRateLimiter initializes a rate limiter for target bits per second (0 = no limit)
func NewRateLimiter(targetBps int64) *RateLimiter {
	if targetBps <= 0 {
		return nil
	}
	bytesPerSec := float64(targetBps) / 8.0
	return &RateLimiter{
		bps:           targetBps,
		tokens:        bytesPerSec * 0.05, // 50ms burst allowance
		maxTokens:     bytesPerSec * 0.10, // 100ms max burst
		lastCheckTime: time.Now(),
	}
}

// Wait preserves the original API; cancellable workers use WaitContext.
func (rl *RateLimiter) Wait(bytes int) { rl.WaitContext(context.Background(), bytes) }

// A frame can exceed the nominal 100ms bucket capacity. Grow that bucket to at
// least one frame; otherwise low-rate TCP tests wait forever for unreachable tokens.
func (rl *RateLimiter) WaitContext(ctx context.Context, bytes int) bool {
	if rl == nil || rl.bps <= 0 {
		return ctx.Err() == nil
	}
	needed := float64(bytes)
	if rl.maxTokens < needed {
		rl.maxTokens = needed
	}
	rate := float64(rl.bps) / 8
	for {
		if ctx.Err() != nil {
			return false
		}
		now := time.Now()
		rl.tokens += now.Sub(rl.lastCheckTime).Seconds() * rate
		rl.lastCheckTime = now
		if rl.tokens > rl.maxTokens {
			rl.tokens = rl.maxTokens
		}
		if rl.tokens >= needed {
			rl.tokens -= needed
			return true
		}
		wait := time.Duration((needed - rl.tokens) / rate * float64(time.Second))
		if wait > 50*time.Millisecond {
			wait = 50 * time.Millisecond
		}
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return false
		case <-timer.C:
		}
	}
}