jpictl/pkg/cluster/rings.go

package cluster

import (
	"fmt"
	"io/fs"
	"net/netip"
	"strconv"

	"git.jpi.io/amery/jpictl/pkg/rings"
	"git.jpi.io/amery/jpictl/pkg/wireguard"
)

const (
	// RingZeroPort is the port wireguard uses for ring0
	RingZeroPort = 51800
	// RingOnePort is the port wireguard uses for ring1
	RingOnePort = 51810
)

// WireguardInterfaceID represents the number in the `wg%v`
// interface name.
type WireguardInterfaceID uint

// AsWireguardInterfaceID returns the [WireguardInterfaceID] for
// a valid [rings.RingID].
func AsWireguardInterfaceID(ring rings.RingID) (WireguardInterfaceID, error) {
	switch ring {
	case rings.RingZeroID:
		return 0, nil
	case rings.RingOneID:
		return 1, nil
	default:
		return 0, ErrInvalidRing(ring)
	}
}

// MustWireguardInterfaceID returns the [WireguardInterfaceID] for
// a valid [rings.RingID], and panics if it's not.
func MustWireguardInterfaceID(ring rings.RingID) WireguardInterfaceID {
	id, err := AsWireguardInterfaceID(ring)
	if err != nil {
		panic(err)
	}
	return id
}

// RingID tells the [rings.RingID] of the [WireguardInterfaceID].
func (wi WireguardInterfaceID) RingID() rings.RingID {
	return rings.RingID(wi + 1)
}

// PubFile returns "wgN.pub"
func (wi WireguardInterfaceID) PubFile() string {
	return fmt.Sprintf("wg%v.pub", wi)
}

// KeyFile returns "wgN.key"
func (wi WireguardInterfaceID) KeyFile() string {
	return fmt.Sprintf("wg%v.key", wi)
}

// ConfFile returns "wgN.conf"
func (wi WireguardInterfaceID) ConfFile() string {
	return fmt.Sprintf("wg%v.conf", wi)
}

// Files returns all wgN.ext file names.
func (wi WireguardInterfaceID) Files() (keyFile, pubFile, confFile string) {
	prefix := "wg" + strconv.Itoa(int(wi))

	return prefix + ".key", prefix + ".pub", prefix + ".conf"
}

// RingInfo contains represents the Wireguard endpoint details
// for a Machine on a particular ring
type RingInfo struct {
	Ring    WireguardInterfaceID
	Enabled bool
	Keys    wireguard.KeyPair
}

// RingID returns the [rings.RingID] for this [RingInfo].
func (ri *RingInfo) RingID() rings.RingID {
	return rings.RingID(ri.Ring + 1)
}

// Merge attempts to combine two RingInfo structs
func (ri *RingInfo) Merge(alter *RingInfo) error {
	switch {
	case alter == nil:
		return nil
	case ri.Ring != alter.Ring:
		// different ring
		return fmt.Errorf("invalid %s: %v ≠ %v", "ring", ri.Ring, alter.Ring)
	case ri.Enabled && !alter.Enabled:
		// can't disable via Merge
		return fmt.Errorf("invalid %s: %v → %v", "enabled", ri.Enabled, alter.Enabled)
	case !canMergeKeyPairs(ri.Keys, alter.Keys):
		// incompatible key pairs
		return fmt.Errorf("invalid %s: %s ≠ %s", "keys", ri.Keys, alter.Keys)
	}

	return ri.unsafeMerge(alter)
}

func (ri *RingInfo) unsafeMerge(alter *RingInfo) error {
	// enable via Merge
	if alter.Enabled {
		ri.Enabled = true
	}

	// fill the gaps on our key pair
	if ri.Keys.PrivateKey.IsZero() {
		ri.Keys.PrivateKey = alter.Keys.PrivateKey
	}
	if ri.Keys.PublicKey.IsZero() {
		ri.Keys.PublicKey = alter.Keys.PublicKey
	}

	return nil
}

func canMergeKeyPairs(p1, p2 wireguard.KeyPair) bool {
	switch {
	case !p1.PrivateKey.IsZero() && !p2.PrivateKey.IsZero() && !p1.PrivateKey.Equal(p2.PrivateKey):
		return false
	case !p1.PublicKey.IsZero() && !p2.PublicKey.IsZero() && !p1.PublicKey.Equal(p2.PublicKey):
		return false
	default:
		return true
	}
}

// RingAddressEncoder provides encoder/decoder access for a particular
// Wireguard ring
type RingAddressEncoder struct {
	ID     rings.RingID
	Port   uint16
	Encode func(zoneID rings.ZoneID, nodeID rings.NodeID) (netip.Addr, bool)
	Decode func(addr netip.Addr) (zoneID rings.ZoneID, nodeID rings.NodeID, ok bool)
}

var (
	// RingZero is a wg0 address encoder/decoder
	RingZero = RingAddressEncoder{
		ID:     rings.RingZeroID,
		Port:   RingZeroPort,
		Decode: ParseRingZeroAddress,
		Encode: RingZeroAddress,
	}
	// RingOne is a wg1 address encoder/decoder
	RingOne = RingAddressEncoder{
		ID:     rings.RingOneID,
		Port:   RingOnePort,
		Decode: ParseRingOneAddress,
		Encode: RingOneAddress,
	}
	// Rings provides indexed access to the ring address encoders
	Rings = []RingAddressEncoder{
		RingZero,
		RingOne,
	}
)

// ValidZoneID checks if the given zoneID is a valid 4 bit zone number.
//
// 0 is reserved, and only allowed when composing CIDRs.
func ValidZoneID(zoneID rings.ZoneID) bool {
	return zoneID == 0 || zoneID.Valid()
}

// ValidNodeID checks if the given nodeID is a valid 8 bit number.
// nodeID is unique within a Zone.
// 0 is reserved, and only allowed when composing CIDRs.
func ValidNodeID(nodeID rings.NodeID) bool {
	return nodeID == 0 || nodeID.Valid()
}

// ParseRingZeroAddress extracts zone and node ID from a wg0 [netip.Addr]
// wg0 addresses are of the form `10.0.{{zoneID}}.{{nodeID}}`
func ParseRingZeroAddress(addr netip.Addr) (zoneID rings.ZoneID, nodeID rings.NodeID, ok bool) {
	if addr.IsValid() {
		a4 := addr.As4()

		if a4[0] == 10 && a4[1] == 0 {
			zoneID = rings.ZoneID(a4[2])
			nodeID = rings.NodeID(a4[3])
			return zoneID, nodeID, true
		}
	}
	return 0, 0, false
}

// RingZeroAddress returns a wg0 IP address
func RingZeroAddress(zoneID rings.ZoneID, nodeID rings.NodeID) (netip.Addr, bool) {
	switch {
	case !ValidZoneID(zoneID) || !ValidNodeID(nodeID):
		return netip.Addr{}, false
	default:
		a4 := [4]uint8{10, 0, uint8(zoneID), uint8(nodeID)}
		return netip.AddrFrom4(a4), true
	}
}

// ParseRingOneAddress extracts zone and node ID from a wg1 [netip.Addr]
// wg1 addresses are of the form `10.{{zoneID << 4}}.{{nodeID}}`
func ParseRingOneAddress(addr netip.Addr) (zoneID rings.ZoneID, nodeID rings.NodeID, ok bool) {
	if addr.IsValid() {
		a4 := addr.As4()

		if a4[0] == 10 && a4[2] == 0 {
			zoneID = rings.ZoneID(a4[1] >> 4)
			nodeID = rings.NodeID(a4[3])
			return zoneID, nodeID, true
		}
	}
	return 0, 0, false
}

// RingOneAddress returns a wg1 IP address
func RingOneAddress(zoneID rings.ZoneID, nodeID rings.NodeID) (netip.Addr, bool) {
	switch {
	case !ValidZoneID(zoneID) || !ValidNodeID(nodeID):
		return netip.Addr{}, false
	default:
		a4 := [4]uint8{10, uint8(zoneID << 4), 0, uint8(nodeID)}
		return netip.AddrFrom4(a4), true
	}
}

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)
	if !ok {
		return false
	}

	nodeID := p.ID
	zoneID := p.Zone()
	addr, _ := r.Encode(zoneID, nodeID)

	rp := &RingPeer{
		Node:       p,
		Address:    addr,
		PrivateKey: ri.Keys.PrivateKey,
		PeerConfig: wireguard.PeerConfig{
			Name:      fmt.Sprintf("%s-%v", p.Name, ri.Ring),
			PublicKey: ri.Keys.PublicKey,
			Endpoint: wireguard.EndpointAddress{
				Host: p.FullName(),
				Port: r.Port,
			},
		},
	}

	switch {
	case r.ID == 0:
		r.setRingZeroAllowedIPs(rp)
	case p.IsGateway():
		r.setRingOneGatewayAllowedIPs(rp)
	default:
		r.setRingOneNodeAllowedIPs(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
// until instructed to terminate the loop
func (r *Ring) ForEachMachine(fn func(*Machine) bool) {
	for _, pp := range r.Peers {
		if fn(pp.Node) {
			return
		}
	}
}

// ExportConfig builds a wgN.conf for the specified machine on the ring
func (r *Ring) ExportConfig(p *Machine) (*wireguard.Config, error) {
	var found bool

	out := &wireguard.Config{
		Interface: wireguard.InterfaceConfig{
			ListenPort: r.Port,
		},
	}

	for _, pp := range r.Peers {
		switch {
		case pp.Node == p:
			// current
			found = true
			out.Interface.Name = pp.PeerConfig.Name
			out.Interface.Address = pp.Address
			out.Interface.PrivateKey = pp.PrivateKey
		default:
			// peer
			pc := pp.PeerConfig
			out.Peer = append(out.Peer, pc)
		}
	}

	if !found {
		return nil, fs.ErrNotExist
	}

	return out, nil
}

// A RingPeer is a node on a [Ring]
type RingPeer struct {
	Node *Machine

	Address    netip.Addr
	PrivateKey wireguard.PrivateKey
	PeerConfig wireguard.PeerConfig
}

// 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)
}

// NewRing composes a new Ring for Wireguard setup
func NewRing(z ZoneIterator, m MachineIterator, ringID rings.RingID) (*Ring, error) {
	var r *Ring
	for _, ring := range Rings {
		if ringID == ring.ID {
			r = &Ring{
				RingAddressEncoder: ring,
				ZoneIterator:       z,
			}
			break
		}
	}

	if r == nil {
		return nil, ErrInvalidRing(ringID)
	}

	m.ForEachMachine(func(p *Machine) bool {
		r.AddPeer(p)
		return false
	})

	return r, nil
}