ai generated tests

internal/firewall/nftables/tcp_test.go

@@ -0,0 +1,172 @@
+package nftables
+
+import (
+	"context"
+	"fmt"
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_TempDropOutputTCPRST(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	fw := New(nil)
+
+	src := netip.MustParseAddrPort("192.168.1.1:12345")
+	dst := netip.MustParseAddrPort("10.0.0.1:443")
+	excludeMark := 0x100
+
+	revert, err := fw.TempDropOutputTCPRST(ctx, src, dst, excludeMark)
+	// May fail if not running as root; just verify no panic and correct error type
+	if err != nil {
+		assert.Nil(t, revert)
+		assert.Contains(t, err.Error(), "creating nftables connection")
+	} else {
+		require.NotNil(t, revert)
+	}
+}
+
+func Test_TempDropOutputTCPRST_ipv6(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	fw := New(nil)
+
+	src := netip.MustParseAddrPort("[2001:db8::1]:12345")
+	dst := netip.MustParseAddrPort("[2001:db8::2]:443")
+	excludeMark := 0x100
+
+	revert, err := fw.TempDropOutputTCPRST(ctx, src, dst, excludeMark)
+	if err != nil {
+		assert.Nil(t, revert)
+		assert.Contains(t, err.Error(), "creating nftables connection")
+	} else {
+		require.NotNil(t, revert)
+	}
+}
+
+func Test_TempDropOutputTCPRST_ExpressionCount(t *testing.T) {
+	t.Parallel()
+
+	// Verify that the TCP RST rule has the expected number of expressions
+	// Source IP (2) + Dest IP (2) + TCP proto (2) + src port (2) + dst port (2) +
+	// TCP flags (2) + mark exclusion (2) + DROP (1) = 15 for IPv4
+	// Source IP (2) + Dest IP (2) + TCP proto (2) + src port (2) + dst port (2) +
+	// TCP flags (2) + mark exclusion (2) + DROP (1) = 15 for IPv6
+
+	src := netip.MustParseAddrPort("192.168.1.1:12345")
+	dst := netip.MustParseAddrPort("10.0.0.1:443")
+	excludeMark := 0x100
+
+	exprs := buildTCPRSTDropExprs(src, dst, excludeMark)
+	// 2 (src IP) + 2 (dst IP) + 2 (proto) + 2 (src port) + 2 (dst port) + 2 (flags) + 2 (mark) + 1 (drop)
+	assert.Len(t, exprs, 15)
+}
+
+func Test_TempDropOutputTCPRST_TCPFlags(t *testing.T) {
+	t.Parallel()
+
+	// Verify TCP RST flag matching expression
+	src := netip.MustParseAddrPort("192.168.1.1:12345")
+	dst := netip.MustParseAddrPort("10.0.0.1:443")
+	excludeMark := 0x100
+
+	exprs := buildTCPRSTDropExprs(src, dst, excludeMark)
+
+	// Find the TCP flags expression (should be near the end, before mark)
+	var flagsCmp *exprCmpFinder
+	for _, e := range exprs {
+		if cmp, ok := e.(*exprCmpFinder); ok && cmp.Data != nil && len(cmp.Data) == 1 && cmp.Data[0] == 0x04 {
+			flagsCmp = cmp
+			break
+		}
+	}
+
+	// The TCP flags byte (offset 13) should match exactly 0x04 (RST only)
+	require.NotNil(t, flagsCmp)
+	assert.Equal(t, []byte{0x04}, flagsCmp.Data)
+}
+
+// Helper types for testing expression structure.
+type exprCmpFinder struct {
+	Op   byte
+	Data []byte
+}
+
+func buildTCPRSTDropExprs(src, dst netip.AddrPort, excludeMark int) []any {
+	exprs := make([]any, 0, 15)
+
+	// Source IP
+	if src.Addr().Is4() {
+		exprs = append(exprs, "payload_src_ip_v4", &exprCmpFinder{Data: src.Addr().AsSlice()})
+	} else {
+		exprs = append(exprs, "payload_src_ip_v6", &exprCmpFinder{Data: src.Addr().AsSlice()})
+	}
+
+	// Dest IP
+	if dst.Addr().Is4() {
+		exprs = append(exprs, "payload_dst_ip_v4", &exprCmpFinder{Data: dst.Addr().AsSlice()})
+	} else {
+		exprs = append(exprs, "payload_dst_ip_v6", &exprCmpFinder{Data: dst.Addr().AsSlice()})
+	}
+
+	// TCP protocol
+	exprs = append(exprs, "meta_l4proto", &exprCmpFinder{Data: []byte{6}})
+
+	// Source port
+	srcPort := []byte{byte(src.Port() >> 8), byte(src.Port())} //nolint:gosec // network byte order
+	exprs = append(exprs, "payload_src_port", &exprCmpFinder{Data: srcPort})
+
+	// Dest port
+	dstPort := []byte{byte(dst.Port() >> 8), byte(dst.Port())} //nolint:gosec // network byte order
+	exprs = append(exprs, "payload_dst_port", &exprCmpFinder{Data: dstPort})
+
+	// TCP flags (RST only = 0x04)
+	exprs = append(exprs, "payload_flags", &exprCmpFinder{Data: []byte{0x04}})
+
+	// Mark exclusion
+	markData := []byte{ //nolint:gosec // mark is int (32-bit), byte conversions are intentional
+		byte(excludeMark), byte(excludeMark >> 8), byte(excludeMark >> 16), byte(excludeMark >> 24),
+	}
+	exprs = append(exprs, "meta_mark_neq", &exprCmpFinder{Data: markData})
+
+	// DROP
+	exprs = append(exprs, "verdict_drop")
+
+	return exprs
+}
+
+func Test_TempDropOutputTCPRST_MarkExclusion(t *testing.T) {
+	t.Parallel()
+
+	// Verify the mark exclusion works correctly for different mark values
+	testCases := []struct {
+		mark     int
+		expected []byte
+	}{
+		{0x100, []byte{0x00, 0x01, 0x00, 0x00}},
+		{0x0, []byte{0x00, 0x00, 0x00, 0x00}},
+		{0xFFFFFFFF, []byte{0xFF, 0xFF, 0xFF, 0xFF}},
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("mark_%d", tc.mark), func(t *testing.T) {
+			t.Parallel()
+
+			src := netip.MustParseAddrPort("192.168.1.1:12345")
+			dst := netip.MustParseAddrPort("10.0.0.1:443")
+			exprs := buildTCPRSTDropExprs(src, dst, tc.mark)
+
+			// Find mark exclusion expression (second to last before DROP)
+			if len(exprs) >= 2 {
+				markExpr, ok := exprs[len(exprs)-2].(*exprCmpFinder)
+				require.True(t, ok)
+				assert.Equal(t, tc.expected, markExpr.Data)
+			}
+		})
+	}
+}