1
0
forked from wrenn/wrenn
Co-authored-by: Tasnim Kabir Sadik <tksadik92@gmail.com>
Reviewed-on: wrenn/sandbox#8
This commit is contained in:
2026-04-09 19:24:49 +00:00
parent 32e5a5a715
commit d3e4812e46
199 changed files with 24552 additions and 2776 deletions

View File

@ -14,10 +14,14 @@ import (
"github.com/redis/go-redis/v9"
"git.omukk.dev/wrenn/sandbox/internal/api"
"git.omukk.dev/wrenn/sandbox/internal/audit"
"git.omukk.dev/wrenn/sandbox/internal/auth"
"git.omukk.dev/wrenn/sandbox/internal/auth/oauth"
"git.omukk.dev/wrenn/sandbox/internal/channels"
"git.omukk.dev/wrenn/sandbox/internal/config"
"git.omukk.dev/wrenn/sandbox/internal/db"
"git.omukk.dev/wrenn/sandbox/proto/hostagent/gen/hostagentv1connect"
"git.omukk.dev/wrenn/sandbox/internal/lifecycle"
"git.omukk.dev/wrenn/sandbox/internal/scheduler"
)
func main() {
@ -66,12 +70,47 @@ func main() {
}
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,
)
// mTLS is mandatory — parse internal CA for CP↔agent communication.
if cfg.CACert == "" || cfg.CAKey == "" {
slog.Error("WRENN_CA_CERT and WRENN_CA_KEY are required — mTLS is mandatory for CP↔agent communication")
os.Exit(1)
}
ca, err := auth.ParseCA(cfg.CACert, cfg.CAKey)
if err != nil {
slog.Error("failed to parse mTLS CA from environment", "error", err)
os.Exit(1)
}
slog.Info("mTLS enabled: CA loaded")
// Host client pool — manages Connect RPC clients to host agents.
cpCertStore, err := auth.NewCPCertStore(ca)
if err != nil {
slog.Error("failed to issue CP client certificate", "error", err)
os.Exit(1)
}
// Renew the CP client certificate periodically so it never expires
// while the control plane is running (TTL = 24h, renewal = every 12h).
go func() {
ticker := time.NewTicker(auth.CPCertRenewInterval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
if err := cpCertStore.Refresh(); err != nil {
slog.Error("failed to renew CP client certificate", "error", err)
} else {
slog.Info("CP client certificate renewed")
}
}
}
}()
hostPool := lifecycle.NewHostClientPoolTLS(auth.CPClientTLSConfig(ca, cpCertStore))
slog.Info("host client pool: mTLS enabled")
// Scheduler — picks a host for each new sandbox (round-robin for now).
hostScheduler := scheduler.NewRoundRobinScheduler(queries)
// OAuth provider registry.
oauthRegistry := oauth.NewRegistry()
@ -86,16 +125,44 @@ func main() {
slog.Info("registered OAuth provider", "provider", "github")
}
// API server.
srv := api.New(queries, agentClient, pool, rdb, []byte(cfg.JWTSecret), oauthRegistry, cfg.OAuthRedirectURL)
// Channels: publisher, service, dispatcher.
if len(cfg.EncryptionKeyHex) != 64 {
slog.Error("WRENN_ENCRYPTION_KEY must be a hex-encoded 32-byte key (64 hex chars)")
os.Exit(1)
}
channelPub := channels.NewPublisher(rdb)
channelSvc := &channels.Service{DB: queries, EncKey: cfg.EncryptionKey}
channelDispatcher := channels.NewDispatcher(rdb, queries, cfg.EncryptionKey)
// Start reconciler.
reconciler := api.NewReconciler(queries, agentClient, "default", 5*time.Second)
reconciler.Start(ctx)
// Shared audit logger with event publishing.
al := audit.NewWithPublisher(queries, channelPub)
// API server.
srv := api.New(queries, hostPool, hostScheduler, pool, rdb, []byte(cfg.JWTSecret), oauthRegistry, cfg.OAuthRedirectURL, ca, al, channelSvc)
// Start template build workers (2 concurrent).
stopBuildWorkers := srv.BuildSvc.StartWorkers(ctx, 2)
defer stopBuildWorkers()
// Start channel event dispatcher.
channelDispatcher.Start(ctx)
// Start host monitor (passive + active reconciliation every 30s).
monitor := api.NewHostMonitor(queries, hostPool, al, 30*time.Second)
monitor.Start(ctx)
// Start metrics sampler (records per-team sandbox stats every 10s).
sampler := api.NewMetricsSampler(queries, 10*time.Second)
sampler.Start(ctx)
// Wrap the API handler with the sandbox proxy so that requests with
// {port}-{sandbox_id}.{domain} Host headers are routed to the sandbox's
// host agent. All other requests pass through to the normal API router.
proxyWrapper := api.NewSandboxProxyWrapper(srv.Handler(), queries, hostPool)
httpServer := &http.Server{
Addr: cfg.ListenAddr,
Handler: srv.Handler(),
Handler: proxyWrapper,
}
// Graceful shutdown on signal.
@ -114,7 +181,7 @@ func main() {
}
}()
slog.Info("control plane starting", "addr", cfg.ListenAddr, "agent", cfg.HostAgentAddr)
slog.Info("control plane starting", "addr", cfg.ListenAddr)
if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
slog.Error("http server error", "error", err)
os.Exit(1)

View File

@ -2,23 +2,32 @@ package main
import (
"context"
"crypto/tls"
"flag"
"log/slog"
"net/http"
"os"
"os/signal"
"path/filepath"
"sync"
"syscall"
"time"
"github.com/joho/godotenv"
"git.omukk.dev/wrenn/sandbox/internal/auth"
"git.omukk.dev/wrenn/sandbox/internal/devicemapper"
"git.omukk.dev/wrenn/sandbox/internal/hostagent"
"git.omukk.dev/wrenn/sandbox/internal/network"
"git.omukk.dev/wrenn/sandbox/internal/sandbox"
"git.omukk.dev/wrenn/sandbox/proto/hostagent/gen/hostagentv1connect"
)
func main() {
registrationToken := flag.String("register", "", "One-time registration token from the control plane")
// Best-effort load — missing .env file is fine.
_ = godotenv.Load()
registrationToken := flag.String("register", "", "One-time registration token from the control plane (required on first run)")
advertiseAddr := flag.String("address", "", "Externally-reachable address (ip:port) for this host agent")
flag.Parse()
@ -36,19 +45,33 @@ func main() {
slog.Warn("failed to enable ip_forward", "error", err)
}
// Clean up any stale dm-snapshot devices from a previous crash.
// Clean up stale resources from a previous crash.
devicemapper.CleanupStaleDevices()
network.CleanupStaleNamespaces()
listenAddr := envOrDefault("AGENT_LISTEN_ADDR", ":50051")
rootDir := envOrDefault("AGENT_FILES_ROOTDIR", "/var/lib/wrenn")
cpURL := os.Getenv("AGENT_CP_URL")
tokenFile := filepath.Join(rootDir, "host-token")
listenAddr := envOrDefault("WRENN_HOST_LISTEN_ADDR", ":50051")
rootDir := envOrDefault("WRENN_DIR", "/var/lib/wrenn")
cpURL := os.Getenv("WRENN_CP_URL")
credsFile := filepath.Join(rootDir, "host-credentials.json")
if cpURL == "" {
slog.Error("WRENN_CP_URL environment variable is required")
os.Exit(1)
}
if *advertiseAddr == "" {
slog.Error("--address flag is required (externally-reachable ip:port)")
os.Exit(1)
}
// Expand base images to the standard disk size (sparse, no extra physical
// disk). This ensures dm-snapshot sandboxes see the full size from boot.
if err := sandbox.EnsureImageSizes(rootDir, sandbox.DefaultDiskSizeMB); err != nil {
slog.Error("failed to expand base images", "error", err)
os.Exit(1)
}
cfg := sandbox.Config{
KernelPath: filepath.Join(rootDir, "kernels", "vmlinux"),
ImagesDir: filepath.Join(rootDir, "images"),
SandboxesDir: filepath.Join(rootDir, "sandboxes"),
SnapshotsDir: filepath.Join(rootDir, "snapshots"),
WrennDir: rootDir,
}
mgr := sandbox.New(cfg)
@ -58,66 +81,116 @@ func main() {
mgr.StartTTLReaper(ctx)
if *advertiseAddr == "" {
slog.Error("--address flag is required (externally-reachable ip:port)")
// Register with the control plane and start heartbeating.
creds, err := hostagent.Register(ctx, hostagent.RegistrationConfig{
CPURL: cpURL,
RegistrationToken: *registrationToken,
TokenFile: credsFile,
Address: *advertiseAddr,
})
if err != nil {
slog.Error("host registration failed", "error", err)
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,
slog.Info("host registered", "host_id", creds.HostID)
// httpServer is declared here so the shutdown func can reference it.
httpServer := &http.Server{Addr: listenAddr}
// mTLS is mandatory — refuse to start without a valid certificate.
var certStore hostagent.CertStore
if creds.CertPEM == "" || creds.KeyPEM == "" || creds.CACertPEM == "" {
slog.Error("mTLS certificate not received from CP — ensure WRENN_CA_CERT and WRENN_CA_KEY are configured on the control plane")
os.Exit(1)
}
if err := certStore.ParseAndStore(creds.CertPEM, creds.KeyPEM); err != nil {
slog.Error("failed to load host TLS certificate", "error", err)
os.Exit(1)
}
tlsCfg := auth.AgentTLSConfigFromPEM(creds.CACertPEM, certStore.GetCert)
if tlsCfg == nil {
slog.Error("failed to build agent TLS config: invalid CA certificate PEM")
os.Exit(1)
}
httpServer.TLSConfig = tlsCfg
slog.Info("mTLS enabled on agent server")
// doShutdown is the single shutdown path. sync.Once ensures mgr.Shutdown
// and httpServer.Shutdown are each called exactly once regardless of
// whether shutdown is triggered by a signal, a heartbeat 404, or the
// Terminate RPC.
var shutdownOnce sync.Once
doShutdown := func(reason string) {
shutdownOnce.Do(func() {
slog.Info("shutting down", "reason", reason)
cancel()
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
defer shutdownCancel()
mgr.Shutdown(shutdownCtx)
if err := httpServer.Shutdown(shutdownCtx); err != nil {
slog.Error("http server shutdown error", "error", err)
}
})
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)
srv := hostagent.NewServer(mgr, func() {
doShutdown("Terminate RPC received")
})
path, handler := hostagentv1connect.NewHostAgentServiceHandler(srv)
proxyHandler := hostagent.NewProxyHandler(mgr)
mux := http.NewServeMux()
mux.Handle(path, handler)
mux.Handle("/proxy/", proxyHandler)
httpServer.Handler = mux
httpServer := &http.Server{
Addr: listenAddr,
Handler: mux,
}
// Start heartbeat loop. Handler must be set before this because the
// immediate beat can trigger doShutdown → httpServer.Shutdown synchronously.
hostagent.StartHeartbeat(ctx, cpURL, credsFile, creds.HostID, 30*time.Second,
// pauseAll: called on 3 consecutive network failures.
func() {
pauseCtx, pauseCancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer pauseCancel()
mgr.PauseAll(pauseCtx)
},
// onDeleted: called when CP returns 404 (host was deleted).
func() {
doShutdown("host deleted from CP")
},
// onCredsRefreshed: hot-swap the TLS certificate after a JWT refresh.
func(tf *hostagent.TokenFile) {
if tf.CertPEM == "" || tf.KeyPEM == "" {
return
}
if err := certStore.ParseAndStore(tf.CertPEM, tf.KeyPEM); err != nil {
slog.Error("failed to hot-swap TLS cert after credentials refresh", "error", err)
} else {
slog.Info("TLS cert hot-swapped after credentials refresh")
}
},
)
// Graceful shutdown on signal.
// Graceful shutdown on SIGINT/SIGTERM.
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()
mgr.Shutdown(shutdownCtx)
if err := httpServer.Shutdown(shutdownCtx); err != nil {
slog.Error("http server shutdown error", "error", err)
}
doShutdown("signal: " + sig.String())
}()
slog.Info("host agent starting", "addr", listenAddr)
if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
slog.Error("http server error", "error", err)
slog.Info("host agent starting", "addr", listenAddr, "host_id", creds.HostID)
// TLSConfig is always set (mTLS is mandatory). Create the TLS listener
// manually because ListenAndServeTLS requires on-disk cert/key paths
// but we use GetCertificate callback for hot-swap support.
ln, err := tls.Listen("tcp", listenAddr, httpServer.TLSConfig)
if err != nil {
slog.Error("failed to start TLS listener", "error", err)
os.Exit(1)
}
if err := httpServer.Serve(ln); err != nil && err != http.ErrServerClosed {
slog.Error("https server error", "error", err)
os.Exit(1)
}