WIP

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

pkg/zones/machine_rings.go

@@ -270,7 +270,7 @@ func (m *Machine) SyncWireguardConfig(ring int) error {
 
 // WriteWireguardConfig ...
 func (m *Machine) WriteWireguardConfig(ring int) error {
-	r, err := NewRing(m.zone.zones, m.zone, ring)
+	r, err := NewRing(m.zone, ring)
 	if err != nil {
 		return err
 	}

pkg/zones/rings.go

@@ -183,26 +183,32 @@ func RingOneAddress(zoneID, nodeID int) (netip.Addr, bool) {
 
 var (
 	_ MachineIterator = (*Ring)(nil)
-	_ ZoneIterator    = (*Ring)(nil)
 )
 
 // A Ring describes all peers on a ring
 type Ring struct {
 	RingAddressEncoder
-	ZoneIterator
 
 	Peers []*RingPeer
 }
 
 // AddPeer adds a [Machine] to the ring
 func (r *Ring) AddPeer(p *Machine) bool {
-	ri, ok := p.getRingInfo(r.ID)
+	_, ok := p.getRingInfo(r.ID)
 	if !ok {
 		return false
 	}
 
+	rp := r.newPeer(p)
+	r.Peers = append(r.Peers, rp)
+	return true
+}
+
+func (r *Ring) newPeer(p *Machine) *RingPeer {
 	nodeID := p.ID
 	zoneID := p.Zone()
+
+	ri, _ := p.getRingInfo(r.ID)
 	addr, _ := r.Encode(zoneID, nodeID)
 
 	rp := &RingPeer{
@@ -218,61 +224,20 @@ func (r *Ring) AddPeer(p *Machine) bool {
 		},
 	}
 
-	switch {
+	if r.ID == 0 {
-	case r.ID == 0:
+		rp.AllowRingOneNetwork(zoneID)
-		r.setRingZeroAllowedIPs(rp)
+		rp.AllowIP(addr)
-	case p.IsGateway():
+	} else {
-		r.setRingOneGatewayAllowedIPs(rp)
+		rp.AllowIP(addr)
-	default:
+
-		r.setRingOneNodeAllowedIPs(rp)
+		if p.IsGateway() {
+			zones := p.zone.zones
+			rp.AllowOtherRingOneNetworks(zones, zoneID)
+			rp.AllowRingZeroGateways(zones)
+		}
 	}
 
-	r.Peers = append(r.Peers, rp)
+	return rp
-	return true
-}
-
-func (r *Ring) setRingZeroAllowedIPs(rp *RingPeer) {
-	zoneID, _, _ := r.Decode(rp.Address)
-
-	// everyone on ring0 is a gateway to ring1
-	addr, _ := RingOneAddress(zoneID, 0)
-	rp.AllowCIDR(addr, 12)
-
-	// peer
-	rp.AllowCIDR(rp.Address, 32)
-}
-
-func (r *Ring) setRingOneGatewayAllowedIPs(rp *RingPeer) {
-	zoneID, _, _ := r.Decode(rp.Address)
-
-	// peer
-	rp.AllowCIDR(rp.Address, 32)
-
-	// ring1 gateways connect to all other ring1 networks
-	r.ForEachZone(func(z *Zone) bool {
-		if z.ID != zoneID {
-			addr, _ := r.Encode(z.ID, 0)
-			rp.AllowCIDR(addr, 12)
-		}
-		return false
-	})
-
-	// ring1 gateways also connect to all ring0 addresses
-	r.ForEachZone(func(z *Zone) bool {
-		z.ForEachMachine(func(p *Machine) bool {
-			if p.IsGateway() {
-				addr, _ := RingZeroAddress(z.ID, p.ID)
-				rp.AllowCIDR(addr, 32)
-			}
-			return false
-		})
-		return false
-	})
-}
-
-func (*Ring) setRingOneNodeAllowedIPs(rp *RingPeer) {
-	// only to the peer itself
-	rp.AllowCIDR(rp.Address, 32)
 }
 
 // ForEachMachine calls a function for each Machine in the ring
@@ -285,35 +250,38 @@ func (r *Ring) ForEachMachine(fn func(*Machine) bool) {
 	}
 }
 
-// ExportConfig builds a wgN.conf for the specified machine on the ring
+// ExportConfig ...
 func (r *Ring) ExportConfig(p *Machine) (*wireguard.Config, error) {
-	var found bool
+	cur, ok := r.getMachine(p)
+	if !ok {
+		return nil, fs.ErrNotExist
+	}
 
 	out := &wireguard.Config{
 		Interface: wireguard.InterfaceConfig{
+			Address:    cur.Address,
+			PrivateKey: cur.PrivateKey,
 			ListenPort: r.Port,
 		},
 	}
 
 	for _, pp := range r.Peers {
-		switch {
+		if pp.Node != p {
-		case pp.Node == p:
-			// current
-			found = true
-			out.Interface.Address = pp.Address
-			out.Interface.PrivateKey = pp.PrivateKey
-		default:
-			// peer
 			pc := pp.PeerConfig
 			out.Peer = append(out.Peer, pc)
 		}
 	}
+	return out, nil
+}
 
-	if !found {
+func (r *Ring) getMachine(p *Machine) (*RingPeer, bool) {
-		return nil, fs.ErrNotExist
+	for _, pp := range r.Peers {
+		if pp.Node == p {
+			return pp, true
+		}
 	}
 
-	return out, nil
+	return nil, false
 }
 
 // A RingPeer is a node on a [Ring]
@@ -325,17 +293,55 @@ type RingPeer struct {
 	PeerConfig wireguard.PeerConfig
 }
 
-// AllowCIDR allows an IP range via this peer
+// AllowIP ...
-func (rp *RingPeer) AllowCIDR(addr netip.Addr, bits int) {
+func (rp *RingPeer) AllowIP(addr netip.Addr) {
-	cidr := netip.PrefixFrom(addr, bits)
+	rp.PeerConfig.AllowedIPs = append(rp.PeerConfig.AllowedIPs,
-	rp.PeerConfig.AllowedIPs = append(rp.PeerConfig.AllowedIPs, cidr)
+		netip.PrefixFrom(addr, 32),
+	)
 }
 
-// NewRing composes a new Ring for Wireguard setup
+// AllowNetwork ...
-func NewRing(z ZoneIterator, m MachineIterator, ring int) (*Ring, error) {
+func (rp *RingPeer) AllowNetwork(addr netip.Addr, bits int) {
+	rp.PeerConfig.AllowedIPs = append(rp.PeerConfig.AllowedIPs,
+		netip.PrefixFrom(addr, bits),
+	)
+}
+
+// AllowRingOneNetwork ...
+func (rp *RingPeer) AllowRingOneNetwork(zoneID int) {
+	addr, _ := RingOneAddress(zoneID, 0)
+	rp.AllowNetwork(addr, 12)
+}
+
+// AllowOtherRingOneNetworks ...
+func (rp *RingPeer) AllowOtherRingOneNetworks(m *Zones, zoneID int) {
+	m.ForEachZone(func(z *Zone) bool {
+		if z.ID != zoneID {
+			addr, _ := RingOneAddress(z.ID, 0)
+			rp.AllowNetwork(addr, 12)
+		}
+		return false
+	})
+}
+
+// AllowRingZeroGateways ...
+func (rp *RingPeer) AllowRingZeroGateways(m *Zones) {
+	m.ForEachMachine(func(p *Machine) bool {
+		switch {
+		case !p.IsGateway():
+			// skip regular nodes
+		default:
+			addr, _ := RingZeroAddress(p.Zone(), p.ID)
+			rp.AllowIP(addr)
+		}
+		return false
+	})
+}
+
+// NewRing ...
+func NewRing(m MachineIterator, ring int) (*Ring, error) {
 	r := &Ring{
 		RingAddressEncoder: Rings[ring],
-		ZoneIterator:       z,
 	}
 
 	m.ForEachMachine(func(p *Machine) bool {

pkg/zones/zone_rings.go

@@ -21,9 +21,9 @@ type machineRinger interface {
 func (z *Zone) SyncWireguardConfig(ring int) error {
 	switch ring {
 	case 0:
-		return syncWireguardConfig(z.zones, z.zones, ring)
+		return syncWireguardConfig(z.zones, ring)
 	case 1:
-		return syncWireguardConfig(z.zones, z, ring)
+		return syncWireguardConfig(z, ring)
 	default:
 		return fs.ErrInvalid
 	}
@@ -45,11 +45,11 @@ func (z *Zone) WriteWireguardKeys(ring int) error {
 func (m *Zones) SyncWireguardConfig(ring int) error {
 	switch ring {
 	case 0:
-		return syncWireguardConfig(m, m, ring)
+		return syncWireguardConfig(m, ring)
 	case 1:
 		var err error
 		m.ForEachZone(func(z *Zone) bool {
-			err = syncWireguardConfig(m, z, ring)
+			err = syncWireguardConfig(z, ring)
 			return err != nil
 		})
 		return err
@@ -69,13 +69,13 @@ func (m *Zones) WriteWireguardKeys(ring int) error {
 	return writeWireguardKeys(m, ring)
 }
 
-func syncWireguardConfig(z ZoneIterator, m machineRinger, ring int) error {
+func syncWireguardConfig(m machineRinger, ring int) error {
 	err := m.PruneWireguardConfig(ring)
 	if err != nil {
 		return err
 	}
 
-	r, err := NewRing(z, m, ring)
+	r, err := NewRing(m, ring)
 	if err != nil {
 		return err
 	}

pkg/zones/zones.go

@@ -13,7 +13,6 @@ import (
 var (
 	_ MachineIterator = (*Zone)(nil)
 	_ MachineIterator = (*Zones)(nil)
-	_ ZoneIterator    = (*Zones)(nil)
 )
 
 // A MachineIterator is a set of Machines we can iterate on
@@ -21,11 +20,6 @@ type MachineIterator interface {
 	ForEachMachine(func(*Machine) bool)
 }
 
-// A ZoneIterator is a set of Zones we can iterate on
-type ZoneIterator interface {
-	ForEachZone(func(*Zone) bool)
-}
-
 // Zone represents one zone in a cluster
 type Zone struct {
 	zones *Zones