From c0ef6ae9c4ecbb77537a448d33b67276b671b148 Mon Sep 17 00:00:00 2001
From: Alejandro Mery <amery@jpi.io>
Date: Sat, 26 Aug 2023 17:14:54 +0000
Subject: [PATCH] zones: rearrange code around WireguardConfigPruner

Signed-off-by: Alejandro Mery <amery@jpi.io>
---
 pkg/zones/wireguard.go  | 49 +++++++++++++++++++++++++++++++++++++++++
 pkg/zones/zone_rings.go | 27 -----------------------
 2 files changed, 49 insertions(+), 27 deletions(-)
 delete mode 100644 pkg/zones/zone_rings.go

diff --git a/pkg/zones/wireguard.go b/pkg/zones/wireguard.go
index 7fd622a..fb46865 100644
--- a/pkg/zones/wireguard.go
+++ b/pkg/zones/wireguard.go
@@ -6,11 +6,60 @@ import (
 )
 
 var (
+	_ WireguardConfigPruner = (*Zones)(nil)
+	_ WireguardConfigPruner = (*Zone)(nil)
+	_ WireguardConfigPruner = (*Machine)(nil)
+
 	_ WireguardKeysWriter = (*Zones)(nil)
 	_ WireguardKeysWriter = (*Zone)(nil)
 	_ WireguardKeysWriter = (*Machine)(nil)
 )
 
+// A WireguardConfigPruner deletes wgN.conf on all machines under
+// its scope with the specified ring disabled
+type WireguardConfigPruner interface {
+	PruneWireguardConfig(ring int) error
+}
+
+// PruneWireguardConfig removes wgN.conf files of machines with
+// the corresponding ring disabled on all zones
+func (m *Zones) PruneWireguardConfig(ring int) error {
+	return pruneWireguardConfig(m, ring)
+}
+
+// PruneWireguardConfig removes wgN.conf files of machines with
+// the corresponding ring disabled.
+func (z *Zone) PruneWireguardConfig(ring int) error {
+	return pruneWireguardConfig(z, ring)
+}
+
+func pruneWireguardConfig(m MachineIterator, ring int) error {
+	var err error
+
+	m.ForEachMachine(func(p *Machine) bool {
+		err = p.zone.PruneWireguardConfig(ring)
+		if os.IsNotExist(err) {
+			// ignore
+			err = nil
+		}
+
+		return err != nil
+	})
+
+	return err
+}
+
+// PruneWireguardConfig deletes the wgN.conf file if its
+// presence on the ring is disabled
+func (m *Machine) PruneWireguardConfig(ring int) error {
+	_, ok := m.getRingInfo(ring)
+	if !ok {
+		return m.RemoveWireguardConfig(ring)
+	}
+
+	return nil
+}
+
 // A WireguardKeysWriter writes the Wireguard Keys for all machines
 // under its scope for the specified ring
 type WireguardKeysWriter interface {
diff --git a/pkg/zones/zone_rings.go b/pkg/zones/zone_rings.go
deleted file mode 100644
index bd3b5c9..0000000
--- a/pkg/zones/zone_rings.go
+++ /dev/null
@@ -1,27 +0,0 @@
-package zones
-
-// PruneWireguardConfig removes wgN.conf files of machines with
-// the corresponding ring disabled.
-func (z *Zone) PruneWireguardConfig(ring int) error {
-	return pruneWireguardConfig(z, ring)
-}
-
-// PruneWireguardConfig removes wgN.conf files of machines with
-// the corresponding ring disabled on all zones
-func (m *Zones) PruneWireguardConfig(ring int) error {
-	return pruneWireguardConfig(m, ring)
-}
-
-func pruneWireguardConfig(m MachineIterator, ring int) error {
-	var err error
-
-	m.ForEachMachine(func(p *Machine) bool {
-		_, ok := p.getRingInfo(ring)
-		if !ok {
-			err = p.RemoveWireguardConfig(ring)
-		}
-		return err != nil
-	})
-
-	return err
-}