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
This commit is contained in:
@ -66,8 +66,6 @@ func main() {
|
||||
}
|
||||
slog.Info("connected to redis")
|
||||
|
||||
_ = rdb // TODO: pass to services that need it (host registration)
|
||||
|
||||
// Connect RPC client for the host agent.
|
||||
agentHTTP := &http.Client{Timeout: 10 * time.Minute}
|
||||
agentClient := hostagentv1connect.NewHostAgentServiceClient(
|
||||
@ -89,7 +87,7 @@ func main() {
|
||||
}
|
||||
|
||||
// API server.
|
||||
srv := api.New(queries, agentClient, pool, []byte(cfg.JWTSecret), oauthRegistry, cfg.OAuthRedirectURL)
|
||||
srv := api.New(queries, agentClient, pool, rdb, []byte(cfg.JWTSecret), oauthRegistry, cfg.OAuthRedirectURL)
|
||||
|
||||
// Start reconciler.
|
||||
reconciler := api.NewReconciler(queries, agentClient, "default", 5*time.Second)
|
||||
|
||||
@ -2,6 +2,7 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
@ -16,6 +17,10 @@ import (
|
||||
)
|
||||
|
||||
func main() {
|
||||
registrationToken := flag.String("register", "", "One-time registration token from the control plane")
|
||||
advertiseAddr := flag.String("address", "", "Externally-reachable address (ip:port) for this host agent")
|
||||
flag.Parse()
|
||||
|
||||
slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
|
||||
Level: slog.LevelDebug,
|
||||
})))
|
||||
@ -38,6 +43,8 @@ func main() {
|
||||
imagesPath := envOrDefault("AGENT_IMAGES_PATH", "/var/lib/wrenn/images")
|
||||
sandboxesPath := envOrDefault("AGENT_SANDBOXES_PATH", "/var/lib/wrenn/sandboxes")
|
||||
snapshotsPath := envOrDefault("AGENT_SNAPSHOTS_PATH", "/var/lib/wrenn/snapshots")
|
||||
cpURL := os.Getenv("AGENT_CP_URL")
|
||||
tokenFile := envOrDefault("AGENT_TOKEN_FILE", "/var/lib/wrenn/host-token")
|
||||
|
||||
cfg := sandbox.Config{
|
||||
KernelPath: kernelPath,
|
||||
@ -53,6 +60,34 @@ func main() {
|
||||
|
||||
mgr.StartTTLReaper(ctx)
|
||||
|
||||
if *advertiseAddr == "" {
|
||||
slog.Error("--address flag is required (externally-reachable ip:port)")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Register with the control plane (if configured).
|
||||
if cpURL != "" {
|
||||
hostToken, err := hostagent.Register(ctx, hostagent.RegistrationConfig{
|
||||
CPURL: cpURL,
|
||||
RegistrationToken: *registrationToken,
|
||||
TokenFile: tokenFile,
|
||||
Address: *advertiseAddr,
|
||||
})
|
||||
if err != nil {
|
||||
slog.Error("host registration failed", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
hostID, err := hostagent.HostIDFromToken(hostToken)
|
||||
if err != nil {
|
||||
slog.Error("failed to extract host ID from token", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
slog.Info("host registered", "host_id", hostID)
|
||||
hostagent.StartHeartbeat(ctx, cpURL, hostID, hostToken, 30*time.Second)
|
||||
}
|
||||
|
||||
srv := hostagent.NewServer(mgr)
|
||||
path, handler := hostagentv1connect.NewHostAgentServiceHandler(srv)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user