From 4f035d3f99043ebdf2287014e6882a079add2365 Mon Sep 17 00:00:00 2001
From: Viktor Liu <viktor@netbird.io>
Date: Fri, 25 Oct 2024 15:26:54 +0200
Subject: [PATCH] Remove legacy forwarding rules in userspace mode

---
 client/firewall/iptables/router_linux.go  |  2 ++
 client/firewall/nftables/manager_linux.go | 18 +-----------------
 client/firewall/nftables/router_linux.go  |  3 +++
 client/firewall/uspfilter/uspfilter.go    |  7 +++++--
 4 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/client/firewall/iptables/router_linux.go b/client/firewall/iptables/router_linux.go
index 90811ae1182..9b75640b4b5 100644
--- a/client/firewall/iptables/router_linux.go
+++ b/client/firewall/iptables/router_linux.go
@@ -296,6 +296,8 @@ func (r *router) RemoveAllLegacyRouteRules() error {
 		}
 		if err := r.iptablesClient.DeleteIfExists(tableFilter, chainRTFWD, rule...); err != nil {
 			merr = multierror.Append(merr, fmt.Errorf("remove legacy forwarding rule: %v", err))
+		} else {
+			delete(r.rules, k)
 		}
 	}
 
diff --git a/client/firewall/nftables/manager_linux.go b/client/firewall/nftables/manager_linux.go
index a4650f3b626..ea8912f27f5 100644
--- a/client/firewall/nftables/manager_linux.go
+++ b/client/firewall/nftables/manager_linux.go
@@ -230,23 +230,7 @@ func (m *Manager) AllowNetbird() error {
 
 // SetLegacyManagement sets the route manager to use legacy management
 func (m *Manager) SetLegacyManagement(isLegacy bool) error {
-	oldLegacy := m.router.legacyManagement
-
-	if oldLegacy != isLegacy {
-		m.router.legacyManagement = isLegacy
-		log.Debugf("Set legacy management to %v", isLegacy)
-	}
-
-	// client reconnected to a newer mgmt, we need to cleanup the legacy rules
-	if !isLegacy && oldLegacy {
-		if err := m.router.RemoveAllLegacyRouteRules(); err != nil {
-			return fmt.Errorf("remove legacy routing rules: %v", err)
-		}
-
-		log.Debugf("Legacy routing rules removed")
-	}
-
-	return nil
+	return firewall.SetLegacyManagement(m.router, isLegacy)
 }
 
 // Reset firewall to the default state
diff --git a/client/firewall/nftables/router_linux.go b/client/firewall/nftables/router_linux.go
index 9b28e4eb213..0e7ea71b774 100644
--- a/client/firewall/nftables/router_linux.go
+++ b/client/firewall/nftables/router_linux.go
@@ -551,7 +551,10 @@ func (r *router) RemoveAllLegacyRouteRules() error {
 		}
 		if err := r.conn.DelRule(rule); err != nil {
 			merr = multierror.Append(merr, fmt.Errorf("remove legacy forwarding rule: %v", err))
+		} else {
+			delete(r.rules, k)
 		}
+
 	}
 	return nberrors.FormatErrorOrNil(merr)
 }
diff --git a/client/firewall/uspfilter/uspfilter.go b/client/firewall/uspfilter/uspfilter.go
index 3829a9baffe..af5dc673393 100644
--- a/client/firewall/uspfilter/uspfilter.go
+++ b/client/firewall/uspfilter/uspfilter.go
@@ -237,8 +237,11 @@ func (m *Manager) DeletePeerRule(rule firewall.Rule) error {
 }
 
 // SetLegacyManagement doesn't need to be implemented for this manager
-func (m *Manager) SetLegacyManagement(_ bool) error {
-	return nil
+func (m *Manager) SetLegacyManagement(isLegacy bool) error {
+	if m.nativeFirewall == nil {
+		return errRouteNotSupported
+	}
+	return m.nativeFirewall.SetLegacyManagement(isLegacy)
 }
 
 // Flush doesn't need to be implemented for this manager