Fix tests

Author: lixmal

client/firewall/uspfilter/nat_stateful_test.go

@@ -10,9 +10,8 @@ import (
 	"github.com/netbirdio/netbird/client/iface/device"
 )
 
-// TestStatefulNATBidirectionalSSH tests that stateful NAT prevents interference
-// when two peers try to SSH to each other simultaneously
-func TestStatefulNATBidirectionalSSH(t *testing.T) {
+// TestPortDNATBasic tests basic port DNAT functionality
+func TestPortDNATBasic(t *testing.T) {
 	manager, err := Create(&IFaceMock{
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
 	}, false, flowLogger)
@@ -30,46 +29,25 @@ func TestStatefulNATBidirectionalSSH(t *testing.T) {
 	require.NoError(t, err)
 
 	// Scenario: Peer A connects to Peer B on port 22 (should get NAT)
-	// This simulates: ssh [email protected]
 	packetAtoB := generateDNATTestPacket(t, peerA, peerB, layers.IPProtocolTCP, 54321, 22)
-	translatedAtoB := manager.translateInboundPortDNAT(packetAtoB, parsePacket(t, packetAtoB))
+	d := parsePacket(t, packetAtoB)
+	translatedAtoB := manager.translateInboundPortDNAT(packetAtoB, d, peerA, peerB)
 	require.True(t, translatedAtoB, "Peer A to Peer B should be translated (NAT applied)")
 
 	// Verify port was translated to 22022
-	d := parsePacket(t, packetAtoB)
+	d = parsePacket(t, packetAtoB)
 	require.Equal(t, uint16(22022), uint16(d.tcp.DstPort), "Port should be rewritten to 22022")
 
-	// Verify NAT connection is tracked (with translated port as key)
-	natConn, exists := manager.portNATTracker.getConnectionNAT(peerA, peerB, 54321, 22022)
-	require.True(t, exists, "NAT connection should be tracked")
-	require.Equal(t, uint16(22), natConn.originalPort, "Original port should be stored")
-
-	// Scenario: Peer B tries to connect to Peer A on port 22 (should NOT get NAT)
-	// This simulates the reverse direction to prevent interference
-	packetBtoA := generateDNATTestPacket(t, peerB, peerA, layers.IPProtocolTCP, 54322, 22)
-	translatedBtoA := manager.translateInboundPortDNAT(packetBtoA, parsePacket(t, packetBtoA))
-	require.False(t, translatedBtoA, "Peer B to Peer A should NOT be translated (prevent interference)")
-
-	// Verify port was NOT translated
-	d2 := parsePacket(t, packetBtoA)
-	require.Equal(t, uint16(22), uint16(d2.tcp.DstPort), "Port should remain 22 (no translation)")
-
-	// Verify no reverse NAT connection is tracked
-	_, reverseExists := manager.portNATTracker.getConnectionNAT(peerB, peerA, 54322, 22)
-	require.False(t, reverseExists, "Reverse NAT connection should NOT be tracked")
-
-	// Scenario: Return traffic from Peer B (SSH server) to Peer A (should be reverse translated)
+	// Scenario: Return traffic from Peer B to Peer A should NOT be translated
+	// (prevents double NAT - original port stored in conntrack)
 	returnPacket := generateDNATTestPacket(t, peerB, peerA, layers.IPProtocolTCP, 22022, 54321)
-	translatedReturn := manager.translateOutboundPortReverse(returnPacket, parsePacket(t, returnPacket))
-	require.True(t, translatedReturn, "Return traffic should be reverse translated")
-
-	// Verify return traffic port was translated back to 22
-	d3 := parsePacket(t, returnPacket)
-	require.Equal(t, uint16(22), uint16(d3.tcp.SrcPort), "Return traffic source port should be 22")
+	d2 := parsePacket(t, returnPacket)
+	translatedReturn := manager.translateInboundPortDNAT(returnPacket, d2, peerB, peerA)
+	require.False(t, translatedReturn, "Return traffic from same IP should not be translated")
 }
 
-// TestStatefulNATConnectionCleanup tests connection cleanup functionality
-func TestStatefulNATConnectionCleanup(t *testing.T) {
+// TestPortDNATMultipleRules tests multiple port DNAT rules
+func TestPortDNATMultipleRules(t *testing.T) {
 	manager, err := Create(&IFaceMock{
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
 	}, false, flowLogger)
@@ -88,24 +66,19 @@ func TestStatefulNATConnectionCleanup(t *testing.T) {
 	err = manager.addPortRedirection(peerB, layers.LayerTypeTCP, 22, 22022)
 	require.NoError(t, err)
 
-	// Establish connection with NAT
-	packet := generateDNATTestPacket(t, peerA, peerB, layers.IPProtocolTCP, 54321, 22)
-	translated := manager.translateInboundPortDNAT(packet, parsePacket(t, packet))
-	require.True(t, translated, "Initial connection should be translated")
-
-	// Verify connection is tracked (using translated port as key)
-	_, exists := manager.portNATTracker.getConnectionNAT(peerA, peerB, 54321, 22022)
-	require.True(t, exists, "Connection should be tracked")
-
-	// Clean up connection
-	manager.portNATTracker.cleanupConnection(peerA, peerB, 54321)
-
-	// Verify connection is no longer tracked (using translated port as key)
-	_, stillExists := manager.portNATTracker.getConnectionNAT(peerA, peerB, 54321, 22022)
-	require.False(t, stillExists, "Connection should be cleaned up")
-
-	// Verify new connection from opposite direction now works
-	reversePacket := generateDNATTestPacket(t, peerB, peerA, layers.IPProtocolTCP, 54322, 22)
-	reverseTranslated := manager.translateInboundPortDNAT(reversePacket, parsePacket(t, reversePacket))
-	require.True(t, reverseTranslated, "Reverse connection should now work after cleanup")
+	// Test traffic to peer B gets translated
+	packetToB := generateDNATTestPacket(t, peerA, peerB, layers.IPProtocolTCP, 54321, 22)
+	d1 := parsePacket(t, packetToB)
+	translatedToB := manager.translateInboundPortDNAT(packetToB, d1, peerA, peerB)
+	require.True(t, translatedToB, "Traffic to peer B should be translated")
+	d1 = parsePacket(t, packetToB)
+	require.Equal(t, uint16(22022), uint16(d1.tcp.DstPort), "Port should be 22022")
+
+	// Test traffic to peer A gets translated
+	packetToA := generateDNATTestPacket(t, peerB, peerA, layers.IPProtocolTCP, 54322, 22)
+	d2 := parsePacket(t, packetToA)
+	translatedToA := manager.translateInboundPortDNAT(packetToA, d2, peerB, peerA)
+	require.True(t, translatedToA, "Traffic to peer A should be translated")
+	d2 = parsePacket(t, packetToA)
+	require.Equal(t, uint16(22022), uint16(d2.tcp.DstPort), "Port should be 22022")
 }