tksadik92/wrenn-releases
1
0
Fork 0
forked from wrenn/wrenn
Code Pull Requests Activity

v0.0.1 (#8)

Browse Source 
Co-authored-by: Tasnim Kabir Sadik <tksadik92@gmail.com>
Reviewed-on: wrenn/sandbox#8
This commit is contained in:
Rafeed M. Bhuiyan
2026-04-09 19:24:49 +00:00
parent 32e5a5a715
commit d3e4812e46
199 changed files with 24552 additions and 2776 deletions
Download Patch File Download Diff File Expand all files Collapse all files

177
cmd/host-agent/main.go
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)
}