From df5f79dccc711dbb598a0683a133b87bec534d85 Mon Sep 17 00:00:00 2001
From: Alex Stockinger <alex@atomicjar.com>
Date: Thu, 28 Jul 2022 07:31:21 +0000
Subject: [PATCH] Make cleanup on shutdown optional

---
 cmd/kg/main.go   |  6 ++++--
 pkg/mesh/mesh.go | 17 +++++++++++++----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/cmd/kg/main.go b/cmd/kg/main.go
index 1834653..7c91698 100644
--- a/cmd/kg/main.go
+++ b/cmd/kg/main.go
@@ -95,6 +95,7 @@ var cmd = &cobra.Command{
 
 var (
 	backend               string
+	cleanUp               bool
 	cleanUpIface          bool
 	createIface           bool
 	cni                   bool
@@ -125,7 +126,8 @@ var (
 
 func init() {
 	cmd.Flags().StringVar(&backend, "backend", k8s.Backend, fmt.Sprintf("The backend for the mesh. Possible values: %s", availableBackends))
-	cmd.Flags().BoolVar(&cleanUpIface, "clean-up-interface", false, "Should Kilo delete its interface when it shuts down?")
+	cmd.Flags().BoolVar(&cleanUp, "clean-up", false, "Clean up network modifications on shutdown.")
+	cmd.Flags().BoolVar(&cleanUpIface, "clean-up-interface", false, "Should Kilo delete its interface when it shuts down? Only relevant when cleanup is true.")
 	cmd.Flags().BoolVar(&createIface, "create-interface", true, "Should kilo create an interface on startup?")
 	cmd.Flags().BoolVar(&cni, "cni", true, "Should Kilo manage the node's CNI configuration?")
 	cmd.Flags().StringVar(&cniPath, "cni-path", mesh.DefaultCNIPath, "Path to CNI config.")
@@ -245,7 +247,7 @@ func runRoot(_ *cobra.Command, _ []string) error {
 	if port < 1 || port > 1<<16-1 {
 		return fmt.Errorf("invalid port: port mus be in range [%d:%d], but got %d", 1, 1<<16-1, port)
 	}
-	m, err := mesh.New(b, enc, gr, hostname, port, s, local, cni, cniPath, iface, cleanUpIface, createIface, mtu, resyncPeriod, prioritisePrivateAddr, iptablesForwardRule, log.With(logger, "component", "kilo"))
+	m, err := mesh.New(b, enc, gr, hostname, port, s, local, cni, cniPath, iface, cleanUp, cleanUpIface, createIface, mtu, resyncPeriod, prioritisePrivateAddr, iptablesForwardRule, log.With(logger, "component", "kilo"))
 	if err != nil {
 		return fmt.Errorf("failed to create Kilo mesh: %v", err)
 	}
diff --git a/pkg/mesh/mesh.go b/pkg/mesh/mesh.go
index 22c54e1..b0b9a74 100644
--- a/pkg/mesh/mesh.go
+++ b/pkg/mesh/mesh.go
@@ -51,6 +51,7 @@ const (
 // Mesh is able to create Kilo network meshes.
 type Mesh struct {
 	Backend
+	cleanup             bool
 	cleanUpIface        bool
 	cni                 bool
 	cniPath             string
@@ -88,7 +89,7 @@ type Mesh struct {
 }
 
 // New returns a new Mesh instance.
-func New(backend Backend, enc encapsulation.Encapsulator, granularity Granularity, hostname string, port int, subnet *net.IPNet, local, cni bool, cniPath, iface string, cleanUpIface bool, createIface bool, mtu uint, resyncPeriod time.Duration, prioritisePrivateAddr, iptablesForwardRule bool, logger log.Logger) (*Mesh, error) {
+func New(backend Backend, enc encapsulation.Encapsulator, granularity Granularity, hostname string, port int, subnet *net.IPNet, local, cni bool, cniPath, iface string, cleanup bool, cleanUpIface bool, createIface bool, mtu uint, resyncPeriod time.Duration, prioritisePrivateAddr, iptablesForwardRule bool, logger log.Logger) (*Mesh, error) {
 	if err := os.MkdirAll(kiloPath, 0700); err != nil {
 		return nil, fmt.Errorf("failed to create directory to store configuration: %v", err)
 	}
@@ -117,9 +118,14 @@ func New(backend Backend, enc encapsulation.Encapsulator, granularity Granularit
 	}
 	var kiloIface int
 	if createIface {
-		kiloIface, _, err = wireguard.New(iface, mtu)
+		link, err := netlink.LinkByName(iface)
 		if err != nil {
-			return nil, fmt.Errorf("failed to create WireGuard interface: %v", err)
+			kiloIface, _, err = wireguard.New(iface, mtu)
+			if err != nil {
+				return nil, fmt.Errorf("failed to create WireGuard interface: %v", err)
+			}
+		} else {
+			kiloIface = link.Attrs().Index
 		}
 	} else {
 		link, err := netlink.LinkByName(iface)
@@ -162,6 +168,7 @@ func New(backend Backend, enc encapsulation.Encapsulator, granularity Granularit
 	}
 	return &Mesh{
 		Backend:             backend,
+		cleanup:             cleanup,
 		cleanUpIface:        cleanUpIface,
 		cni:                 cni,
 		cniPath:             cniPath,
@@ -248,7 +255,9 @@ func (m *Mesh) Run(ctx context.Context) error {
 			}
 		}
 	}()
-	defer m.cleanUp()
+	if m.cleanup {
+		defer m.cleanUp()
+	}
 	resync := time.NewTimer(m.resyncPeriod)
 	checkIn := time.NewTimer(checkInPeriod)
 	nw := m.Nodes().Watch()