felhom.eu/hub/internal/store/wg.go

package store

// S1 offsite connectivity (doc 06 §3.2): the WG endpoint record + the per-peer /32 allocator.
// The hub DB is the source of truth for the endpoint's peer list; internal/wgsync pushes the
// FULL list to the endpoint (declarative — the endpoint converges, drift is erased). These
// methods never touch SSH; they are pure persistence + allocation.

import (
	"database/sql"
	"errors"
	"fmt"
	"net/netip"
	"strings"
)

// ErrWGEndpointUnset is returned by AddWGPeer when no endpoint record exists yet — allocation
// needs the tunnel subnet, so peers cannot be registered before PUT /admin/wg/endpoint.
var ErrWGEndpointUnset = errors.New("wg endpoint not configured")

// ErrWGSubnetExhausted is returned when the tunnel subnet has no free host address left.
var ErrWGSubnetExhausted = errors.New("tunnel subnet exhausted")

// WGEndpoint is the public endpoint's coordinates (doc 06 §3.2 "endpoint record"). One row
// expected ("ep0") for now; the schema allows more for a later multi-endpoint world.
type WGEndpoint struct {
	EndpointID   string
	DNSName      string
	WGPort       int
	ServerPubkey string
	TunnelSubnet string // CIDR, e.g. "10.77.0.0/24"
	PBSTunnelIP  string // the endpoint's own in-tunnel address, e.g. "10.77.0.1"
}

// WGPeer is one registered peer (doc 06 §3.2 "per-host peer entry"). HostID is inert until the
// S2 host-join; presence of the row IS the desired state (no status column by design).
type WGPeer struct {
	Pubkey     string
	AssignedIP string // bare IP — the /32 suffix is presentation, not storage
	HostID     string
	Note       string
}

// SetWGEndpoint upserts the (single expected) endpoint record.
func (s *Store) SetWGEndpoint(e *WGEndpoint) error {
	if e.EndpointID == "" {
		e.EndpointID = "ep0"
	}
	_, err := s.db.Exec(`
		INSERT INTO wg_endpoints (endpoint_id, dns_name, wg_port, server_pubkey, tunnel_subnet, pbs_tunnel_ip)
		VALUES (?, ?, ?, ?, ?, ?)
		ON CONFLICT(endpoint_id) DO UPDATE SET
			dns_name = excluded.dns_name, wg_port = excluded.wg_port,
			server_pubkey = excluded.server_pubkey, tunnel_subnet = excluded.tunnel_subnet,
			pbs_tunnel_ip = excluded.pbs_tunnel_ip, updated_at = datetime('now')`,
		e.EndpointID, e.DNSName, e.WGPort, e.ServerPubkey, e.TunnelSubnet, e.PBSTunnelIP)
	return err
}

// GetWGEndpoint returns the endpoint record, or sql.ErrNoRows when unset. With multiple rows
// (not expected in S1) the lowest endpoint_id wins, deterministically.
func (s *Store) GetWGEndpoint() (*WGEndpoint, error) {
	var e WGEndpoint
	err := s.db.QueryRow(`
		SELECT endpoint_id, dns_name, wg_port, server_pubkey, tunnel_subnet, pbs_tunnel_ip
		FROM wg_endpoints ORDER BY endpoint_id LIMIT 1`).
		Scan(&e.EndpointID, &e.DNSName, &e.WGPort, &e.ServerPubkey, &e.TunnelSubnet, &e.PBSTunnelIP)
	if err != nil {
		return nil, err
	}
	return &e, nil
}

// AddWGPeer registers a peer pubkey and allocates the lowest free host address in the tunnel
// subnet — skipping the network address, the endpoint's own pbs_tunnel_ip, and (v4) the
// broadcast. Idempotent: an already-registered pubkey returns its existing IP with existed=true
// and allocates nothing. One transaction; the UNIQUE(assigned_ip) constraint backstops a
// concurrent-allocation race (the loser retries once internally).
func (s *Store) AddWGPeer(pubkey, hostID, note string) (assignedIP string, existed bool, err error) {
	for attempt := 0; attempt < 2; attempt++ {
		assignedIP, existed, err = s.addWGPeerOnce(pubkey, hostID, note)
		if err != nil && attempt == 0 && strings.Contains(err.Error(), "UNIQUE constraint failed: wg_peers.assigned_ip") {
			continue // race with a concurrent insert — re-read the assigned set and retry once
		}
		return assignedIP, existed, err
	}
	return assignedIP, existed, err
}

func (s *Store) addWGPeerOnce(pubkey, hostID, note string) (string, bool, error) {
	tx, err := s.db.Begin()
	if err != nil {
		return "", false, err
	}
	defer tx.Rollback()

	// Idempotency: existing pubkey → its IP, no allocation.
	var existingIP string
	err = tx.QueryRow(`SELECT assigned_ip FROM wg_peers WHERE pubkey = ?`, pubkey).Scan(&existingIP)
	if err == nil {
		return existingIP, true, nil
	}
	if err != sql.ErrNoRows {
		return "", false, err
	}

	var subnet, pbsIP string
	err = tx.QueryRow(`SELECT tunnel_subnet, pbs_tunnel_ip FROM wg_endpoints ORDER BY endpoint_id LIMIT 1`).
		Scan(&subnet, &pbsIP)
	if err == sql.ErrNoRows {
		return "", false, ErrWGEndpointUnset
	}
	if err != nil {
		return "", false, err
	}

	taken := map[string]bool{}
	rows, err := tx.Query(`SELECT assigned_ip FROM wg_peers`)
	if err != nil {
		return "", false, err
	}
	for rows.Next() {
		var ip string
		if err := rows.Scan(&ip); err != nil {
			rows.Close()
			return "", false, err
		}
		taken[ip] = true
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return "", false, err
	}

	ip, err := lowestFreeHost(subnet, pbsIP, taken)
	if err != nil {
		return "", false, err
	}

	if _, err := tx.Exec(`INSERT INTO wg_peers (pubkey, assigned_ip, host_id, note) VALUES (?, ?, ?, ?)`,
		pubkey, ip, hostID, note); err != nil {
		return "", false, err
	}
	if err := tx.Commit(); err != nil {
		return "", false, err
	}
	return ip, false, nil
}

// lowestFreeHost walks the subnet's host addresses from lowest and returns the first not in
// `taken`, skipping the network address, the endpoint's own address, and (v4) the broadcast.
func lowestFreeHost(subnet, endpointIP string, taken map[string]bool) (string, error) {
	prefix, err := netip.ParsePrefix(subnet)
	if err != nil {
		return "", fmt.Errorf("bad tunnel_subnet %q: %w", subnet, err)
	}
	prefix = prefix.Masked()
	last := lastAddr(prefix)
	for a := prefix.Addr().Next(); prefix.Contains(a); a = a.Next() {
		if a.Is4() && a == last {
			break // v4 broadcast — never allocated
		}
		ip := a.String()
		if ip == endpointIP || taken[ip] {
			continue
		}
		return ip, nil
	}
	return "", ErrWGSubnetExhausted
}

// lastAddr returns the highest address in the prefix (the v4 broadcast).
func lastAddr(p netip.Prefix) netip.Addr {
	b := p.Addr().AsSlice()
	bits := p.Bits()
	for i := range b {
		hostBits := 8*(i+1) - bits
		if hostBits <= 0 {
			continue
		}
		if hostBits > 8 {
			hostBits = 8
		}
		b[i] |= byte(0xFF >> (8 - hostBits)) //nolint:gosec // hostBits is 1..8 here
	}
	a, _ := netip.AddrFromSlice(b)
	return a
}

// RemoveWGPeer deletes a peer registration. sql.ErrNoRows when the pubkey is unknown.
func (s *Store) RemoveWGPeer(pubkey string) error {
	res, err := s.db.Exec(`DELETE FROM wg_peers WHERE pubkey = ?`, pubkey)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return sql.ErrNoRows
	}
	return nil
}

// ListWGPeers returns all registered peers ordered by assigned_ip — a deterministic order so
// the sync payload bytes are stable for identical registries.
func (s *Store) ListWGPeers() ([]WGPeer, error) {
	rows, err := s.db.Query(`SELECT pubkey, assigned_ip, host_id, note FROM wg_peers ORDER BY assigned_ip`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var peers []WGPeer
	for rows.Next() {
		var p WGPeer
		if err := rows.Scan(&p.Pubkey, &p.AssignedIP, &p.HostID, &p.Note); err != nil {
			return nil, err
		}
		peers = append(peers, p)
	}
	return peers, rows.Err()
}