Files
pptx704 2c66959b92 Add host registration, heartbeat, and multi-host management
Implements the full host ↔ control plane connection flow:

- Host CRUD endpoints (POST/GET/DELETE /v1/hosts) with role-based access:
  regular hosts admin-only, BYOC hosts for admins and team owners
- One-time registration token flow: admin creates host → gets token (1hr TTL
  in Redis + Postgres audit trail) → host agent registers with specs → gets
  long-lived JWT (1yr)
- Host agent registration client with automatic spec detection (arch, CPU,
  memory, disk) and token persistence to disk
- Periodic heartbeat (30s) via POST /v1/hosts/{id}/heartbeat with X-Host-Token
  auth and host ID cross-check
- Token regeneration endpoint (POST /v1/hosts/{id}/token) for retry after
  failed registration
- Tag management (add/remove/list) with team-scoped access control
- Host JWT with typ:"host" claim, cross-use prevention in both VerifyJWT and
  VerifyHostJWT
- requireHostToken middleware for host agent authentication
- DB-level race protection: RegisterHost uses AND status='pending' with
  rows-affected check; Redis GetDel for atomic token consume
- Migration for future mTLS support (cert_fingerprint, mtls_enabled columns)
- Host agent flags: --register (one-time token), --address (required ip:port)
- serviceErrToHTTP extended with "forbidden" → 403 mapping
- OpenAPI spec, .env.example, and README updated
2026-03-17 05:51:28 +06:00

125 lines
3.3 KiB
Go

package main
import (
"context"
"log/slog"
"net/http"
"os"
"os/signal"
"strings"
"syscall"
"time"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/redis/go-redis/v9"
"git.omukk.dev/wrenn/sandbox/internal/api"
"git.omukk.dev/wrenn/sandbox/internal/auth/oauth"
"git.omukk.dev/wrenn/sandbox/internal/config"
"git.omukk.dev/wrenn/sandbox/internal/db"
"git.omukk.dev/wrenn/sandbox/proto/hostagent/gen/hostagentv1connect"
)
func main() {
slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
Level: slog.LevelDebug,
})))
cfg := config.Load()
if len(cfg.JWTSecret) < 32 {
slog.Error("JWT_SECRET must be at least 32 characters")
os.Exit(1)
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
// Database connection pool.
pool, err := pgxpool.New(ctx, cfg.DatabaseURL)
if err != nil {
slog.Error("failed to connect to database", "error", err)
os.Exit(1)
}
defer pool.Close()
if err := pool.Ping(ctx); err != nil {
slog.Error("failed to ping database", "error", err)
os.Exit(1)
}
slog.Info("connected to database")
queries := db.New(pool)
// Redis client.
redisOpts, err := redis.ParseURL(cfg.RedisURL)
if err != nil {
slog.Error("failed to parse REDIS_URL", "error", err)
os.Exit(1)
}
rdb := redis.NewClient(redisOpts)
defer rdb.Close()
if err := rdb.Ping(ctx).Err(); err != nil {
slog.Error("failed to ping redis", "error", err)
os.Exit(1)
}
slog.Info("connected to redis")
// Connect RPC client for the host agent.
agentHTTP := &http.Client{Timeout: 10 * time.Minute}
agentClient := hostagentv1connect.NewHostAgentServiceClient(
agentHTTP,
cfg.HostAgentAddr,
)
// OAuth provider registry.
oauthRegistry := oauth.NewRegistry()
if cfg.OAuthGitHubClientID != "" && cfg.OAuthGitHubClientSecret != "" {
if cfg.CPPublicURL == "" {
slog.Error("CP_PUBLIC_URL must be set when OAuth providers are configured")
os.Exit(1)
}
callbackURL := strings.TrimRight(cfg.CPPublicURL, "/") + "/v1/auth/oauth/github/callback"
ghProvider := oauth.NewGitHubProvider(cfg.OAuthGitHubClientID, cfg.OAuthGitHubClientSecret, callbackURL)
oauthRegistry.Register(ghProvider)
slog.Info("registered OAuth provider", "provider", "github")
}
// API server.
srv := api.New(queries, agentClient, pool, rdb, []byte(cfg.JWTSecret), oauthRegistry, cfg.OAuthRedirectURL)
// Start reconciler.
reconciler := api.NewReconciler(queries, agentClient, "default", 5*time.Second)
reconciler.Start(ctx)
httpServer := &http.Server{
Addr: cfg.ListenAddr,
Handler: srv.Handler(),
}
// Graceful shutdown on signal.
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
go func() {
sig := <-sigCh
slog.Info("received signal, shutting down", "signal", sig)
cancel()
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
defer shutdownCancel()
if err := httpServer.Shutdown(shutdownCtx); err != nil {
slog.Error("http server shutdown error", "error", err)
}
}()
slog.Info("control plane starting", "addr", cfg.ListenAddr, "agent", cfg.HostAgentAddr)
if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
slog.Error("http server error", "error", err)
os.Exit(1)
}
slog.Info("control plane stopped")
}