packet.go

package testsupport

import (
	"encoding/binary"
	"net"
	"testing"

	"github.com/gopacket/gopacket"
	"github.com/gopacket/gopacket/layers"
	"github.com/hetznercloud/virtio-go/virtio"
	"golang.org/x/sys/unix"
)

// PrependPacket encodes the given [virtio.NetHdr], puts it in front of the
// given packet slice and returns the new slice.
func PrependPacket(tb testing.TB, vnethdr virtio.NetHdr, pkt []byte) []byte {
	tb.Helper()

	result := make([]byte, virtio.NetHdrSize+len(pkt))
	if err := vnethdr.Encode(result); err != nil {
		tb.Fatalf("Encode vnethdr: %v", err)
	}
	copy(result[virtio.NetHdrSize:], pkt)

	return result
}

// TestPacket returns a new test packet which consists of Ethernet, IPv4 and UDP
// headers and a variable length payload. The given MAC address will be used as
// the source and destination MAC in the Ethernet header and the packet will
// have the given total length.
//
// Additionally, a [virtio.NetHdr] is returned that configures the necessary
// offloads that would be needed to transmit this packet over an MTU 1500 link.
// It configures checksum offloads and GSO, if needed.
func TestPacket(tb testing.TB, mac net.HardwareAddr, length int) (virtio.NetHdr, []byte) {
	tb.Helper()

	if length < 64 {
		tb.Fatal("Packets should be at least 64 bytes long.")
	}

	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}

	eth := layers.Ethernet{
		SrcMAC:       mac,
		DstMAC:       mac,
		EthernetType: layers.EthernetTypeIPv4,
	}
	ipv4 := layers.IPv4{
		Version:  4,
		TTL:      64,
		Protocol: layers.IPProtocolUDP,
		SrcIP:    net.IPv4(127, 0, 0, 1),
		DstIP:    net.IPv4(127, 0, 0, 1),
	}
	udp := layers.UDP{
		SrcPort: 1234,
		DstPort: 5678,
	}
	if err := udp.SetNetworkLayerForChecksum(&ipv4); err != nil {
		tb.Fatalf("Set network layer for UDP header: %v", err)
	}

	buf := gopacket.NewSerializeBuffer()
	if err := gopacket.SerializeLayers(
		buf,
		opts,
		&eth, &ipv4, &udp, Payload(tb, length-42),
	); err != nil {
		tb.Fatalf("Serialize packet: %v", err)
	}

	packet := buf.Bytes()
	if len(packet) != length {
		tb.Fatal("Unexpected packet length mismatch.")
	}

	// Clear the UDP checksum to make the offload work.
	// When the offloads are not used, the UDP checksum will just remain zero,
	// which is also okay, because it is optional for IPv4 UDP.
	binary.BigEndian.PutUint16(packet[14+20+6:], uint16(0))

	vnethdr := virtio.NetHdr{
		// The packet was already serialized with valid L4 checksums, but we
		// still configure the checksum offload here for testing purposes.
		// Also, it's needed for GSO to work.
		Flags:      unix.VIRTIO_NET_HDR_F_NEEDS_CSUM,
		CsumStart:  14 + 20, // Checksum everything after the IP header.
		CsumOffset: 6,       // Insert the checksum into the UDP header.
	}

	// In case the packet would not fit a MTU 1500 link, also configure GSO.
	if length > 1514 {
		vnethdr.GSOType = unix.VIRTIO_NET_HDR_GSO_UDP_L4
		vnethdr.GSOSize = 1500 - 20 - 8 // Segmented payload size.
		vnethdr.HdrLen = 14 + 20 + 8    // Length of the headers to be replicated.
	}

	return vnethdr, packet
}

// Payload returns a new [gopacket.Payload] with the given length which can be
// used for test packets.
func Payload(tb testing.TB, length int) gopacket.Payload {
	tb.Helper()

	str := "HELLO_VIRTIO_"
	payload := make(gopacket.Payload, length)
	for i := range length {
		payload[i] = str[i%len(str)]
	}

	return payload
}