packet.go

package main

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

type Protocol string

const (
	ProtocolTCP  Protocol = "tcp"
	ProtocolUDP  Protocol = "udp"
	ProtocolICMP Protocol = "icmp"
)

type ProtocolChooser struct {
	TCPBoundary float64
	UDPBoundary float64
}

func NewProtocolChooser(mix ProtocolMix) (ProtocolChooser, error) {
	normalized, err := mix.normalized()
	if err != nil {
		return ProtocolChooser{}, err
	}
	return ProtocolChooser{
		TCPBoundary: normalized.TCP,
		UDPBoundary: normalized.TCP + normalized.UDP,
	}, nil
}

func (c ProtocolChooser) Next(rng *rand.Rand) Protocol {
	value := rng.Float64()
	if value < c.TCPBoundary {
		return ProtocolTCP
	}
	if value < c.UDPBoundary {
		return ProtocolUDP
	}
	return ProtocolICMP
}

type PacketSizeSampler struct {
	Config PacketConfig
	RNG    *rand.Rand
}

func NewPacketSizeSampler(cfg PacketConfig, rng *rand.Rand) PacketSizeSampler {
	return PacketSizeSampler{Config: cfg, RNG: rng}
}

func (s *PacketSizeSampler) Next(minimum int) int {
	minimum = max(minimum, s.Config.MinSize)
	for attempts := 0; attempts < 16; attempts++ {
		value := int(float64(s.Config.MeanSize) + s.RNG.NormFloat64()*s.Config.StdDev)
		if value >= minimum && value <= s.Config.MaxSize {
			return value
		}
	}
	return min(max(s.Config.MeanSize, minimum), s.Config.MaxSize)
}

func minimumFrameLength(ipv6 bool, protocol Protocol) int {
	ipHeader := 20
	if ipv6 {
		ipHeader = 40
	}
	transport := 8
	if protocol == ProtocolTCP {
		transport = 20
	}
	return 14 + ipHeader + transport + 4
}

func buildPacketHeader(spec PacketSpec, rng *rand.Rand) ([]byte, error) {
	var source, destination netip.Addr
	if spec.Direction == DirectionIncoming {
		source, destination = spec.Peer, spec.Internal
	} else {
		source, destination = spec.Internal, spec.Peer
	}
	if source.Is6() != destination.Is6() {
		return nil, fmt.Errorf("source and destination address families differ")
	}

	frameLength := max(spec.FrameLength, minimumFrameLength(source.Is6(), spec.Protocol))
	sourceMAC := macFromAddress(source, 0x11)
	destinationMAC := macFromAddress(destination, 0x22)

	header := make([]byte, 14)
	copy(header[0:6], destinationMAC[:])
	copy(header[6:12], sourceMAC[:])
	if source.Is6() {
		binary.BigEndian.PutUint16(header[12:14], 0x86dd)
		ipHeader, err := buildIPv6Header(source, destination, spec, frameLength, rng)
		if err != nil {
			return nil, err
		}
		header = append(header, ipHeader...)
	} else {
		binary.BigEndian.PutUint16(header[12:14], 0x0800)
		ipHeader, err := buildIPv4Header(source, destination, spec, frameLength, rng)
		if err != nil {
			return nil, err
		}
		header = append(header, ipHeader...)
	}
	return header, nil
}

func buildIPv4Header(source, destination netip.Addr, spec PacketSpec, frameLength int, rng *rand.Rand) ([]byte, error) {
	protocolNumber, err := ipProtocolNumber(spec.Protocol, false)
	if err != nil {
		return nil, err
	}
	layer3Length := frameLength - 14 - 4
	if layer3Length > 65535 {
		return nil, fmt.Errorf("IPv4 packet length %d exceeds 65535", layer3Length)
	}

	header := make([]byte, 20)
	header[0] = 0x45
	binary.BigEndian.PutUint16(header[2:4], uint16(layer3Length))
	binary.BigEndian.PutUint16(header[4:6], uint16(rng.Uint32()))
	binary.BigEndian.PutUint16(header[6:8], 0x4000)
	header[8] = 64
	header[9] = protocolNumber
	sourceBytes := source.As4()
	destinationBytes := destination.As4()
	copy(header[12:16], sourceBytes[:])
	copy(header[16:20], destinationBytes[:])
	binary.BigEndian.PutUint16(header[10:12], checksum(header))

	transport, err := buildTransportHeader(spec, layer3Length-20, false, rng)
	if err != nil {
		return nil, err
	}
	return append(header, transport...), nil
}

func buildIPv6Header(source, destination netip.Addr, spec PacketSpec, frameLength int, rng *rand.Rand) ([]byte, error) {
	protocolNumber, err := ipProtocolNumber(spec.Protocol, true)
	if err != nil {
		return nil, err
	}
	layer3Length := frameLength - 14 - 4
	payloadLength := layer3Length - 40
	if payloadLength < 0 || payloadLength > 65535 {
		return nil, fmt.Errorf("invalid IPv6 payload length %d", payloadLength)
	}

	header := make([]byte, 40)
	flowLabel := rng.Uint32() & 0x000fffff
	binary.BigEndian.PutUint32(header[0:4], 0x60000000|flowLabel)
	binary.BigEndian.PutUint16(header[4:6], uint16(payloadLength))
	header[6] = protocolNumber
	header[7] = 64
	sourceBytes := source.As16()
	destinationBytes := destination.As16()
	copy(header[8:24], sourceBytes[:])
	copy(header[24:40], destinationBytes[:])

	transport, err := buildTransportHeader(spec, payloadLength, true, rng)
	if err != nil {
		return nil, err
	}
	return append(header, transport...), nil
}

func buildTransportHeader(spec PacketSpec, layer4Length int, ipv6 bool, rng *rand.Rand) ([]byte, error) {
	destinationPort := spec.DestinationPort
	if destinationPort == 0 {
		destinationPort = randomDestinationPort(spec.Protocol, rng)
	}
	sourcePort := uint16(1024 + rng.Intn(65535-1024))

	switch spec.Protocol {
	case ProtocolTCP:
		header := make([]byte, 20)
		binary.BigEndian.PutUint16(header[0:2], sourcePort)
		binary.BigEndian.PutUint16(header[2:4], destinationPort)
		binary.BigEndian.PutUint32(header[4:8], rng.Uint32())
		binary.BigEndian.PutUint32(header[8:12], rng.Uint32())
		header[12] = 5 << 4
		switch spec.TCPFlags {
		case "syn":
			header[13] = 0x02
		case "syn-ack":
			header[13] = 0x12
		default:
			header[13] = 0x10
		}
		binary.BigEndian.PutUint16(header[14:16], 65535)
		binary.BigEndian.PutUint16(header[16:18], 0xffff)
		return header, nil
	case ProtocolUDP:
		if layer4Length < 8 || layer4Length > 65535 {
			return nil, fmt.Errorf("invalid UDP length %d", layer4Length)
		}
		header := make([]byte, 8)
		binary.BigEndian.PutUint16(header[0:2], sourcePort)
		binary.BigEndian.PutUint16(header[2:4], destinationPort)
		binary.BigEndian.PutUint16(header[4:6], uint16(layer4Length))
		if ipv6 {
			binary.BigEndian.PutUint16(header[6:8], 0xffff)
		}
		return header, nil
	case ProtocolICMP:
		header := make([]byte, 8)
		if ipv6 {
			header[0] = 128 // ICMPv6 echo request
		} else {
			header[0] = 8 // ICMP echo request
		}
		binary.BigEndian.PutUint16(header[4:6], uint16(rng.Uint32()))
		binary.BigEndian.PutUint16(header[6:8], uint16(rng.Uint32()))
		binary.BigEndian.PutUint16(header[2:4], checksum(header))
		return header, nil
	default:
		return nil, fmt.Errorf("unsupported protocol %q", spec.Protocol)
	}
}

func ipProtocolNumber(protocol Protocol, ipv6 bool) (byte, error) {
	switch protocol {
	case ProtocolTCP:
		return 6, nil
	case ProtocolUDP:
		return 17, nil
	case ProtocolICMP:
		if ipv6 {
			return 58, nil
		}
		return 1, nil
	default:
		return 0, fmt.Errorf("unsupported protocol %q", protocol)
	}
}

func randomDestinationPort(protocol Protocol, rng *rand.Rand) uint16 {
	var common []uint16
	switch protocol {
	case ProtocolTCP:
		common = []uint16{443, 80, 22, 8443, 25, 3306}
	case ProtocolUDP:
		common = []uint16{53, 123, 443, 500, 4500, 5060}
	default:
		return 0
	}
	if rng.Float64() < 0.85 {
		return common[rng.Intn(len(common))]
	}
	return uint16(1024 + rng.Intn(65535-1024))
}

func checksum(data []byte) uint16 {
	var sum uint32
	for len(data) >= 2 {
		sum += uint32(binary.BigEndian.Uint16(data[:2]))
		data = data[2:]
	}
	if len(data) == 1 {
		sum += uint32(data[0]) << 8
	}
	for sum>>16 != 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return ^uint16(sum)
}