orama/pkg/gateway/handlers/wireguard/handler.go

package wireguard

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/DeBrosOfficial/network/pkg/auth"
	"github.com/DeBrosOfficial/network/pkg/rqlite"
	"go.uber.org/zap"
)

// PeerRecord represents a WireGuard peer stored in RQLite
type PeerRecord struct {
	NodeID    string `json:"node_id" db:"node_id"`
	WGIP      string `json:"wg_ip" db:"wg_ip"`
	PublicKey string `json:"public_key" db:"public_key"`
	PublicIP  string `json:"public_ip" db:"public_ip"`
	WGPort    int    `json:"wg_port" db:"wg_port"`
}

// RegisterPeerRequest is the request body for peer registration
type RegisterPeerRequest struct {
	NodeID       string `json:"node_id"`
	PublicKey    string `json:"public_key"`
	PublicIP     string `json:"public_ip"`
	WGPort       int    `json:"wg_port,omitempty"`
	ClusterSecret string `json:"cluster_secret"`
}

// RegisterPeerResponse is the response for peer registration
type RegisterPeerResponse struct {
	AssignedWGIP string       `json:"assigned_wg_ip"`
	Peers        []PeerRecord `json:"peers"`
}

// Handler handles WireGuard peer exchange endpoints
type Handler struct {
	logger        *zap.Logger
	rqliteClient  rqlite.Client
	clusterSecret string // expected cluster secret for auth
}

// NewHandler creates a new WireGuard handler
func NewHandler(logger *zap.Logger, rqliteClient rqlite.Client, clusterSecret string) *Handler {
	return &Handler{
		logger:        logger,
		rqliteClient:  rqliteClient,
		clusterSecret: clusterSecret,
	}
}

// HandleRegisterPeer handles POST /v1/internal/wg/peer
// A new node calls this to register itself and get all existing peers.
func (h *Handler) HandleRegisterPeer(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	r.Body = http.MaxBytesReader(w, r.Body, 1<<20) // 1MB
	var req RegisterPeerRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid request body", http.StatusBadRequest)
		return
	}

	// Validate cluster secret
	if h.clusterSecret != "" && req.ClusterSecret != h.clusterSecret {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	if req.NodeID == "" || req.PublicKey == "" || req.PublicIP == "" {
		http.Error(w, "node_id, public_key, and public_ip are required", http.StatusBadRequest)
		return
	}

	if req.WGPort == 0 {
		req.WGPort = 51820
	}

	ctx := r.Context()

	// Assign next available WG IP
	wgIP, err := h.assignNextWGIP(ctx)
	if err != nil {
		h.logger.Error("failed to assign WG IP", zap.Error(err))
		http.Error(w, "failed to assign WG IP", http.StatusInternalServerError)
		return
	}

	// Insert peer record
	_, err = h.rqliteClient.Exec(ctx,
		"INSERT OR REPLACE INTO wireguard_peers (node_id, wg_ip, public_key, public_ip, wg_port) VALUES (?, ?, ?, ?, ?)",
		req.NodeID, wgIP, req.PublicKey, req.PublicIP, req.WGPort)
	if err != nil {
		h.logger.Error("failed to insert WG peer", zap.Error(err))
		http.Error(w, "failed to register peer", http.StatusInternalServerError)
		return
	}

	// Get all peers (including the one just added)
	peers, err := h.ListPeers(ctx)
	if err != nil {
		h.logger.Error("failed to list WG peers", zap.Error(err))
		http.Error(w, "failed to list peers", http.StatusInternalServerError)
		return
	}

	resp := RegisterPeerResponse{
		AssignedWGIP: wgIP,
		Peers:        peers,
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(resp)

	h.logger.Info("registered WireGuard peer",
		zap.String("node_id", req.NodeID),
		zap.String("wg_ip", wgIP),
		zap.String("public_ip", req.PublicIP))
}

// HandleListPeers handles GET /v1/internal/wg/peers
func (h *Handler) HandleListPeers(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	if !h.validateInternalRequest(r) {
		http.Error(w, "unauthorized", http.StatusForbidden)
		return
	}

	peers, err := h.ListPeers(r.Context())
	if err != nil {
		h.logger.Error("failed to list WG peers", zap.Error(err))
		http.Error(w, "failed to list peers", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(peers)
}

// HandleRemovePeer handles DELETE /v1/internal/wg/peer?node_id=xxx
func (h *Handler) HandleRemovePeer(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodDelete {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	if !h.validateInternalRequest(r) {
		http.Error(w, "unauthorized", http.StatusForbidden)
		return
	}

	nodeID := r.URL.Query().Get("node_id")
	if nodeID == "" {
		http.Error(w, "node_id parameter required", http.StatusBadRequest)
		return
	}

	_, err := h.rqliteClient.Exec(r.Context(),
		"DELETE FROM wireguard_peers WHERE node_id = ?", nodeID)
	if err != nil {
		h.logger.Error("failed to remove WG peer", zap.Error(err))
		http.Error(w, "failed to remove peer", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
	h.logger.Info("removed WireGuard peer", zap.String("node_id", nodeID))
}

// validateInternalRequest checks that the request comes from a WireGuard peer
// and includes a valid cluster secret. Both conditions must be met.
func (h *Handler) validateInternalRequest(r *http.Request) bool {
	if !auth.IsWireGuardPeer(r.RemoteAddr) {
		return false
	}
	if h.clusterSecret == "" {
		return true
	}
	return r.Header.Get("X-Cluster-Secret") == h.clusterSecret
}

// ListPeers returns all registered WireGuard peers
func (h *Handler) ListPeers(ctx context.Context) ([]PeerRecord, error) {
	var peers []PeerRecord
	err := h.rqliteClient.Query(ctx, &peers,
		"SELECT node_id, wg_ip, public_key, public_ip, wg_port FROM wireguard_peers ORDER BY wg_ip")
	if err != nil {
		return nil, fmt.Errorf("failed to query wireguard_peers: %w", err)
	}
	return peers, nil
}

// assignNextWGIP finds the next available 10.0.0.x IP by querying all peers
// and finding the numerically highest IP. Avoids lexicographic MAX() issues.
func (h *Handler) assignNextWGIP(ctx context.Context) (string, error) {
	var rows []struct {
		WGIP string `db:"wg_ip"`
	}

	err := h.rqliteClient.Query(ctx, &rows, "SELECT wg_ip FROM wireguard_peers")
	if err != nil {
		return "", fmt.Errorf("failed to query WG IPs: %w", err)
	}

	if len(rows) == 0 {
		return "10.0.0.1", nil
	}

	maxA, maxB, maxC, maxD := 0, 0, 0, 0
	for _, row := range rows {
		var a, b, c, d int
		if _, err := fmt.Sscanf(row.WGIP, "%d.%d.%d.%d", &a, &b, &c, &d); err != nil {
			continue
		}
		if c > maxC || (c == maxC && d > maxD) {
			maxA, maxB, maxC, maxD = a, b, c, d
		}
	}

	if maxA == 0 {
		return "10.0.0.1", nil
	}

	maxD++
	if maxD > 254 {
		maxC++
		maxD = 1
		if maxC > 255 {
			return "", fmt.Errorf("WireGuard IP space exhausted")
		}
	}

	return fmt.Sprintf("%d.%d.%d.%d", maxA, maxB, maxC, maxD), nil
}