Merge pull request 'zones: FilterMachines() and Zone.GatewayIDs()' (#3)

Reviewed-on: #3

pkg/zones/env.go

@@ -22,23 +22,11 @@ func (m *Zones) Env(export bool) *Env {
 	}
 }
 
-// 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.writeEnvVarInts(&buf, m.Zones(), "ZONES")
+	m.writeEnvVarFn(&buf, genEnvZones, "ZONES")
 	m.ForEachZone(func(z *Zone) bool {
 		m.writeEnvZone(&buf, z)
 		return false
@@ -57,22 +45,25 @@ func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
 	m.writeEnvVar(w, z.Name, "ZONE%v_%s", zoneID, "NAME")
 
 	// ZONE{zoneID}_GW
-	gateways, _ := z.GatewayIDs()
-	m.writeEnvVarInts(w, gateways, "ZONE%v_%s", 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")
+		}
+	}
 }
 
-func (m *Env) writeEnvVarInts(w io.Writer, value []int, name string, args ...any) {
-	var s string
+func (m *Env) writeEnvVarFn(w io.Writer, fn func(*Env) string, name string, args ...any) {
+	var value string
 
-	if n := len(value); n > 0 {
-		ss := make([]string, len(value))
-		for i, v := range value {
-			ss[i] = fmt.Sprintf("%v", v)
-		}
-		s = strings.Join(ss, " ")
+	if fn != nil {
+		value = fn(m)
 	}
 
-	m.writeEnvVar(w, s, name, args...)
+	m.writeEnvVar(w, value, name, args...)
 }
 
 func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
@@ -93,16 +84,35 @@ func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	}
 }
 
-func genEnvZoneNodes(z *Zone) string {
-	if n := z.Len(); n > 0 {
-		s := make([]string, 0, n)
+func genEnvZones(m *Env) string {
+	var s []string
 
-		z.ForEachMachine(func(p *Machine) bool {
-			s = append(s, p.Name)
-			return false
-		})
+	m.ForEachZone(func(z *Zone) bool {
+		s = append(s, fmt.Sprintf("%v", z.ID))
+		return false
+	})
 
-		return strings.Join(s, " ")
-	}
-	return ""
+	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)
+	}
+	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
 }