forked from wrenn/wrenn
v0.0.1 (#8)
Co-authored-by: Tasnim Kabir Sadik <tksadik92@gmail.com> Reviewed-on: wrenn/sandbox#8
This commit is contained in:
@ -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)
|
||||
|
||||
@ -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)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user