From 07b4a22752101dbe2c3a5ce44eafcbfd7af1adf6 Mon Sep 17 00:00:00 2001
From: Alejandro Mery <amery@jpi.io>
Date: Sat, 26 Aug 2023 00:09:45 +0000
Subject: [PATCH] zones: introduce MachineIterator interface

Signed-off-by: Alejandro Mery <amery@jpi.io>
---
 pkg/zones/zone_rings.go | 52 ++++++++++++++++++-----------------------
 pkg/zones/zones.go      | 10 ++++++++
 2 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/pkg/zones/zone_rings.go b/pkg/zones/zone_rings.go
index 4bd2e3b..25db1bc 100644
--- a/pkg/zones/zone_rings.go
+++ b/pkg/zones/zone_rings.go
@@ -5,9 +5,29 @@ import "os"
 // PruneWireguardConfig removes wgN.conf files of machines with
 // the corresponding ring disabled.
 func (z *Zone) PruneWireguardConfig(ring int) error {
+	return pruneWireguardConfig(z, ring)
+}
+
+// WriteWireguardKeys rewrites all wgN.{key,pub} files on this zone
+func (z *Zone) WriteWireguardKeys(ring int) error {
+	return writeWireguardKeys(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)
+}
+
+// WriteWireguardKeys rewrites all wgN.{key,pub} files
+func (m *Zones) WriteWireguardKeys(ring int) error {
+	return writeWireguardKeys(m, ring)
+}
+
+func pruneWireguardConfig(m MachineIterator, ring int) error {
 	var err error
 
-	z.ForEachMachine(func(p *Machine) bool {
+	m.ForEachMachine(func(p *Machine) bool {
 		_, ok := p.getRingInfo(ring)
 		if !ok {
 			err = p.RemoveWireguardConfig(ring)
@@ -18,11 +38,10 @@ func (z *Zone) PruneWireguardConfig(ring int) error {
 	return err
 }
 
-// WriteWireguardKeys rewrites all wgN.{key,pub} files on this zone
-func (z *Zone) WriteWireguardKeys(ring int) error {
+func writeWireguardKeys(m MachineIterator, ring int) error {
 	var err error
 
-	z.ForEachMachine(func(p *Machine) bool {
+	m.ForEachMachine(func(p *Machine) bool {
 		err = p.WriteWireguardKeys(ring)
 		if os.IsNotExist(err) {
 			// ignore
@@ -34,28 +53,3 @@ func (z *Zone) WriteWireguardKeys(ring int) error {
 
 	return err
 }
-
-// PruneWireguardConfig removes wgN.conf files of machines with
-// the corresponding ring disabled on all zones
-func (m *Zones) PruneWireguardConfig(ring int) error {
-	var err error
-
-	m.ForEachZone(func(z *Zone) bool {
-		err = z.PruneWireguardConfig(ring)
-		return err != nil
-	})
-
-	return err
-}
-
-// WriteWireguardKeys rewrites all wgN.{key,pub} files
-func (m *Zones) WriteWireguardKeys(ring int) error {
-	var err error
-
-	m.ForEachZone(func(z *Zone) bool {
-		err = z.WriteWireguardKeys(ring)
-		return err != nil
-	})
-
-	return err
-}
diff --git a/pkg/zones/zones.go b/pkg/zones/zones.go
index 8fd5af1..a092d9f 100644
--- a/pkg/zones/zones.go
+++ b/pkg/zones/zones.go
@@ -10,6 +10,16 @@ import (
 	"darvaza.org/resolver"
 )
 
+var (
+	_ MachineIterator = (*Zone)(nil)
+	_ MachineIterator = (*Zones)(nil)
+)
+
+// A MachineIterator is a set of Machines we can iterate on
+type MachineIterator interface {
+	ForEachMachine(func(*Machine) bool)
+}
+
 // Zone represents one zone in a cluster
 type Zone struct {
 	zones *Zones