package nebula

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

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"

	"github.com/slackhq/nebula/firewall"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/ipv4"
)

func Test_newPacket(t *testing.T) {
	p := &firewall.ParsedPacket{}

	// length fails
	err := newPacket([]byte{}, true, p)
	require.ErrorIs(t, err, ErrPacketTooShort)

	err = newPacket([]byte{0x40}, true, p)
	require.ErrorIs(t, err, ErrIPv4PacketTooShort)

	err = newPacket([]byte{0x60}, true, p)
	require.ErrorIs(t, err, ErrIPv6PacketTooShort)

	// length fail with ip options
	h := ipv4.Header{
		Version: 1,
		Len:     100,
		Src:     net.IPv4(10, 0, 0, 1),
		Dst:     net.IPv4(10, 0, 0, 2),
		Options: []byte{0, 1, 0, 2},
	}

	b, _ := h.Marshal()
	err = newPacket(b, true, p)
	require.ErrorIs(t, err, ErrIPv4InvalidHeaderLength)

	// not an ipv4 packet
	err = newPacket([]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, true, p)
	require.ErrorIs(t, err, ErrUnknownIPVersion)

	// invalid ihl
	err = newPacket([]byte{4<<4 | (8 >> 2 & 0x0f), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, true, p)
	require.ErrorIs(t, err, ErrIPv4InvalidHeaderLength)

	// account for variable ip header length - incoming
	h = ipv4.Header{
		Version:  1,
		Len:      100,
		Src:      net.IPv4(10, 0, 0, 1),
		Dst:      net.IPv4(10, 0, 0, 2),
		Options:  []byte{0, 1, 0, 2},
		Protocol: firewall.ProtoTCP,
	}

	b, _ = h.Marshal()
	b = append(b, []byte{0, 3, 0, 4}...)
	err = newPacket(b, true, p)

	require.NoError(t, err)
	assert.Equal(t, uint8(firewall.ProtoTCP), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("10.0.0.2"), p.LocalAddr)
	assert.Equal(t, netip.MustParseAddr("10.0.0.1"), p.RemoteAddr)
	assert.Equal(t, uint16(3), p.RemotePort)
	assert.Equal(t, uint16(4), p.LocalPort)
	assert.False(t, p.Fragment)

	// account for variable ip header length - outgoing
	h = ipv4.Header{
		Version:  1,
		Protocol: 2,
		Len:      100,
		Src:      net.IPv4(10, 0, 0, 1),
		Dst:      net.IPv4(10, 0, 0, 2),
		Options:  []byte{0, 1, 0, 2},
	}

	b, _ = h.Marshal()
	b = append(b, []byte{0, 5, 0, 6}...)
	err = newPacket(b, false, p)

	require.NoError(t, err)
	assert.Equal(t, uint8(2), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("10.0.0.1"), p.LocalAddr)
	assert.Equal(t, netip.MustParseAddr("10.0.0.2"), p.RemoteAddr)
	assert.Equal(t, uint16(6), p.RemotePort)
	assert.Equal(t, uint16(5), p.LocalPort)
	assert.False(t, p.Fragment)
}

func Test_newPacket_v6(t *testing.T) {
	p := &firewall.ParsedPacket{}

	// invalid ipv6
	ip := layers.IPv6{
		Version:  6,
		HopLimit: 128,
		SrcIP:    net.IPv6linklocalallrouters,
		DstIP:    net.IPv6linklocalallnodes,
	}

	buffer := gopacket.NewSerializeBuffer()
	opt := gopacket.SerializeOptions{
		ComputeChecksums: false,
		FixLengths:       false,
	}
	err := gopacket.SerializeLayers(buffer, opt, &ip)
	require.NoError(t, err)

	err = newPacket(buffer.Bytes(), true, p)
	require.ErrorIs(t, err, ErrIPv6CouldNotFindPayload)

	// A v6 packet with a hop-by-hop extension
	// ICMPv6 Payload (Echo Request)
	icmpLayer := layers.ICMPv6{
		TypeCode: layers.ICMPv6TypeEchoRequest,
	}
	// Hop-by-Hop Extension Header
	hopOption := layers.IPv6HopByHopOption{}
	hopOption.OptionData = []byte{0, 0, 0, 0}
	hopByHop := layers.IPv6HopByHop{}
	hopByHop.Options = append(hopByHop.Options, &hopOption)

	ip = layers.IPv6{
		Version:    6,
		HopLimit:   128,
		NextHeader: layers.IPProtocolIPv6Destination,
		SrcIP:      net.IPv6linklocalallrouters,
		DstIP:      net.IPv6linklocalallnodes,
	}

	buffer.Clear()
	err = gopacket.SerializeLayers(buffer, gopacket.SerializeOptions{
		ComputeChecksums: false,
		FixLengths:       true,
	}, &ip, &hopByHop, &icmpLayer)
	if err != nil {
		panic(err)
	}
	// Ensure buffer length checks during parsing with the next 2 tests.

	// A full IPv6 header and 1 byte in the first extension, but missing
	// the length byte.
	err = newPacket(buffer.Bytes()[:41], true, p)
	require.ErrorIs(t, err, ErrIPv6CouldNotFindPayload)

	// A full IPv6 header plus 1 full extension, but only 1 byte of the
	// next layer, missing length byte
	err = newPacket(buffer.Bytes()[:49], true, p)
	require.ErrorIs(t, err, ErrIPv6CouldNotFindPayload)
	err = nil

	// A good ICMP packet
	ip = layers.IPv6{
		Version:    6,
		NextHeader: layers.IPProtocolICMPv6,
		HopLimit:   128,
		SrcIP:      net.IPv6linklocalallrouters,
		DstIP:      net.IPv6linklocalallnodes,
	}

	icmp := layers.ICMPv6{
		TypeCode: layers.ICMPv6TypeEchoRequest,
		Checksum: 0x1234,
	}

	buffer.Clear()
	require.NoError(t, gopacket.SerializeLayers(buffer, opt, &ip, &icmp))
	require.Error(t, newPacket(buffer.Bytes(), true, p))

	buffer.Clear()
	echo := layers.ICMPv6Echo{
		Identifier: 0xabcd,
		SeqNumber:  1234,
	}
	require.NoError(t, gopacket.SerializeLayers(buffer, opt, &ip, &icmp, &echo))
	require.NoError(t, newPacket(buffer.Bytes(), true, p))
	assert.Equal(t, uint8(layers.IPProtocolICMPv6), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.RemoteAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.LocalAddr)
	assert.Equal(t, uint16(0xabcd), p.RemotePort)
	assert.Equal(t, uint16(0), p.LocalPort)
	assert.False(t, p.Fragment)

	// A good ESP packet
	b := buffer.Bytes()
	b[6] = byte(layers.IPProtocolESP)
	err = newPacket(b, true, p)
	require.NoError(t, err)
	assert.Equal(t, uint8(layers.IPProtocolESP), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.RemoteAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.LocalAddr)
	assert.Equal(t, uint16(0), p.RemotePort)
	assert.Equal(t, uint16(0), p.LocalPort)
	assert.False(t, p.Fragment)

	// A good None packet
	b = buffer.Bytes()
	b[6] = byte(layers.IPProtocolNoNextHeader)
	err = newPacket(b, true, p)
	require.NoError(t, err)
	assert.Equal(t, uint8(layers.IPProtocolNoNextHeader), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.RemoteAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.LocalAddr)
	assert.Equal(t, uint16(0), p.RemotePort)
	assert.Equal(t, uint16(0), p.LocalPort)
	assert.False(t, p.Fragment)

	// An unknown protocol packet
	b = buffer.Bytes()
	b[6] = 255 // 255 is a reserved protocol number
	err = newPacket(b, true, p)
	require.ErrorIs(t, err, ErrIPv6CouldNotFindPayload)

	// A good UDP packet
	ip = layers.IPv6{
		Version:    6,
		NextHeader: firewall.ProtoUDP,
		HopLimit:   128,
		SrcIP:      net.IPv6linklocalallrouters,
		DstIP:      net.IPv6linklocalallnodes,
	}

	udp := layers.UDP{
		SrcPort: layers.UDPPort(36123),
		DstPort: layers.UDPPort(22),
	}
	err = udp.SetNetworkLayerForChecksum(&ip)
	require.NoError(t, err)

	buffer.Clear()
	err = gopacket.SerializeLayers(buffer, opt, &ip, &udp, gopacket.Payload([]byte{0xde, 0xad, 0xbe, 0xef}))
	if err != nil {
		panic(err)
	}
	b = buffer.Bytes()

	// incoming
	err = newPacket(b, true, p)
	require.NoError(t, err)
	assert.Equal(t, uint8(firewall.ProtoUDP), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.RemoteAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.LocalAddr)
	assert.Equal(t, uint16(36123), p.RemotePort)
	assert.Equal(t, uint16(22), p.LocalPort)
	assert.False(t, p.Fragment)

	// outgoing
	err = newPacket(b, false, p)
	require.NoError(t, err)
	assert.Equal(t, uint8(firewall.ProtoUDP), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.LocalAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.RemoteAddr)
	assert.Equal(t, uint16(36123), p.LocalPort)
	assert.Equal(t, uint16(22), p.RemotePort)
	assert.False(t, p.Fragment)

	// Too short UDP packet
	err = newPacket(b[:len(b)-10], false, p) // pull off the last 10 bytes
	require.ErrorIs(t, err, ErrIPv6PacketTooShort)

	// A good TCP packet
	b[6] = byte(layers.IPProtocolTCP)

	// incoming
	err = newPacket(b, true, p)
	require.NoError(t, err)
	assert.Equal(t, uint8(firewall.ProtoTCP), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.RemoteAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.LocalAddr)
	assert.Equal(t, uint16(36123), p.RemotePort)
	assert.Equal(t, uint16(22), p.LocalPort)
	assert.False(t, p.Fragment)

	// outgoing
	err = newPacket(b, false, p)
	require.NoError(t, err)
	assert.Equal(t, uint8(firewall.ProtoTCP), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.LocalAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.RemoteAddr)
	assert.Equal(t, uint16(36123), p.LocalPort)
	assert.Equal(t, uint16(22), p.RemotePort)
	assert.False(t, p.Fragment)

	// Too short TCP packet
	err = newPacket(b[:len(b)-10], false, p) // pull off the last 10 bytes
	require.ErrorIs(t, err, ErrIPv6PacketTooShort)

	// A good UDP packet with an AH header
	ip = layers.IPv6{
		Version:    6,
		NextHeader: layers.IPProtocolAH,
		HopLimit:   128,
		SrcIP:      net.IPv6linklocalallrouters,
		DstIP:      net.IPv6linklocalallnodes,
	}

	ah := layers.IPSecAH{
		AuthenticationData: []byte{0xde, 0xad, 0xbe, 0xef, 0xde, 0xad, 0xbe, 0xef},
	}
	ah.NextHeader = layers.IPProtocolUDP

	udpHeader := []byte{
		0x8d, 0x1b, // Source port 36123
		0x00, 0x16, // Destination port 22
		0x00, 0x00, // Length
		0x00, 0x00, // Checksum
	}

	buffer.Clear()
	err = ip.SerializeTo(buffer, opt)
	if err != nil {
		panic(err)
	}

	b = buffer.Bytes()
	ahb := serializeAH(&ah)
	b = append(b, ahb...)
	b = append(b, udpHeader...)

	err = newPacket(b, true, p)
	require.NoError(t, err)
	assert.Equal(t, uint8(firewall.ProtoUDP), p.Protocol)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.RemoteAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.LocalAddr)
	assert.Equal(t, uint16(36123), p.RemotePort)
	assert.Equal(t, uint16(22), p.LocalPort)
	assert.False(t, p.Fragment)

	// Ensure buffer bounds checking during processing
	err = newPacket(b[:41], true, p)
	require.ErrorIs(t, err, ErrIPv6PacketTooShort)

	// Invalid AH header
	b = buffer.Bytes()
	err = newPacket(b, true, p)
	require.ErrorIs(t, err, ErrIPv6CouldNotFindPayload)
}

func Test_newPacket_ipv6Fragment(t *testing.T) {
	p := &firewall.ParsedPacket{}

	ip := &layers.IPv6{
		Version:    6,
		NextHeader: layers.IPProtocolIPv6Fragment,
		HopLimit:   64,
		SrcIP:      net.IPv6linklocalallrouters,
		DstIP:      net.IPv6linklocalallnodes,
	}

	// First fragment
	fragHeader1 := []byte{
		uint8(layers.IPProtocolUDP), // Next Header (UDP)
		0x00,                        // Reserved
		0x00,                        // Fragment Offset high byte (0)
		0x01,                        // Fragment Offset low byte & flags (M=1)
		0x00, 0x00, 0x00, 0x01,      // Identification
	}

	udpHeader := []byte{
		0x8d, 0x1b, // Source port 36123
		0x00, 0x16, // Destination port 22
		0x00, 0x00, // Length
		0x00, 0x00, // Checksum
	}

	buffer := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}

	err := ip.SerializeTo(buffer, opts)
	if err != nil {
		t.Fatal(err)
	}

	firstFrag := buffer.Bytes()
	firstFrag = append(firstFrag, fragHeader1...)
	firstFrag = append(firstFrag, udpHeader...)
	firstFrag = append(firstFrag, []byte{0xde, 0xad, 0xbe, 0xef}...)

	// Test first fragment incoming
	err = newPacket(firstFrag, true, p)
	require.NoError(t, err)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.RemoteAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.LocalAddr)
	assert.Equal(t, uint8(layers.IPProtocolUDP), p.Protocol)
	assert.Equal(t, uint16(36123), p.RemotePort)
	assert.Equal(t, uint16(22), p.LocalPort)
	assert.False(t, p.Fragment)

	// Test first fragment outgoing
	err = newPacket(firstFrag, false, p)
	require.NoError(t, err)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.LocalAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.RemoteAddr)
	assert.Equal(t, uint8(layers.IPProtocolUDP), p.Protocol)
	assert.Equal(t, uint16(36123), p.LocalPort)
	assert.Equal(t, uint16(22), p.RemotePort)
	assert.False(t, p.Fragment)

	// Second fragment
	fragHeader2 := []byte{
		uint8(layers.IPProtocolUDP), // Next Header (UDP)
		0x00,                        // Reserved
		0xb9,                        // Fragment Offset high byte (185)
		0x01,                        // Fragment Offset low byte & flags (M=1)
		0x00, 0x00, 0x00, 0x01,      // Identification
	}

	buffer.Clear()
	err = ip.SerializeTo(buffer, opts)
	if err != nil {
		t.Fatal(err)
	}

	secondFrag := buffer.Bytes()
	secondFrag = append(secondFrag, fragHeader2...)
	secondFrag = append(secondFrag, []byte{0xde, 0xad, 0xbe, 0xef}...)

	// Test second fragment incoming
	err = newPacket(secondFrag, true, p)
	require.NoError(t, err)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.RemoteAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.LocalAddr)
	assert.Equal(t, uint8(layers.IPProtocolUDP), p.Protocol)
	assert.Equal(t, uint16(0), p.RemotePort)
	assert.Equal(t, uint16(0), p.LocalPort)
	assert.True(t, p.Fragment)

	// Test second fragment outgoing
	err = newPacket(secondFrag, false, p)
	require.NoError(t, err)
	assert.Equal(t, netip.MustParseAddr("ff02::2"), p.LocalAddr)
	assert.Equal(t, netip.MustParseAddr("ff02::1"), p.RemoteAddr)
	assert.Equal(t, uint8(layers.IPProtocolUDP), p.Protocol)
	assert.Equal(t, uint16(0), p.LocalPort)
	assert.Equal(t, uint16(0), p.RemotePort)
	assert.True(t, p.Fragment)

	// Too short of a fragment packet
	err = newPacket(secondFrag[:len(secondFrag)-10], false, p)
	require.ErrorIs(t, err, ErrIPv6PacketTooShort)
}

func BenchmarkParseV6(b *testing.B) {
	// Regular UDP packet
	ip := &layers.IPv6{
		Version:    6,
		NextHeader: layers.IPProtocolUDP,
		HopLimit:   64,
		SrcIP:      net.IPv6linklocalallrouters,
		DstIP:      net.IPv6linklocalallnodes,
	}

	udp := &layers.UDP{
		SrcPort: layers.UDPPort(36123),
		DstPort: layers.UDPPort(22),
	}

	buffer := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		ComputeChecksums: false,
		FixLengths:       true,
	}

	err := gopacket.SerializeLayers(buffer, opts, ip, udp)
	if err != nil {
		b.Fatal(err)
	}
	normalPacket := buffer.Bytes()

	// First Fragment packet
	ipFrag := &layers.IPv6{
		Version:    6,
		NextHeader: layers.IPProtocolIPv6Fragment,
		HopLimit:   64,
		SrcIP:      net.IPv6linklocalallrouters,
		DstIP:      net.IPv6linklocalallnodes,
	}

	fragHeader := []byte{
		uint8(layers.IPProtocolUDP), // Next Header (UDP)
		0x00,                        // Reserved
		0x00,                        // Fragment Offset high byte (0)
		0x01,                        // Fragment Offset low byte & flags (M=1)
		0x00, 0x00, 0x00, 0x01,      // Identification
	}

	udpHeader := []byte{
		0x8d, 0x7b, // Source port 36123
		0x00, 0x16, // Destination port 22
		0x00, 0x00, // Length
		0x00, 0x00, // Checksum
	}

	buffer.Clear()
	err = ipFrag.SerializeTo(buffer, opts)
	if err != nil {
		b.Fatal(err)
	}

	firstFrag := buffer.Bytes()
	firstFrag = append(firstFrag, fragHeader...)
	firstFrag = append(firstFrag, udpHeader...)
	firstFrag = append(firstFrag, []byte{0xde, 0xad, 0xbe, 0xef}...)

	// Second Fragment packet
	fragHeader[2] = 0xb9 // offset 185
	buffer.Clear()
	err = ipFrag.SerializeTo(buffer, opts)
	if err != nil {
		b.Fatal(err)
	}

	secondFrag := buffer.Bytes()
	secondFrag = append(secondFrag, fragHeader...)
	secondFrag = append(secondFrag, []byte{0xde, 0xad, 0xbe, 0xef}...)

	fp := &firewall.ParsedPacket{}

	b.Run("Normal", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			if err = parseV6(normalPacket, true, fp); err != nil {
				b.Fatal(err)
			}
		}
	})

	b.Run("FirstFragment", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			if err = parseV6(firstFrag, true, fp); err != nil {
				b.Fatal(err)
			}
		}
	})

	b.Run("SecondFragment", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			if err = parseV6(secondFrag, true, fp); err != nil {
				b.Fatal(err)
			}
		}
	})

	// Evil packet
	evilPacket := &layers.IPv6{
		Version:    6,
		NextHeader: layers.IPProtocolIPv6HopByHop,
		HopLimit:   64,
		SrcIP:      net.IPv6linklocalallrouters,
		DstIP:      net.IPv6linklocalallnodes,
	}

	hopHeader := []byte{
		uint8(layers.IPProtocolIPv6HopByHop), // Next Header (HopByHop)
		0x00,                                 // Length
		0x00, 0x00,                           // Options and padding
		0x00, 0x00, 0x00, 0x00, // More options and padding
	}

	lastHopHeader := []byte{
		uint8(layers.IPProtocolUDP), // Next Header (UDP)
		0x00,                        // Length
		0x00, 0x00,                  // Options and padding
		0x00, 0x00, 0x00, 0x00, // More options and padding
	}

	buffer.Clear()
	err = evilPacket.SerializeTo(buffer, opts)
	if err != nil {
		b.Fatal(err)
	}

	evilBytes := buffer.Bytes()
	for range 200 {
		evilBytes = append(evilBytes, hopHeader...)
	}
	evilBytes = append(evilBytes, lastHopHeader...)
	evilBytes = append(evilBytes, udpHeader...)
	evilBytes = append(evilBytes, []byte{0xde, 0xad, 0xbe, 0xef}...)

	b.Run("200 HopByHop headers", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			if err = parseV6(evilBytes, false, fp); err != nil {
				b.Fatal(err)
			}
		}
	})
}

// Ensure authentication data is a multiple of 8 bytes by padding if necessary
func padAuthData(authData []byte) []byte {
	// Length of Authentication Data must be a multiple of 8 bytes
	paddingLength := (8 - (len(authData) % 8)) % 8 // Only pad if necessary
	if paddingLength > 0 {
		authData = append(authData, make([]byte, paddingLength)...)
	}
	return authData
}

// Custom function to manually serialize IPSecAH for both IPv4 and IPv6
func serializeAH(ah *layers.IPSecAH) []byte {
	buf := new(bytes.Buffer)

	// Ensure Authentication Data is a multiple of 8 bytes
	ah.AuthenticationData = padAuthData(ah.AuthenticationData)
	// Calculate Payload Length (in 32-bit words, minus 2)
	payloadLen := uint8((12+len(ah.AuthenticationData))/4) - 2

	// Serialize fields
	if err := binary.Write(buf, binary.BigEndian, ah.NextHeader); err != nil {
		panic(err)
	}
	if err := binary.Write(buf, binary.BigEndian, payloadLen); err != nil {
		panic(err)
	}
	if err := binary.Write(buf, binary.BigEndian, ah.Reserved); err != nil {
		panic(err)
	}
	if err := binary.Write(buf, binary.BigEndian, ah.SPI); err != nil {
		panic(err)
	}
	if err := binary.Write(buf, binary.BigEndian, ah.Seq); err != nil {
		panic(err)
	}
	if len(ah.AuthenticationData) > 0 {
		if err := binary.Write(buf, binary.BigEndian, ah.AuthenticationData); err != nil {
			panic(err)
		}
	}

	return buf.Bytes()
}

// Test_newPacket_v6ExtHeaderOverflow is a regression test for the IPv6 extension-header
// length uint8 overflow in parseV6. A Destination-Options header with HdrExtLen=255 spans
// (255+1)*8 = 2048 bytes, so the real transport header sits at offset 2088. Before the fix
// the advance was computed in uint8 and wrapped to 0 (then clamped to 8), so the firewall
// read the transport header ~2KB too early from attacker-controlled option bytes while the
// host OS parses the real header, a firewall port/proto bypass. The fix makes parseV6 land
// on the same offset the host does.
func Test_newPacket_v6ExtHeaderOverflow(t *testing.T) {
	p := &firewall.ParsedPacket{}

	const (
		hdrLen      = 40              // IPv6 header
		extLen      = 2048            // (255+1)*8, the true Destination-Options header size
		realTCPAt   = hdrLen + extLen // 2088, where the host reads the transport header
		forgedTCPAt = hdrLen + 8      // 48, where the pre-fix wrapped+clamped walk landed
	)

	pkt := make([]byte, realTCPAt+4)
	pkt[0] = 0x60                                   // version 6
	pkt[6] = byte(layers.IPProtocolIPv6Destination) // NextHeader -> Destination Options
	pkt[40] = byte(firewall.ProtoTCP)               // Dest-Options NextHeader -> TCP
	pkt[41] = 255                                   // HdrExtLen = 255

	// Forged transport header at the pre-fix (wrong) offset: dst port 443.
	binary.BigEndian.PutUint16(pkt[forgedTCPAt+2:forgedTCPAt+4], 443)
	// Real transport header at the offset the host actually uses: dst port 22.
	binary.BigEndian.PutUint16(pkt[realTCPAt+2:realTCPAt+4], 22)

	require.NoError(t, newPacket(pkt, true, p))
	assert.Equal(t, uint8(firewall.ProtoTCP), p.Protocol)
	// LocalPort is the destination port for incoming traffic. It must be the real port (22)
	// the host delivers to, not the forged 443 at the overflowed offset.
	assert.Equal(t, uint16(22), p.LocalPort, "firewall must parse the real transport header, not the overflowed offset")
}

// Test_newPacket_parsedFields pins the ParsedPacket byproducts the RX
// batcher consumes: IPHdrLen (the true L4 offset) and FragAny (any fragment
// shape at all — unlike Packet.Fragment, which is port-oriented and true
// only for non-first fragments).
func Test_newPacket_parsedFields(t *testing.T) {
	p := &firewall.ParsedPacket{}

	// Plain IPv4 TCP, IHL 20: L4 offset 20, no fragment shape.
	v4 := make([]byte, 28)
	v4[0] = 0x45
	v4[9] = firewall.ProtoTCP
	binary.BigEndian.PutUint16(v4[6:8], 0x4000) // DF only
	require.NoError(t, newPacket(v4, true, p))
	assert.Equal(t, 20, p.IPHdrLen)
	assert.False(t, p.FragAny)
	assert.False(t, p.Fragment)

	// IPv4 first fragment (MF set, offset 0): the firewall can read ports
	// (Fragment false) but the coalescer must not touch it (FragAny true).
	ff := make([]byte, 28)
	ff[0] = 0x45
	ff[9] = firewall.ProtoUDP
	binary.BigEndian.PutUint16(ff[6:8], 0x2000) // MF, offset 0
	require.NoError(t, newPacket(ff, true, p))
	assert.False(t, p.Fragment)
	assert.True(t, p.FragAny)
	assert.Equal(t, 20, p.IPHdrLen)

	// IPv4 non-first fragment (nonzero offset): both flags set.
	nf := make([]byte, 28)
	nf[0] = 0x45
	nf[9] = firewall.ProtoUDP
	binary.BigEndian.PutUint16(nf[6:8], 0x00b9)
	require.NoError(t, newPacket(nf, true, p))
	assert.True(t, p.Fragment)
	assert.True(t, p.FragAny)

	// IPv4 with options (IHL 24): IPHdrLen tracks the real L4 offset.
	opts := make([]byte, 32)
	opts[0] = 0x46
	opts[9] = firewall.ProtoTCP
	binary.BigEndian.PutUint16(opts[6:8], 0x4000)
	require.NoError(t, newPacket(opts, true, p))
	assert.Equal(t, 24, p.IPHdrLen)
	assert.False(t, p.FragAny)

	// Plain IPv6 TCP: L4 at 40.
	v6 := make([]byte, 60)
	v6[0] = 0x60
	v6[6] = firewall.ProtoTCP
	require.NoError(t, newPacket(v6, true, p))
	assert.Equal(t, 40, p.IPHdrLen)
	assert.False(t, p.FragAny)

	// IPv6 hop-by-hop then TCP: IPHdrLen lands past the extension header.
	hbh := make([]byte, 60)
	hbh[0] = 0x60
	hbh[6] = 0 // hop-by-hop
	hbh[40] = firewall.ProtoTCP
	hbh[41] = 0 // HdrExtLen 0 -> 8-byte header
	require.NoError(t, newPacket(hbh, true, p))
	assert.Equal(t, 48, p.IPHdrLen)
	assert.False(t, p.FragAny)

	// IPv6 first fragment: terminal proto resolved, FragAny set, Fragment not.
	f6 := make([]byte, 60)
	f6[0] = 0x60
	f6[6] = 44 // fragment extension header
	f6[40] = firewall.ProtoUDP
	require.NoError(t, newPacket(f6, true, p))
	assert.True(t, p.FragAny)
	assert.False(t, p.Fragment)
	assert.Equal(t, uint8(firewall.ProtoUDP), p.Protocol)

	// IPv6 non-first fragment: both set, walk stops at the fragment header.
	f6n := make([]byte, 60)
	f6n[0] = 0x60
	f6n[6] = 44
	f6n[40] = firewall.ProtoUDP
	binary.BigEndian.PutUint16(f6n[42:44], 0x0008)
	require.NoError(t, newPacket(f6n, true, p))
	assert.True(t, p.Fragment)
	assert.True(t, p.FragAny)
}