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"
	"golang.org/x/net/ipv6"
)

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, ErrIPv6PacketTooShort)

	// A v6 packet with a hop-by-hop extension
	// ICMPv6 Payload (Echo Request)
	icmpLayer := layers.ICMPv6{
		TypeCode: layers.CreateICMPv6TypeCode(layers.ICMPv6TypeEchoRequest, 0),
	}
	// 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, ErrIPv6PacketTooShort)

	// 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, ErrIPv6PacketTooShort)
	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.CreateICMPv6TypeCode(layers.ICMPv6TypeEchoRequest, 0),
		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 minimal 4 byte non-echo ICMPv6 message (type, code, checksum), no identifier to read
	icmpMin := make([]byte, ipv6.HeaderLen+4)
	copy(icmpMin, buffer.Bytes()[:ipv6.HeaderLen])
	icmpMin[6] = byte(layers.IPProtocolICMPv6)
	icmpMin[ipv6.HeaderLen] = 1 // type 1, destination unreachable, not echo
	err = newPacket(icmpMin, true, p)
	require.NoError(t, err)
	assert.Equal(t, uint8(layers.IPProtocolICMPv6), p.Protocol)
	assert.Equal(t, uint16(0), 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, we don't dissect it so we fail closed on its true protocol with no ports
	b = buffer.Bytes()
	b[6] = 255 // 255 is a reserved protocol number
	err = newPacket(b, true, p)
	require.NoError(t, err)
	assert.Equal(t, uint8(255), p.Protocol)
	assert.Equal(t, uint16(0), p.RemotePort)
	assert.Equal(t, uint16(0), p.LocalPort)
	assert.False(t, p.Fragment)

	// 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, a truncated AH header can't reach the payload
	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, ErrIPv6PacketTooShort)
}

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_v6ExtHeaderPastBuffer is a regression test for an extension header whose declared length
// advances the walk past the end of the packet. The upper layer protocol's header isn't actually present,
// so parseV6 must drop the packet rather than classify it as the terminal protocol with no ports.
func Test_newPacket_v6ExtHeaderPastBuffer(t *testing.T) {
	p := &firewall.ParsedPacket{}

	pkt := make([]byte, 48)
	pkt[0] = 0x60
	pkt[6] = byte(layers.IPProtocolIPv6Destination) // Destination Options
	pkt[7] = 64                                     // hop limit
	pkt[40] = byte(layers.IPProtocolSCTP)           // Dest Options next header = SCTP
	pkt[41] = 255                                   // declared length (255+1)*8 = 2048, past the 48 byte buffer

	require.ErrorIs(t, newPacket(pkt, true, p), ErrIPv6PacketTooShort)
}

// Test_newPacket_v6ExtHeaderConfusion is a regression test for parseV6 walking any unrecognized
// Next Header as if it were an ipv6 extension header. A real upper layer protocol Nebula doesn't
// dissect (SCTP here) is not walkable, so applying the (len+1)*8 formula marched into the SCTP
// payload and landed on a byte that looked like UDP, forging a protocol/port pair the firewall
// would trust while the host delivered the real SCTP datagram. The fix fails closed: the packet
// is classified as its true protocol with no ports, so it only matches an `any` rule.
func Test_newPacket_v6ExtHeaderConfusion(t *testing.T) {
	p := &firewall.ParsedPacket{}

	pkt := make([]byte, 52)
	pkt[0] = 0x60                        // version 6
	pkt[6] = byte(layers.IPProtocolSCTP) // NextHeader = SCTP, a real protocol, not an extension header
	pkt[7] = 64                          // hop limit

	// Real SCTP header at offset 40. Pre-fix parseV6 walked SCTP as an extension header: byte 41 (0x00, the
	// low byte of the src port below) was read as the header length, giving next=(0+1)*8=8, which landed the
	// walk on byte 40 (0x11), misread as NextHeader=UDP, then bytes 48-51 as ports.
	binary.BigEndian.PutUint16(pkt[40:42], 0x1100) // SCTP src port; byte 40=0x11, byte 41=0x00
	binary.BigEndian.PutUint16(pkt[42:44], 445)    // SCTP dst port, never read by parseV6
	binary.BigEndian.PutUint16(pkt[48:50], 53)     // SCTP checksum bytes, pre-fix forged RemotePort
	binary.BigEndian.PutUint16(pkt[50:52], 53)     // pre-fix forged LocalPort

	require.NoError(t, newPacket(pkt, true, p))
	assert.Equal(t, uint8(layers.IPProtocolSCTP), p.Protocol, "must classify as the true protocol, not the forged UDP")
	assert.Equal(t, uint16(0), p.RemotePort)
	assert.Equal(t, uint16(0), p.LocalPort)
	assert.False(t, p.Fragment)

	// Same confusion, but the unknown protocol sits after a real extension header. The HopByHop is walked
	// correctly, then SCTP must still fail closed instead of being walked into its own payload. Protocol is
	// the only assertion that discriminates the fix here, a regression that walked SCTP would misclassify it.
	chained := make([]byte, 60)
	chained[0] = 0x60                                  // version 6
	chained[6] = byte(layers.IPProtocolIPv6HopByHop)   // NextHeader = HopByHop extension
	chained[7] = 64                                    // hop limit
	chained[40] = byte(layers.IPProtocolSCTP)          // HopByHop NextHeader = SCTP
	chained[41] = 0                                    // HopByHop length 0 -> 8 bytes, SCTP begins at offset 48
	binary.BigEndian.PutUint16(chained[48:50], 0x1100) // SCTP src port, pre-fix forged NextHeader/length bait
	binary.BigEndian.PutUint16(chained[50:52], 445)    // SCTP dst port, never read by parseV6

	require.NoError(t, newPacket(chained, true, p))
	assert.Equal(t, uint8(layers.IPProtocolSCTP), p.Protocol, "must fail closed on the unknown protocol after the extension header")
	assert.Equal(t, uint16(0), p.RemotePort)
	assert.Equal(t, uint16(0), p.LocalPort)
	assert.False(t, p.Fragment)
}

// 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)
}