zones: Env: allow multiple gateways on a Zone

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

cmd/jpictl/env.go

@@ -16,10 +16,19 @@ var envCmd = &cobra.Command{
 			return err
 		}
 
-		return m.WriteEnv(os.Stdout)
+		_, err = m.Env(*envExport).WriteTo(os.Stdout)
+		return err
 	},
 }
 
+// Command Flags
+var (
+	envExport *bool
+)
+
 func init() {
 	rootCmd.AddCommand(envCmd)
+
+	envExport = envCmd.PersistentFlags().BoolP("export", "e", false,
+		"export generated variables")
 }

pkg/zones/env.go

@@ -7,21 +7,47 @@ import (
 	"strings"
 )
 
-// WriteEnv generates environment variables for shell scripts
-func (m *Zones) WriteEnv(w io.Writer) error {
+// Env is a shell environment factory for this cluster
+type Env struct {
+	ZoneIterator
+
+	export bool
+}
+
+// Env returns a shell environment factory
+func (m *Zones) Env(export bool) *Env {
+	return &Env{
+		ZoneIterator: m,
+		export:       export,
+	}
+}
+
+// Zones returns the list of Zone IDs
+func (m *Env) Zones() []int {
+	var zones []int
+
+	m.ForEachZone(func(z *Zone) bool {
+		zones = append(zones, z.ID)
+		return false
+	})
+
+	return zones
+}
+
+// WriteTo generates environment variables for shell scripts
+func (m *Env) WriteTo(w io.Writer) (int64, error) {
 	var buf bytes.Buffer
 
-	m.writeEnvVarFn(&buf, genEnvZones, "ZONES")
+	m.writeEnvVar(&buf, genEnvInts(m.Zones()), "ZONES")
 	m.ForEachZone(func(z *Zone) bool {
 		m.writeEnvZone(&buf, z)
 		return false
 	})
 
-	_, err := buf.WriteTo(w)
-	return err
+	return buf.WriteTo(w)
 }
 
-func (m *Zones) writeEnvZone(w io.Writer, z *Zone) {
+func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
 	zoneID := z.ID
 
 	// ZONE{zoneID}
@@ -31,28 +57,19 @@ func (m *Zones) writeEnvZone(w io.Writer, z *Zone) {
 	m.writeEnvVar(w, z.Name, "ZONE%v_%s", zoneID, "NAME")
 
 	// ZONE{zoneID}_GW
-	gatewayID := getRingZeroGatewayID(z)
-	if gatewayID > 0 {
-		m.writeEnvVar(w, fmt.Sprintf("%v", gatewayID), "ZONE%v_%s", zoneID, "GW")
-
-		// ZONE{zoneID}_IP
-		if ip, ok := RingZeroAddress(zoneID, gatewayID); ok {
-			m.writeEnvVar(w, ip.String(), "ZONE%v_%s", zoneID, "IP")
-		}
+	gateways := z.GatewayIDs()
+	if len(gateways) > 0 {
+		m.writeEnvVar(w, genEnvInts(gateways), "ZONE%v_%s", zoneID, "GW")
 	}
 }
 
-func (m *Zones) writeEnvVarFn(w io.Writer, fn func(*Zones) string, name string, args ...any) {
-	var value string
+func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
+	var prefix string
 
-	if fn != nil {
-		value = fn(m)
+	if m.export {
+		prefix = "export "
 	}
 
-	m.writeEnvVar(w, value, name, args...)
-}
-
-func (*Zones) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	if len(args) > 0 {
 		name = fmt.Sprintf(name, args...)
 	}
@@ -60,37 +77,29 @@ func (*Zones) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	if name != "" {
 		value = strings.TrimSpace(value)
 
-		_, _ = fmt.Fprintf(w, "%s=%q\n", name, value)
+		_, _ = fmt.Fprintf(w, "%s%s=%q\n", prefix, name, value)
 	}
 }
 
-func genEnvZones(m *Zones) string {
-	s := make([]string, 0, len(m.Zones))
-	for _, z := range m.Zones {
-		s = append(s, fmt.Sprintf("%v", z.ID))
+func genEnvInts(a []int) string {
+	s := make([]string, len(a))
+	for i, id := range a {
+		s[i] = fmt.Sprintf("%v", id)
 	}
 
 	return strings.Join(s, " ")
 }
 
 func genEnvZoneNodes(z *Zone) string {
-	s := make([]string, 0, len(z.Machines))
-	for _, p := range z.Machines {
-		s = append(s, p.Name)
+	if n := z.Len(); n > 0 {
+		s := make([]string, 0, n)
+
+		z.ForEachMachine(func(p *Machine) bool {
+			s = append(s, p.Name)
+			return false
+		})
+
+		return strings.Join(s, " ")
 	}
-	return strings.Join(s, " ")
-}
-
-func getRingZeroGatewayID(z *Zone) int {
-	var gatewayID int
-
-	z.ForEachMachine(func(p *Machine) bool {
-		if p.IsGateway() {
-			gatewayID = p.ID
-		}
-
-		return gatewayID != 0
-	})
-
-	return gatewayID
+	return ""
 }

pkg/zones/wireguard.go

@@ -45,7 +45,7 @@ func pruneWireguardConfig(m MachineIterator, ring int) error {
 	var err error
 
 	m.ForEachMachine(func(p *Machine) bool {
-		err = p.zone.PruneWireguardConfig(ring)
+		err = p.PruneWireguardConfig(ring)
 		if os.IsNotExist(err) {
 			// ignore
 			err = nil