Fix tests

netbirdio · Jan 15, 2025 · 7f08282 · 7f08282
1 parent 2b9c610
commit 7f08282
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 26 deletions.
diff --git a/client/firewall/iptables/manager_linux_test.go b/client/firewall/iptables/manager_linux_test.go
@@ -215,11 +215,7 @@ func TestIptablesCreatePerformance(t *testing.T) {
 			start := time.Now()
 			for i := 0; i < testMax; i++ {
 				port := &fw.Port{Values: []int{1000 + i}}
-				if i%2 == 0 {
-					_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
-				} else {
-					_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
-				}
+				_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
 
 				require.NoError(t, err, "failed to add rule")
 			}

diff --git a/client/firewall/nftables/manager_linux_test.go b/client/firewall/nftables/manager_linux_test.go
@@ -201,11 +201,7 @@ func TestNFtablesCreatePerformance(t *testing.T) {
 			start := time.Now()
 			for i := 0; i < testMax; i++ {
 				port := &fw.Port{Values: []int{1000 + i}}
-				if i%2 == 0 {
-					_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
-				} else {
-					_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
-				}
+				_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
 				require.NoError(t, err, "failed to add rule")
 
 				if i%100 == 0 {

diff --git a/client/firewall/uspfilter/uspfilter.go b/client/firewall/uspfilter/uspfilter.go
@@ -578,19 +578,22 @@ func (m *Manager) AddUDPPacketHook(
 
 // RemovePacketHook removes packet hook by given ID
 func (m *Manager) RemovePacketHook(hookID string) error {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
 	for _, arr := range m.incomingRules {
 		for _, r := range arr {
 			if r.id == hookID {
-				rule := r
-				return m.DeletePeerRule(&rule)
+				delete(arr, r.id)
+				return nil
 			}
 		}
 	}
 	for _, arr := range m.outgoingRules {
 		for _, r := range arr {
 			if r.id == hookID {
-				rule := r
-				return m.DeletePeerRule(&rule)
+				delete(arr, r.id)
+				return nil
 			}
 		}
 	}

diff --git a/client/firewall/uspfilter/uspfilter_bench_test.go b/client/firewall/uspfilter/uspfilter_bench_test.go
@@ -91,7 +91,7 @@ func BenchmarkCoreFiltering(b *testing.B) {
 			setupFunc: func(m *Manager) {
 				// Single rule allowing all traffic
 				_, err := m.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolALL, nil, nil,
-					fw.RuleDirectionIN, fw.ActionAccept, "", "allow all")
+					fw.ActionAccept, "", "allow all")
 				require.NoError(b, err)
 			},
 			desc: "Baseline: Single 'allow all' rule without connection tracking",
@@ -114,7 +114,7 @@ func BenchmarkCoreFiltering(b *testing.B) {
 					_, err := m.AddPeerFiltering(ip, fw.ProtocolTCP,
 						&fw.Port{Values: []int{1024 + i}},
 						&fw.Port{Values: []int{80}},
-						fw.RuleDirectionIN, fw.ActionAccept, "", "explicit return")
+						fw.ActionAccept, "", "explicit return")
 					require.NoError(b, err)
 				}
 			},
@@ -126,7 +126,7 @@ func BenchmarkCoreFiltering(b *testing.B) {
 			setupFunc: func(m *Manager) {
 				// Add some basic rules but rely on state for established connections
 				_, err := m.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP, nil, nil,
-					fw.RuleDirectionIN, fw.ActionDrop, "", "default drop")
+					fw.ActionDrop, "", "default drop")
 				require.NoError(b, err)
 			},
 			desc: "Connection tracking with established connections",
@@ -590,7 +590,7 @@ func BenchmarkLongLivedConnections(b *testing.B) {
 				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
 					&fw.Port{Values: []int{80}},
 					nil,
-					fw.RuleDirectionIN, fw.ActionAccept, "", "return traffic")
+					fw.ActionAccept, "", "return traffic")
 				require.NoError(b, err)
 			}
 
@@ -681,7 +681,7 @@ func BenchmarkShortLivedConnections(b *testing.B) {
 				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
 					&fw.Port{Values: []int{80}},
 					nil,
-					fw.RuleDirectionIN, fw.ActionAccept, "", "return traffic")
+					fw.ActionAccept, "", "return traffic")
 				require.NoError(b, err)
 			}
 
@@ -799,7 +799,7 @@ func BenchmarkParallelLongLivedConnections(b *testing.B) {
 				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
 					&fw.Port{Values: []int{80}},
 					nil,
-					fw.RuleDirectionIN, fw.ActionAccept, "", "return traffic")
+					fw.ActionAccept, "", "return traffic")
 				require.NoError(b, err)
 			}
 
@@ -886,7 +886,7 @@ func BenchmarkParallelShortLivedConnections(b *testing.B) {
 				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
 					&fw.Port{Values: []int{80}},
 					nil,
-					fw.RuleDirectionIN, fw.ActionAccept, "", "return traffic")
+					fw.ActionAccept, "", "return traffic")
 				require.NoError(b, err)
 			}
 

diff --git a/client/firewall/uspfilter/uspfilter_test.go b/client/firewall/uspfilter/uspfilter_test.go
@@ -261,7 +261,6 @@ func TestNotMatchByIP(t *testing.T) {
 
 	ip := net.ParseIP("0.0.0.0")
 	proto := fw.ProtocolUDP
-	direction := fw.RuleDirectionOUT
 	action := fw.ActionAccept
 	comment := "Test rule"
 

diff --git a/client/internal/acl/manager_test.go b/client/internal/acl/manager_test.go
@@ -119,8 +119,8 @@ func TestDefaultManager(t *testing.T) {
 
 		networkMap.FirewallRulesIsEmpty = false
 		acl.ApplyFiltering(networkMap)
-		if len(acl.peerRulesPairs) != 2 {
-			t.Errorf("rules should contain 2 rules if FirewallRulesIsEmpty is not set, got: %v", len(acl.peerRulesPairs))
+		if len(acl.peerRulesPairs) != 1 {
+			t.Errorf("rules should contain 1 rules if FirewallRulesIsEmpty is not set, got: %v", len(acl.peerRulesPairs))
 			return
 		}
 	})
@@ -356,8 +356,8 @@ func TestDefaultManagerEnableSSHRules(t *testing.T) {
 
 	acl.ApplyFiltering(networkMap)
 
-	if len(acl.peerRulesPairs) != 4 {
-		t.Errorf("expect 4 rules (last must be SSH), got: %d", len(acl.peerRulesPairs))
+	if len(acl.peerRulesPairs) != 3 {
+		t.Errorf("expect 3 rules (last must be SSH), got: %d", len(acl.peerRulesPairs))
 		return
 	}
 }