From 614212a24cd4a94d185907f8237eb220a95fea43 Mon Sep 17 00:00:00 2001
From: Pasha Sviderski <me@psviderski.name>
Date: Fri, 25 Jul 2025 12:37:54 +1000
Subject: [PATCH] chore: clean up custom iptables chains on machine reset

---
 internal/machine/cluster.go                 |  5 ++--
 internal/machine/firewall/iptables_linux.go | 32 +++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/internal/machine/cluster.go b/internal/machine/cluster.go
index ed681d2e..f6d6c4ce 100644
--- a/internal/machine/cluster.go
+++ b/internal/machine/cluster.go
@@ -408,8 +408,9 @@ func (cc *clusterController) Cleanup() error {
 	if err := cc.wgnet.Cleanup(); err != nil {
 		errs = append(errs, fmt.Errorf("cleanup WireGuard network: %w", err))
 	}
-	// TODO: cleanup custom iptables chains. They're flushed when the machine is initialised again but it would
-	//  cleaner to delete them here.
+	if err := firewall.CleanupIptablesChains(); err != nil {
+		errs = append(errs, fmt.Errorf("cleanup iptables chains: %w", err))
+	}
 
 	return errors.Join(errs...)
 }
diff --git a/internal/machine/firewall/iptables_linux.go b/internal/machine/firewall/iptables_linux.go
index 7928a71f..f0e25339 100644
--- a/internal/machine/firewall/iptables_linux.go
+++ b/internal/machine/firewall/iptables_linux.go
@@ -2,6 +2,7 @@ package firewall
 
 import (
 	"fmt"
+	"log/slog"
 	"strconv"
 	"strings"
 
@@ -69,3 +70,34 @@ func ConfigureIptablesChains() error {
 
 	return nil
 }
+
+// CleanupIptablesChains removes the custom iptables chains and rules created by ConfigureIptablesChains.
+func CleanupIptablesChains() error {
+	ipt := iptables.GetIptable(iptables.IPv4)
+
+	// First, remove the jump rule from INPUT chain to UNCLOUD-INPUT.
+	jumpRule := []string{"-m", "comment", "--comment", "Uncloud-managed", "-j", UncloudInputChain}
+	if err := ipt.ProgramRule(iptables.Filter, "INPUT", iptables.Delete, jumpRule); err != nil {
+		return fmt.Errorf("delete iptables jump rule from INPUT: %w", err)
+	}
+
+	// Flush all rules from UNCLOUD-INPUT chain as it must be empty before deletion.
+	if err := ipt.RawCombinedOutput("-t", string(iptables.Filter), "-F", UncloudInputChain); err != nil {
+		// Chain might not exist which is fine.
+		if !strings.Contains(err.Error(), "No chain") {
+			return fmt.Errorf("flush iptables chain '%s': %w", UncloudInputChain, err)
+		}
+	}
+
+	// Delete the UNCLOUD-INPUT chain.
+	if err := ipt.RawCombinedOutput("-t", string(iptables.Filter), "-X", UncloudInputChain); err != nil {
+		// Chain might not exist which is fine.
+		if !strings.Contains(err.Error(), "No chain") {
+			return fmt.Errorf("delete iptables chain '%s': %w", UncloudInputChain, err)
+		}
+	} else {
+		slog.Info("Deleted iptables chain.", "chain", UncloudInputChain)
+	}
+
+	return nil
+}