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, ring)
+	r, err := NewRing(m.zone.zones, m.zone, ring)
 	if err != nil {
 		return err
 	}

pkg/zones/rings.go

@@ -183,32 +183,26 @@ 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 {
-	_, ok := p.getRingInfo(r.ID)
+	ri, 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{
@@ -224,20 +218,61 @@ func (r *Ring) newPeer(p *Machine) *RingPeer {
 		},
 	}
 
-	if r.ID == 0 {
-		rp.AllowRingOneNetwork(zoneID)
-		rp.AllowIP(addr)
-	} else {
-		rp.AllowIP(addr)
-
-		if p.IsGateway() {
-			zones := p.zone.zones
-			rp.AllowOtherRingOneNetworks(zones, zoneID)
-			rp.AllowRingZeroGateways(zones)
-		}
+	switch {
+	case r.ID == 0:
+		r.setRingZeroAllowedIPs(rp)
+	case p.IsGateway():
+		r.setRingOneGatewayAllowedIPs(rp)
+	default:
+		r.setRingOneNodeAllowedIPs(rp)
 	}
 
-	return rp
+	r.Peers = append(r.Peers, 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
@@ -250,38 +285,35 @@ func (r *Ring) ForEachMachine(fn func(*Machine) bool) {
 	}
 }
 
-// ExportConfig ...
+// ExportConfig builds a wgN.conf for the specified machine on the ring
 func (r *Ring) ExportConfig(p *Machine) (*wireguard.Config, error) {
-	cur, ok := r.getMachine(p)
-	if !ok {
-		return nil, fs.ErrNotExist
-	}
+	var found bool
 
 	out := &wireguard.Config{
 		Interface: wireguard.InterfaceConfig{
-			Address:    cur.Address,
-			PrivateKey: cur.PrivateKey,
 			ListenPort: r.Port,
 		},
 	}
 
 	for _, pp := range r.Peers {
-		if pp.Node != p {
+		switch {
+		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
-}
 
-func (r *Ring) getMachine(p *Machine) (*RingPeer, bool) {
-	for _, pp := range r.Peers {
-		if pp.Node == p {
-			return pp, true
-		}
+	if !found {
+		return nil, fs.ErrNotExist
 	}
 
-	return nil, false
+	return out, nil
 }
 
 // A RingPeer is a node on a [Ring]
@@ -293,55 +325,17 @@ type RingPeer struct {
 	PeerConfig wireguard.PeerConfig
 }
 
-// AllowIP ...
-func (rp *RingPeer) AllowIP(addr netip.Addr) {
-	rp.PeerConfig.AllowedIPs = append(rp.PeerConfig.AllowedIPs,
-		netip.PrefixFrom(addr, 32),
-	)
+// AllowCIDR allows an IP range via this peer
+func (rp *RingPeer) AllowCIDR(addr netip.Addr, bits int) {
+	cidr := netip.PrefixFrom(addr, bits)
+	rp.PeerConfig.AllowedIPs = append(rp.PeerConfig.AllowedIPs, cidr)
 }
 
-// AllowNetwork ...
-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) {
+// NewRing composes a new Ring for Wireguard setup
+func NewRing(z ZoneIterator, 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, ring)
+		return syncWireguardConfig(z.zones, z.zones, ring)
 	case 1:
-		return syncWireguardConfig(z, ring)
+		return syncWireguardConfig(z.zones, 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, ring)
+		return syncWireguardConfig(m, m, ring)
 	case 1:
 		var err error
 		m.ForEachZone(func(z *Zone) bool {
-			err = syncWireguardConfig(z, ring)
+			err = syncWireguardConfig(m, z, ring)
 			return err != nil
 		})
 		return err
@@ -69,13 +69,13 @@ func (m *Zones) WriteWireguardKeys(ring int) error {
 	return writeWireguardKeys(m, ring)
 }
 
-func syncWireguardConfig(m machineRinger, ring int) error {
+func syncWireguardConfig(z ZoneIterator, m machineRinger, ring int) error {
 	err := m.PruneWireguardConfig(ring)
 	if err != nil {
 		return err
 	}
 
-	r, err := NewRing(m, ring)
+	r, err := NewRing(z, m, ring)
 	if err != nil {
 		return err
 	}

pkg/zones/zones.go

@@ -13,6 +13,7 @@ import (
 var (
 	_ MachineIterator = (*Zone)(nil)
 	_ MachineIterator = (*Zones)(nil)
+	_ ZoneIterator    = (*Zones)(nil)
 )
 
 // A MachineIterator is a set of Machines we can iterate on
@@ -20,6 +21,11 @@ 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