zones: Machine.Zone()

Signed-off-by: Alejandro Mery <amery@jpi.io>

pkg/wireguard/config.go

@@ -1,17 +1,40 @@
 package wireguard
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"io"
 	"net/netip"
 	"strconv"
 	"strings"
+	"text/template"
 
 	"darvaza.org/core"
 	"gopkg.in/gcfg.v1"
 )
 
+var configTemplate = template.Must(template.New("config").Funcs(template.FuncMap{
+	"PrefixJoin": func(a []netip.Prefix, sep string) string {
+		s := make([]string, len(a))
+		for i, p := range a {
+			s[i] = p.String()
+		}
+		return strings.Join(s, sep)
+	},
+}).Parse(`[Interface]
+Address = {{.Interface.Address}}
+PrivateKey = {{.Interface.PrivateKey}}
+ListenPort = {{.Interface.ListenPort}}
+{{- range .Peer }}
+
+[Peer]
+PublicKey = {{.PublicKey}}
+Endpoint = {{.Endpoint}}
+AllowedIPs = {{ PrefixJoin .AllowedIPs ", "}}
+{{- end }}
+`))
+
 // Config represents a wgN.conf file
 type Config struct {
 	Interface InterfaceConfig
@@ -28,6 +51,17 @@ func (f *Config) Peers() int {
 	return len(f.Peer)
 }
 
+// WriteTo writes a Wireguard [Config] onto the provided [io.Writer]
+func (f *Config) WriteTo(w io.Writer) (int64, error) {
+	var buf bytes.Buffer
+
+	if err := configTemplate.Execute(&buf, f); err != nil {
+		return 0, err
+	}
+
+	return buf.WriteTo(w)
+}
+
 // InterfaceConfig represents the [Interface] section
 type InterfaceConfig struct {
 	Address    netip.Addr

pkg/zones/machine.go

@@ -43,6 +43,11 @@ func (m *Machine) IsGateway() bool {
 	return ok
 }
 
+// Zone indicates the [Zone] this machine belongs to
+func (m *Machine) Zone() int {
+	return m.zone.ID
+}
+
 func (m *Machine) getPeerByName(name string) (*Machine, bool) {
 	return m.zone.zones.GetMachineByName(name)
 }

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 writeWireguardKeys(m MachineIterator, ring int) error {
-func (z *Zone) WriteWireguardKeys(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
-}

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