Fix tests

This commit is contained in:
Viktor Liu
2025-10-28 21:11:45 +01:00
parent 576b4a779c
commit a7a85d4dc8

View File

@@ -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 user@100.10.0.51
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")
// 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")
// 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 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")
}