// SPDX-FileCopyrightText: 2026 Echolot contributors
// SPDX-License-Identifier: GPL-3.0-or-later

//go:build linux

package dataplane

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

// Fragment headers are the kind of thing that is either exactly right or silently useless: a
// wrong offset unit, a missing MF bit or a bad checksum produces packets that leave the machine
// and are dropped by the receiver's IP stack without a word. Nothing downstream would notice —
// the client would simply record "fragments do not get through", which is a wrong answer rather
// than a missing one. Hence these check the bytes.

func testAddrs() (netip.AddrPort, netip.AddrPort) {
	return netip.MustParseAddrPort("192.0.2.1:8442"), netip.MustParseAddrPort("198.51.100.9:41000")
}

func TestSplitCoversThePayloadExactlyOnce(t *testing.T) {
	src, dst := testAddrs()
	udp := buildUDP(src, dst, make([]byte, 2000))

	frags := splitIPv4(src.Addr(), dst.Addr(), udp, 576, 0x1234)
	if len(frags) < 3 {
		t.Fatalf("expected several fragments for %d bytes, got %d", len(udp), len(frags))
	}

	// Reassemble the way a receiver would: place each fragment's payload at its offset.
	rebuilt := make([]byte, len(udp))
	covered := make([]bool, len(udp))
	for _, f := range frags {
		flagsOff := binary.BigEndian.Uint16(f[6:8])
		off := int(flagsOff&0x1FFF) * 8
		body := f[20:]
		if off+len(body) > len(udp) {
			t.Fatalf("fragment at offset %d overruns the datagram", off)
		}
		for i, b := range body {
			if covered[off+i] {
				t.Fatalf("byte %d delivered twice", off+i)
			}
			covered[off+i] = true
			rebuilt[off+i] = b
		}
	}
	for i, c := range covered {
		if !c {
			t.Fatalf("byte %d was never sent", i)
		}
	}
	for i := range udp {
		if rebuilt[i] != udp[i] {
			t.Fatalf("reassembled byte %d differs", i)
		}
	}
}

func TestFragmentHeadersAreWellFormed(t *testing.T) {
	src, dst := testAddrs()
	udp := buildUDP(src, dst, make([]byte, 3000))
	frags := splitIPv4(src.Addr(), dst.Addr(), udp, 800, 0xBEEF)

	for i, f := range frags {
		if got := f[0]; got != 0x45 {
			t.Errorf("fragment %d: version/IHL = %#x, want 0x45", i, got)
		}
		if got := f[9]; got != 17 {
			t.Errorf("fragment %d: protocol = %d, want 17 (UDP)", i, got)
		}
		if got := binary.BigEndian.Uint16(f[4:6]); got != 0xBEEF {
			t.Errorf("fragment %d: IP ID = %#x — all fragments of one datagram must share it", i, got)
		}
		if got := binary.BigEndian.Uint16(f[2:4]); int(got) != len(f) {
			t.Errorf("fragment %d: total length = %d, actual %d", i, got, len(f))
		}
		flagsOff := binary.BigEndian.Uint16(f[6:8])
		mf := flagsOff&0x2000 != 0
		wantMF := i < len(frags)-1
		if mf != wantMF {
			t.Errorf("fragment %d: MF = %v, want %v", i, mf, wantMF)
		}
	}
}

// Offsets are counted in 8-byte units, so every fragment but the last must be a multiple of 8.
// A 100-byte "fragment size" that silently becomes 100 bytes on the wire produces a datagram no
// host will ever reassemble.
func TestNonFinalFragmentsAreEightByteMultiples(t *testing.T) {
	src, dst := testAddrs()
	udp := buildUDP(src, dst, make([]byte, 2500))
	for _, size := range []int{8, 100, 576, 999, 1400} {
		frags := splitIPv4(src.Addr(), dst.Addr(), udp, (size/8)*8, 1)
		for i, f := range frags[:len(frags)-1] {
			if body := len(f) - 20; body%8 != 0 {
				t.Errorf("size %d: non-final fragment %d carries %d bytes, not a multiple of 8",
					size, i, body)
			}
		}
	}
}

// The UDP checksum is optional in IPv4, and sending zero would be less code — but a
// zero-checksum datagram is dropped by some middleboxes, and that drop would be recorded as a
// fragmentation failure. So it must be present and correct.
func TestUDPChecksumVerifies(t *testing.T) {
	src, dst := testAddrs()
	for _, n := range []int{0, 1, 7, 8, 100, 1001} { // odd lengths exercise the tail-byte path
		udp := buildUDP(src, dst, make([]byte, n))
		if got := binary.BigEndian.Uint16(udp[6:8]); got == 0 {
			t.Fatalf("payload %d: checksum is zero, which means 'not computed'", n)
		}
		if sum := verifyUDPChecksum(src.Addr(), dst.Addr(), udp); sum != 0xFFFF {
			t.Errorf("payload %d: checksum does not verify (one's complement sum %#x)", n, sum)
		}
		if got := binary.BigEndian.Uint16(udp[4:6]); int(got) != len(udp) {
			t.Errorf("payload %d: UDP length field %d, actual %d", n, got, len(udp))
		}
	}
}

func TestUDPPortsComeFromTheSessionAddresses(t *testing.T) {
	src, dst := testAddrs()
	udp := buildUDP(src, dst, []byte("x"))
	if got := binary.BigEndian.Uint16(udp[0:2]); got != src.Port() {
		t.Errorf("source port = %d, want %d", got, src.Port())
	}
	// The destination port must be the client's observed source port, or the datagram arrives
	// at the machine and is discarded before any socket sees it.
	if got := binary.BigEndian.Uint16(udp[2:4]); got != dst.Port() {
		t.Errorf("destination port = %d, want %d", got, dst.Port())
	}
}

// Recomputes the one's complement sum over the pseudo-header and datagram; a correct checksum
// makes the total 0xFFFF.
func verifyUDPChecksum(src, dst netip.Addr, udp []byte) uint16 {
	var sum uint32
	s4, d4 := src.Unmap().As4(), dst.Unmap().As4()
	for _, b := range [][]byte{s4[:], d4[:]} {
		sum += uint32(binary.BigEndian.Uint16(b[0:2]))
		sum += uint32(binary.BigEndian.Uint16(b[2:4]))
	}
	sum += 17
	sum += uint32(len(udp))
	for i := 0; i+1 < len(udp); i += 2 {
		sum += uint32(binary.BigEndian.Uint16(udp[i : i+2]))
	}
	if len(udp)%2 == 1 {
		sum += uint32(udp[len(udp)-1]) << 8
	}
	for sum>>16 != 0 {
		sum = (sum & 0xFFFF) + (sum >> 16)
	}
	return uint16(sum)
}