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")

// ErrWGPubkeyBoundElsewhere is returned by RegisterWGPeerForHost when the pubkey is already
// registered for a different owner — a key is never silently stolen (S2 Scenario C-c2).
var ErrWGPubkeyBoundElsewhere = errors.New("wg pubkey already registered elsewhere")

// 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
	CreatedAt  string // as stored (SQLite datetime text); display-only
}

// SetWGEndpoint upserts an endpoint record. Multiple endpoints are legal since v0.47.0
// (the /offsite management UI); allocation + sync still use the LOWEST endpoint_id only
// (GetWGEndpoint) — per-endpoint allocation is a future arc.
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
}

// ListWGEndpoints returns every endpoint record ordered by endpoint_id (v0.47.0 — the
// /offsite management UI lists them all; the lowest id stays THE allocation/sync endpoint).
func (s *Store) ListWGEndpoints() ([]WGEndpoint, error) {
	rows, err := s.db.Query(`
		SELECT endpoint_id, dns_name, wg_port, server_pubkey, tunnel_subnet, pbs_tunnel_ip
		FROM wg_endpoints ORDER BY endpoint_id`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var eps []WGEndpoint
	for rows.Next() {
		var e WGEndpoint
		if err := rows.Scan(&e.EndpointID, &e.DNSName, &e.WGPort, &e.ServerPubkey, &e.TunnelSubnet, &e.PBSTunnelIP); err != nil {
			return nil, err
		}
		eps = append(eps, e)
	}
	return eps, rows.Err()
}

// DeleteWGEndpoint removes an endpoint record. Plain delete — the peers-in-subnet guard
// lives in the web handler (where the refusal message is built); sql.ErrNoRows when the
// endpoint is unknown.
func (s *Store) DeleteWGEndpoint(endpointID string) error {
	res, err := s.db.Exec(`DELETE FROM wg_endpoints WHERE endpoint_id = ?`, endpointID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return sql.ErrNoRows
	}
	return 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
	}

	ip, err := allocateWGPeerTx(tx, pubkey, hostID, note)
	if err != nil {
		return "", false, err
	}
	if err := tx.Commit(); err != nil {
		return "", false, err
	}
	return ip, false, nil
}

// allocateWGPeerTx is the tx-scoped allocation core (S2 extraction — behavior-identical to the
// S1 inline code): load the endpoint, read the taken set, pick the lowest free host address,
// INSERT. Callers own the transaction and the pubkey-existence checks.
func allocateWGPeerTx(tx *sql.Tx, pubkey, hostID, note string) (string, error) {
	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 "", ErrWGEndpointUnset
	}
	if err != nil {
		return "", err
	}

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

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

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

// RegisterWGPeerForHost binds a pubkey to a host (S2 — the box-facing registration core).
// One transaction; outcomes:
//   - host already has this pubkey → (ip, changed=false): idempotent, NO generation/sync side effects
//   - host has a DIFFERENT pubkey → re-key IN PLACE: pubkey replaced, /32 KEPT (stable tunnel
//     addressing across rotation/DR); conflict with any existing row for the new pubkey → typed error
//   - no row for host: pubkey exists bound elsewhere → ErrWGPubkeyBoundElsewhere; exists unbound
//     (an S1 admin/test row) → ADOPT it (bind, keep its ip); absent → allocate the lowest free /32
//
// The idx_wg_peers_host partial unique index backstops the one-bound-peer-per-host invariant.
func (s *Store) RegisterWGPeerForHost(hostID, pubkey string) (ip string, changed bool, err error) {
	for attempt := 0; attempt < 2; attempt++ {
		ip, changed, err = s.registerWGPeerForHostOnce(hostID, pubkey)
		if err != nil && attempt == 0 && strings.Contains(err.Error(), "UNIQUE constraint failed: wg_peers.assigned_ip") {
			continue // allocation race — retry once against the fresh taken set
		}
		return ip, changed, err
	}
	return ip, changed, err
}

func (s *Store) registerWGPeerForHostOnce(hostID, pubkey string) (string, bool, error) {
	if hostID == "" {
		return "", false, fmt.Errorf("wg: RegisterWGPeerForHost needs a host_id")
	}
	tx, err := s.db.Begin()
	if err != nil {
		return "", false, err
	}
	defer tx.Rollback()

	// Existing binding for this host?
	var curPK, curIP string
	err = tx.QueryRow(`SELECT pubkey, assigned_ip FROM wg_peers WHERE host_id = ?`, hostID).Scan(&curPK, &curIP)
	if err == nil {
		if curPK == pubkey {
			return curIP, false, nil // idempotent re-register
		}
		// Re-key in place. The new pubkey must not exist anywhere (bound OR unbound) — a key is
		// never stolen, and adopting here would change the host's ip (stable-IP rule forbids it).
		var n int
		if err := tx.QueryRow(`SELECT COUNT(*) FROM wg_peers WHERE pubkey = ?`, pubkey).Scan(&n); err != nil {
			return "", false, err
		}
		if n > 0 {
			return "", false, ErrWGPubkeyBoundElsewhere
		}
		if _, err := tx.Exec(`UPDATE wg_peers SET pubkey = ?, updated_at = datetime('now') WHERE host_id = ?`,
			pubkey, hostID); err != nil {
			return "", false, err
		}
		if err := tx.Commit(); err != nil {
			return "", false, err
		}
		return curIP, true, nil
	}
	if err != sql.ErrNoRows {
		return "", false, err
	}

	// No binding for this host — does the pubkey already exist?
	var pkHost, pkIP string
	err = tx.QueryRow(`SELECT host_id, assigned_ip FROM wg_peers WHERE pubkey = ?`, pubkey).Scan(&pkHost, &pkIP)
	if err == nil {
		if pkHost != "" {
			return "", false, ErrWGPubkeyBoundElsewhere
		}
		// Unbound S1 row → adopt: bind it to the host, keep its ip.
		if _, err := tx.Exec(`UPDATE wg_peers SET host_id = ?, updated_at = datetime('now') WHERE pubkey = ?`,
			hostID, pubkey); err != nil {
			return "", false, err
		}
		if err := tx.Commit(); err != nil {
			return "", false, err
		}
		return pkIP, true, nil
	}
	if err != sql.ErrNoRows {
		return "", false, err
	}

	// Fresh registration → allocate.
	ip, err := allocateWGPeerTx(tx, pubkey, hostID, "")
	if err != nil {
		return "", false, err
	}
	if err := tx.Commit(); err != nil {
		return "", false, err
	}
	return ip, true, nil
}

// GetWGPeerForHost returns the host's bound peer, or sql.ErrNoRows when none.
func (s *Store) GetWGPeerForHost(hostID string) (*WGPeer, error) {
	var p WGPeer
	err := s.db.QueryRow(`SELECT pubkey, assigned_ip, host_id, note FROM wg_peers WHERE host_id = ?`, hostID).
		Scan(&p.Pubkey, &p.AssignedIP, &p.HostID, &p.Note)
	if err != nil {
		return nil, err
	}
	return &p, 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, created_at 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, &p.CreatedAt); err != nil {
			return nil, err
		}
		peers = append(peers, p)
	}
	return peers, rows.Err()
}