diff --git a/pkg/zones/env.go b/pkg/zones/env.go index d362278..a1da645 100644 --- a/pkg/zones/env.go +++ b/pkg/zones/env.go @@ -22,11 +22,23 @@ 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.writeEnvVarFn(&buf, genEnvZones, "ZONES") + m.writeEnvVarInts(&buf, m.Zones(), "ZONES") m.ForEachZone(func(z *Zone) bool { m.writeEnvZone(&buf, z) return false @@ -45,25 +57,27 @@ func (m *Env) 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() + m.writeEnvVarInts(w, gateways, "ZONE%v_%s", zoneID, "GW") } -func (m *Env) writeEnvVarFn(w io.Writer, fn func(*Env) string, name string, args ...any) { - var value string +func (m *Env) writeEnvVarInts(w io.Writer, value []int, name string, args ...any) { + var s string + + if n := len(value); n > 0 { + var buf bytes.Buffer - if fn != nil { - value = fn(m) + for i, v := range value { + if i != 0 { + _, _ = fmt.Fprint(&buf, " ") + } + _, _ = fmt.Fprintf(&buf, "%v", v) + } + + s = buf.String() } - m.writeEnvVar(w, value, name, args...) + m.writeEnvVar(w, s, name, args...) } func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) { @@ -84,35 +98,16 @@ func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) { } } -func genEnvZones(m *Env) string { - var s []string - - m.ForEachZone(func(z *Zone) bool { - s = append(s, fmt.Sprintf("%v", z.ID)) - return false - }) - - 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 - } + if n := z.Len(); n > 0 { + s := make([]string, 0, n) - return gatewayID != 0 - }) + z.ForEachMachine(func(p *Machine) bool { + s = append(s, p.Name) + return false + }) - return gatewayID + return strings.Join(s, " ") + } + return "" }