From e5639b2f4ec68ed7087914f9024fa401060f4515 Mon Sep 17 00:00:00 2001
From: Alejandro Mery <amery@jpi.io>
Date: Fri, 27 Oct 2023 19:00:06 +0000
Subject: [PATCH] cluster: generate ring keys on scan if missing

Signed-off-by: Alejandro Mery <amery@jpi.io>
---
 pkg/cluster/cluster_scan.go  |  4 ++++
 pkg/cluster/machine_rings.go | 17 +++++++++++++++++
 pkg/cluster/machine_scan.go  | 19 ++++++++++++++++++-
 pkg/cluster/rings.go         |  2 +-
 4 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/pkg/cluster/cluster_scan.go b/pkg/cluster/cluster_scan.go
index d148674..781619b 100644
--- a/pkg/cluster/cluster_scan.go
+++ b/pkg/cluster/cluster_scan.go
@@ -73,6 +73,10 @@ func (m *Cluster) scanMachines(opts *ScanOptions) error {
 		err = p.scan(opts)
 		return err != nil
 	})
+	m.ForEachMachine(func(p *Machine) bool {
+		err = p.scanWrapUp(opts)
+		return err != nil
+	})
 	return err
 }
 
diff --git a/pkg/cluster/machine_rings.go b/pkg/cluster/machine_rings.go
index a2f608b..cceae9b 100644
--- a/pkg/cluster/machine_rings.go
+++ b/pkg/cluster/machine_rings.go
@@ -240,6 +240,23 @@ func (m *Machine) applyZoneNodeID(zoneID, nodeID int) error {
 	return nil
 }
 
+func (m *Machine) setRingDefaults(ri *RingInfo) error {
+	if ri.Keys.PrivateKey.IsZero() {
+		m.info().
+			WithField("subsystem", "wireguard").
+			WithField("node", m.Name).
+			WithField("ring", ri.Ring).
+			Print("generating key pair")
+
+		kp, err := wireguard.NewKeyPair()
+		if err != nil {
+			return err
+		}
+		ri.Keys = kp
+	}
+	return nil
+}
+
 // RemoveWireguardConfig deletes wgN.conf from the machine's
 // config directory.
 func (m *Machine) RemoveWireguardConfig(ring int) error {
diff --git a/pkg/cluster/machine_scan.go b/pkg/cluster/machine_scan.go
index e282800..ef3f221 100644
--- a/pkg/cluster/machine_scan.go
+++ b/pkg/cluster/machine_scan.go
@@ -68,7 +68,8 @@ func (m *Machine) setID() error {
 	return nil
 }
 
-func (m *Machine) scan(opts *ScanOptions) error {
+// scan is called once we know about all zones and machine names
+func (m *Machine) scan(_ *ScanOptions) error {
 	for i := 0; i < RingsCount; i++ {
 		if err := m.tryApplyWireguardConfig(i); err != nil {
 			m.error(err).
@@ -80,6 +81,22 @@ func (m *Machine) scan(opts *ScanOptions) error {
 		}
 	}
 
+	return nil
+}
+
+// scanWrapUp is called once all machines have been scanned
+func (m *Machine) scanWrapUp(opts *ScanOptions) error {
+	for _, ri := range m.Rings {
+		if err := m.setRingDefaults(ri); err != nil {
+			m.error(err).
+				WithField("subsystem", "wireguard").
+				WithField("node", m.Name).
+				WithField("ring", ri.Ring).
+				Print()
+			return err
+		}
+	}
+
 	if !opts.DontResolvePublicAddresses {
 		return m.UpdatePublicAddresses()
 	}
diff --git a/pkg/cluster/rings.go b/pkg/cluster/rings.go
index 64cae1c..28301a9 100644
--- a/pkg/cluster/rings.go
+++ b/pkg/cluster/rings.go
@@ -41,7 +41,7 @@ func (ri *RingInfo) Merge(alter *RingInfo) error {
 		// can't disable via Merge
 		return fmt.Errorf("invalid %s: %v → %v", "enabled", ri.Enabled, alter.Enabled)
 	case !canMergeKeyPairs(ri.Keys, alter.Keys):
-		// incompatible keypairs
+		// incompatible key pairs
 		return fmt.Errorf("invalid %s: %s ≠ %s", "keys", ri.Keys, alter.Keys)
 	}