From ce02cd822ade54178a1151ea84ec5d38276aba64 Mon Sep 17 00:00:00 2001 From: Nicolas Vargas Date: Mon, 23 Feb 2026 22:20:57 -0500 Subject: [PATCH 1/6] Feat(core): Add Docker client options for worker mode --- cmd/api/main.go | 2 +- internal/api/integration_test.go | 2 +- internal/docker/client.go | 332 +++++++++++++++++++++---------- internal/docker/names.go | 4 + models/sandbox.go | 1 + 5 files changed, 236 insertions(+), 105 deletions(-) diff --git a/cmd/api/main.go b/cmd/api/main.go index 25580f5..e6a208c 100644 --- a/cmd/api/main.go +++ b/cmd/api/main.go @@ -38,7 +38,7 @@ func main() { db := database.New("sandbox.db") repo := database.NewRepository(db) - dc := docker.New(repo) + dc := docker.New(docker.WithRepository(repo)) // --- Reverse proxy (multi-listen) --- proxyServer := proxy.New(cfg.BaseDomain, repo) diff --git a/internal/api/integration_test.go b/internal/api/integration_test.go index 3cab1f6..b1ae4bd 100644 --- a/internal/api/integration_test.go +++ b/internal/api/integration_test.go @@ -23,7 +23,7 @@ func realRouter() *gin.Engine { db := database.New(":memory:") repo := database.NewRepository(db) r := gin.New() - h := api.New(docker.New(repo), "localhost", ":3000") + h := api.New(docker.New(docker.WithRepository(repo)), "localhost", ":3000") h.RegisterHealthCheck(r) h.RegisterRoutes(r.Group("/v1")) return r diff --git a/internal/docker/client.go b/internal/docker/client.go index af2718b..87c9042 100644 --- a/internal/docker/client.go +++ b/internal/docker/client.go @@ -29,24 +29,43 @@ import ( // Client wraps the Docker SDK and exposes sandbox operations. type Client struct { cli *moby.Client - repo *database.Repository - timers sync.Map // map[containerID]*timerEntry - commands sync.Map // map[cmdID]*runningCommand - onCacheInvalid func(name string) // called when a sandbox's ports change or it is removed + repo *database.Repository // nil in worker mode (no DB) + hostIP string // bind IP for container ports ("127.0.0.1" or "0.0.0.0") + timers sync.Map // map[containerID]*timerEntry + commands sync.Map // map[cmdID]*runningCommand + onCacheInvalid func(name string) // called when a sandbox's ports change or it is removed +} + +// Option configures a Client. +type Option func(*Client) + +// WithRepository enables database persistence for sandboxes and commands. +// When nil (default), the client operates in stateless mode (worker). +func WithRepository(repo *database.Repository) Option { + return func(c *Client) { c.repo = repo } +} + +// WithHostIP sets the bind IP for container port mappings. +// Use "127.0.0.1" (default) for local/all-in-one, "0.0.0.0" for worker mode. +func WithHostIP(ip string) Option { + return func(c *Client) { c.hostIP = ip } } // runningCommand tracks a command that is currently executing. type runningCommand struct { - execID string // Docker exec instance ID - sandboxID string // parent sandbox container ID - cmd []string // original command (for pkill pattern) - cancel context.CancelFunc // cancels the exec context - stdout *ringBuffer // captures stdout - stderr *ringBuffer // captures stderr - done chan struct{} // closed when command finishes - mu sync.Mutex - exitCode int - finished bool + execID string // Docker exec instance ID + sandboxID string // parent sandbox container ID + cmd []string // original command (for pkill pattern) + cwd string // working directory + startedAt int64 // unix milliseconds + cancel context.CancelFunc // cancels the exec context + stdout *ringBuffer // captures stdout + stderr *ringBuffer // captures stderr + done chan struct{} // closed when command finishes + mu sync.Mutex + exitCode int + finishedAt int64 // unix milliseconds, 0 while running + finished bool } // timerEntry holds a timer and a cancel channel to avoid goroutine leaks. @@ -76,10 +95,10 @@ var ( mobyClient *moby.Client ) -// New creates a Docker Client with the given repository. -// The underlying Docker connection is a singleton (created once), -// but each Client gets its own repository. -func New(repo *database.Repository) *Client { +// New creates a Docker Client with the given options. +// The underlying Docker connection is a singleton (created once). +// Without WithRepository, the client runs in stateless mode (no DB). +func New(opts ...Option) *Client { once.Do(func() { cli, err := moby.NewClientWithOpts(moby.FromEnv, moby.WithAPIVersionNegotiation()) if err != nil { @@ -87,7 +106,11 @@ func New(repo *database.Repository) *Client { } mobyClient = cli }) - return &Client{cli: mobyClient, repo: repo} + c := &Client{cli: mobyClient, hostIP: "127.0.0.1"} + for _, o := range opts { + o(c) + } + return c } // SetCacheInvalidator registers a callback invoked when a sandbox's ports @@ -101,9 +124,11 @@ func (c *Client) invalidateCache(containerID string) { if c.onCacheInvalid == nil { return } - sb, err := c.repo.FindByID(containerID) - if err == nil && sb != nil && sb.Name != "" { - c.onCacheInvalid(sb.Name) + if c.repo != nil { + sb, err := c.repo.FindByID(containerID) + if err == nil && sb != nil && sb.Name != "" { + c.onCacheInvalid(sb.Name) + } } } @@ -113,19 +138,11 @@ func (c *Client) Ping(ctx context.Context) error { return err } -// List returns all sandboxes tracked in the database, enriched with live -// state from Docker. Stopped containers are always included. +// List returns all sandboxes, enriched with live state from Docker. +// With a repository: uses DB records as source of truth (includes stopped containers). +// Without a repository: lists directly from the Docker daemon. func (c *Client) List(ctx context.Context) ([]models.SandboxSummary, error) { - // Fetch all persisted sandboxes from the database. - dbSandboxes, err := c.repo.FindAll() - if err != nil { - return nil, err - } - if len(dbSandboxes) == 0 { - return []models.SandboxSummary{}, nil - } - - // Fetch all containers (including stopped) to build a lookup map. + // Fetch all containers (including stopped) from Docker. result, err := c.cli.ContainerList(ctx, moby.ContainerListOptions{All: true}) if err != nil { return nil, err @@ -155,6 +172,36 @@ func (c *Client) List(ctx context.Context) ([]models.SandboxSummary, error) { } } + // Without DB: build summaries directly from Docker containers. + if c.repo == nil { + summaries := make([]models.SandboxSummary, 0, len(lookup)) + for id, info := range lookup { + s := models.SandboxSummary{ + ID: id, + Name: info.Name, + Image: info.Image, + Status: info.Status, + State: info.State, + Ports: portKeys(info.Ports), + } + if entry := c.getTimerEntry(id); entry != nil { + ea := entry.expiresAt + s.ExpiresAt = &ea + } + summaries = append(summaries, s) + } + return summaries, nil + } + + // With DB: use persisted records as source of truth. + dbSandboxes, err := c.repo.FindAll() + if err != nil { + return nil, err + } + if len(dbSandboxes) == 0 { + return []models.SandboxSummary{}, nil + } + summaries := make([]models.SandboxSummary, 0, len(dbSandboxes)) for _, db := range dbSandboxes { s := models.SandboxSummary{ @@ -217,7 +264,7 @@ func (c *Client) Create(ctx context.Context, req models.CreateSandboxRequest) (m } hostCfg := &container.HostConfig{ - PortBindings: buildPortBindings(ports), + PortBindings: buildPortBindings(ports, c.hostIP), } // Apply resource limits (defaults: 1GB RAM, 1 vCPU) @@ -236,11 +283,18 @@ func (c *Client) Create(ctx context.Context, req models.CreateSandboxRequest) (m NanoCPUs: int64(cpus * 1e9), } - // Auto-generate a unique sandbox name. - name := generateUniqueName(func(n string) bool { - sb, _ := c.repo.FindByName(n) - return sb != nil - }) + // Determine sandbox name: use pre-generated name if provided, otherwise auto-generate. + name := req.Name + if name == "" { + if c.repo != nil { + name = generateUniqueName(func(n string) bool { + sb, _ := c.repo.FindByName(n) + return sb != nil + }) + } else { + name = generateUniqueName(nil) + } + } result, err := c.cli.ContainerCreate(ctx, moby.ContainerCreateOptions{ Config: cfg, @@ -270,15 +324,17 @@ func (c *Client) Create(ctx context.Context, req models.CreateSandboxRequest) (m assignedPorts := extractPorts(info.Container.NetworkSettings.Ports) - // Persist sandbox (fire-and-forget: log errors, don't block). - if err := c.repo.Save(database.Sandbox{ - ID: result.ID, - Name: name, - Image: req.Image, - Ports: database.JSONMap(assignedPorts), - Port: mainPort, - }); err != nil { - log.Printf("database: failed to persist sandbox %s: %v", result.ID, err) + // Persist sandbox if repository is available. + if c.repo != nil { + if err := c.repo.Save(database.Sandbox{ + ID: result.ID, + Name: name, + Image: req.Image, + Ports: database.JSONMap(assignedPorts), + Port: mainPort, + }); err != nil { + log.Printf("database: failed to persist sandbox %s: %v", result.ID, err) + } } return models.CreateSandboxResponse{ @@ -350,8 +406,10 @@ func (c *Client) Start(ctx context.Context, id string) (models.RestartResponse, ports := extractPorts(info.Container.NetworkSettings.Ports) - if dbErr := c.repo.UpdatePorts(id, database.JSONMap(ports)); dbErr != nil { - log.Printf("database: failed to update ports for sandbox %s: %v", id, dbErr) + if c.repo != nil { + if dbErr := c.repo.UpdatePorts(id, database.JSONMap(ports)); dbErr != nil { + log.Printf("database: failed to update ports for sandbox %s: %v", id, dbErr) + } } c.invalidateCache(id) @@ -406,8 +464,10 @@ func (c *Client) Restart(ctx context.Context, id string) (models.RestartResponse ports := extractPorts(info.Container.NetworkSettings.Ports) // Update persisted ports after restart (they may change). - if dbErr := c.repo.UpdatePorts(id, database.JSONMap(ports)); dbErr != nil { - log.Printf("database: failed to update ports for sandbox %s: %v", id, dbErr) + if c.repo != nil { + if dbErr := c.repo.UpdatePorts(id, database.JSONMap(ports)); dbErr != nil { + log.Printf("database: failed to update ports for sandbox %s: %v", id, dbErr) + } } c.invalidateCache(id) @@ -438,13 +498,14 @@ func (c *Client) Remove(ctx context.Context, id string) error { return err } - // Clean up command records from DB. - if dbErr := c.repo.DeleteCommandsBySandbox(id); dbErr != nil { - log.Printf("database: failed to delete commands for sandbox %s: %v", id, dbErr) - } - - if dbErr := c.repo.Delete(id); dbErr != nil { - log.Printf("database: failed to delete sandbox %s: %v", id, dbErr) + // Clean up DB records if repository is available. + if c.repo != nil { + if dbErr := c.repo.DeleteCommandsBySandbox(id); dbErr != nil { + log.Printf("database: failed to delete commands for sandbox %s: %v", id, dbErr) + } + if dbErr := c.repo.Delete(id); dbErr != nil { + log.Printf("database: failed to delete sandbox %s: %v", id, dbErr) + } } return nil } @@ -584,17 +645,19 @@ func (c *Client) ExecCommand(ctx context.Context, sandboxID string, req models.E return models.CommandDetail{}, wrapNotFound(err) } - // Persist command to DB. - argsJSON, _ := json.Marshal(req.Args) - if err := c.repo.SaveCommand(database.Command{ - ID: cmdID, - SandboxID: sandboxID, - Name: req.Command, - Args: string(argsJSON), - Cwd: req.Cwd, - StartedAt: now, - }); err != nil { - return models.CommandDetail{}, fmt.Errorf("save command: %w", err) + // Persist command to DB if repository is available. + if c.repo != nil { + argsJSON, _ := json.Marshal(req.Args) + if err := c.repo.SaveCommand(database.Command{ + ID: cmdID, + SandboxID: sandboxID, + Name: req.Command, + Args: string(argsJSON), + Cwd: req.Cwd, + StartedAt: now, + }); err != nil { + return models.CommandDetail{}, fmt.Errorf("save command: %w", err) + } } // Set up ring buffers and tracking. @@ -606,6 +669,8 @@ func (c *Client) ExecCommand(ctx context.Context, sandboxID string, req models.E execID: execCfg.ID, sandboxID: sandboxID, cmd: fullCmd, + cwd: req.Cwd, + startedAt: now, cancel: cancel, stdout: stdoutBuf, stderr: stderrBuf, @@ -629,11 +694,15 @@ func (c *Client) ExecCommand(ctx context.Context, sandboxID string, req models.E attached, err := c.cli.ExecAttach(execCtx, execCfg.ID, moby.ExecAttachOptions{}) if err != nil { log.Printf("exec attach %s: %v", cmdID, err) + fin := time.Now().UnixMilli() rc.mu.Lock() rc.exitCode = -1 + rc.finishedAt = fin rc.finished = true rc.mu.Unlock() - c.repo.UpdateCommandFinished(cmdID, -1, time.Now().UnixMilli()) + if c.repo != nil { + c.repo.UpdateCommandFinished(cmdID, -1, fin) + } return } defer attached.Close() @@ -648,13 +717,16 @@ func (c *Client) ExecCommand(ctx context.Context, sandboxID string, req models.E exitCode = inspect.ExitCode } - finishedAt := time.Now().UnixMilli() + fin := time.Now().UnixMilli() rc.mu.Lock() rc.exitCode = exitCode + rc.finishedAt = fin rc.finished = true rc.mu.Unlock() - c.repo.UpdateCommandFinished(cmdID, exitCode, finishedAt) + if c.repo != nil { + c.repo.UpdateCommandFinished(cmdID, exitCode, fin) + } }() return models.CommandDetail{ @@ -669,18 +741,31 @@ func (c *Client) ExecCommand(ctx context.Context, sandboxID string, req models.E // GetCommand returns command details by ID. func (c *Client) GetCommand(ctx context.Context, sandboxID, cmdID string) (models.CommandDetail, error) { - dbCmd, err := c.repo.FindCommandByID(cmdID) - if err != nil { - return models.CommandDetail{}, err - } - if dbCmd == nil { - return models.CommandDetail{}, ErrCommandNotFound + // Try in-memory first (works in both modes). + if v, ok := c.commands.Load(cmdID); ok { + rc := v.(*runningCommand) + if rc.sandboxID != sandboxID { + return models.CommandDetail{}, ErrCommandNotFound + } + return c.runningCommandToDetail(cmdID, rc), nil } - if dbCmd.SandboxID != sandboxID { - return models.CommandDetail{}, ErrCommandNotFound + + // Fall back to DB if available. + if c.repo != nil { + dbCmd, err := c.repo.FindCommandByID(cmdID) + if err != nil { + return models.CommandDetail{}, err + } + if dbCmd == nil { + return models.CommandDetail{}, ErrCommandNotFound + } + if dbCmd.SandboxID != sandboxID { + return models.CommandDetail{}, ErrCommandNotFound + } + return c.dbCommandToDetail(*dbCmd), nil } - return c.dbCommandToDetail(*dbCmd), nil + return models.CommandDetail{}, ErrCommandNotFound } // ListCommands returns all commands for a sandbox. @@ -690,32 +775,49 @@ func (c *Client) ListCommands(ctx context.Context, sandboxID string) ([]models.C return nil, wrapNotFound(err) } - dbCmds, err := c.repo.FindCommandsBySandbox(sandboxID) - if err != nil { - return nil, err + // With DB: use persisted records as source of truth. + if c.repo != nil { + dbCmds, err := c.repo.FindCommandsBySandbox(sandboxID) + if err != nil { + return nil, err + } + details := make([]models.CommandDetail, 0, len(dbCmds)) + for _, cmd := range dbCmds { + details = append(details, c.dbCommandToDetail(cmd)) + } + return details, nil } - details := make([]models.CommandDetail, 0, len(dbCmds)) - for _, cmd := range dbCmds { - details = append(details, c.dbCommandToDetail(cmd)) - } + // Without DB: list from in-memory sync.Map. + var details []models.CommandDetail + c.commands.Range(func(key, value any) bool { + cmdID := key.(string) + rc := value.(*runningCommand) + if rc.sandboxID == sandboxID { + details = append(details, c.runningCommandToDetail(cmdID, rc)) + } + return true + }) return details, nil } // KillCommand sends a signal to a running command. func (c *Client) KillCommand(ctx context.Context, sandboxID, cmdID string, signal int) (models.CommandDetail, error) { - // Look up running command. + // Look up running command in memory. v, ok := c.commands.Load(cmdID) if !ok { - // Check if it exists in DB. - dbCmd, err := c.repo.FindCommandByID(cmdID) - if err != nil { - return models.CommandDetail{}, err - } - if dbCmd == nil { - return models.CommandDetail{}, ErrCommandNotFound + // Not in memory — check DB if available, otherwise not found. + if c.repo != nil { + dbCmd, err := c.repo.FindCommandByID(cmdID) + if err != nil { + return models.CommandDetail{}, err + } + if dbCmd == nil { + return models.CommandDetail{}, ErrCommandNotFound + } + return models.CommandDetail{}, ErrCommandFinished } - return models.CommandDetail{}, ErrCommandFinished + return models.CommandDetail{}, ErrCommandNotFound } rc := v.(*runningCommand) @@ -838,6 +940,29 @@ func (c *Client) dbCommandToDetail(cmd database.Command) models.CommandDetail { return detail } +// runningCommandToDetail converts an in-memory runningCommand to models.CommandDetail. +// Used when no DB is available (worker mode). +func (c *Client) runningCommandToDetail(cmdID string, rc *runningCommand) models.CommandDetail { + rc.mu.Lock() + defer rc.mu.Unlock() + + detail := models.CommandDetail{ + ID: cmdID, + Name: rc.cmd[0], + Args: rc.cmd[1:], + Cwd: rc.cwd, + SandboxID: rc.sandboxID, + StartedAt: rc.startedAt, + } + if rc.finished { + ec := rc.exitCode + detail.ExitCode = &ec + fa := rc.finishedAt + detail.FinishedAt = &fa + } + return detail +} + // ReadFile reads the content of a file inside a sandbox. func (c *Client) ReadFile(ctx context.Context, id, path string) (string, error) { result, err := c.execWithStdin(ctx, id, []string{"cat", path}, nil) @@ -1130,9 +1255,10 @@ func buildExposedPorts(ports []string) network.PortSet { return ps } -// buildPortBindings creates port bindings that only listen on 127.0.0.1 (loopback). -// This ensures container ports are only reachable through the reverse proxy, not directly. -func buildPortBindings(ports []string) network.PortMap { +// buildPortBindings creates port bindings on the configured host IP. +// "127.0.0.1" (default/all-in-one) ensures ports are only reachable through the proxy. +// "0.0.0.0" (worker mode) allows the orchestrator to reach container ports over the network. +func buildPortBindings(ports []string, hostIP string) network.PortMap { if len(ports) == 0 { return nil } @@ -1142,7 +1268,7 @@ func buildPortBindings(ports []string) network.PortMap { if err != nil { continue } - pm[parsed] = []network.PortBinding{{HostIP: netip.MustParseAddr("127.0.0.1")}} + pm[parsed] = []network.PortBinding{{HostIP: netip.MustParseAddr(hostIP)}} } if len(pm) == 0 { return nil diff --git a/internal/docker/names.go b/internal/docker/names.go index 7309c9c..065bc4d 100644 --- a/internal/docker/names.go +++ b/internal/docker/names.go @@ -92,8 +92,12 @@ func generateName() string { } // generateUniqueName returns a name that does not collide with existing names. +// If exists is nil (no DB), returns the first generated name without collision checks. // After 10 attempts, it appends a random 4-digit suffix. func generateUniqueName(exists func(string) bool) string { + if exists == nil { + return generateName() + } for range 10 { name := generateName() if !exists(name) { diff --git a/models/sandbox.go b/models/sandbox.go index 60e598e..d2d30b9 100644 --- a/models/sandbox.go +++ b/models/sandbox.go @@ -11,6 +11,7 @@ type ResourceLimits struct { // CreateSandboxRequest is the body for POST /v1/sandboxes type CreateSandboxRequest struct { Image string `json:"image" binding:"required" example:"node:24"` + Name string `json:"name,omitempty"` // pre-generated name (used by orchestrator → worker). Empty = auto-generate. Ports []string `json:"ports" example:"3000,8080"` // container ports to expose, e.g. ["3000", "8080/tcp"]. First port is the default for proxy routing. Timeout int `json:"timeout" example:"900"` // seconds until auto-stop, 0 = default (900s) Resources *ResourceLimits `json:"resources"` // CPU/memory limits, nil = defaults (1GB RAM, 1 vCPU) From fdadf817bc1ab3991d936da29fc153754bd91f35 Mon Sep 17 00:00:00 2001 From: Nicolas Vargas Date: Mon, 23 Feb 2026 22:28:16 -0500 Subject: [PATCH 2/6] Feat(core): Add WorkerConfig and LoadWorker --- cmd/worker/main.go | 152 ++++++++++++ internal/config/config.go | 27 +- internal/worker/handler.go | 447 ++++++++++++++++++++++++++++++++++ internal/worker/middleware.go | 24 ++ internal/worker/router.go | 44 ++++ 5 files changed, 692 insertions(+), 2 deletions(-) create mode 100644 cmd/worker/main.go create mode 100644 internal/worker/handler.go create mode 100644 internal/worker/middleware.go create mode 100644 internal/worker/router.go diff --git a/cmd/worker/main.go b/cmd/worker/main.go new file mode 100644 index 0000000..fe04780 --- /dev/null +++ b/cmd/worker/main.go @@ -0,0 +1,152 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "log" + "net/http" + "os/signal" + "syscall" + "time" + + "github.com/gin-gonic/gin" + "open-sandbox/internal/config" + "open-sandbox/internal/docker" + "open-sandbox/internal/worker" +) + +func main() { + cfg := config.LoadWorker() + + dc := docker.New(docker.WithHostIP(cfg.HostIP)) + + r := gin.New() + r.Use(gin.Logger(), gin.Recovery()) + + internal := r.Group("/internal/v1") + if cfg.APIKey != "" { + internal.Use(worker.APIKeyAuth(cfg.APIKey)) + } + + h := worker.NewHandler(dc) + h.RegisterRoutes(internal) + + r.NoRoute(func(c *gin.Context) { + c.JSON(http.StatusNotFound, gin.H{ + "code": "NOT_FOUND", + "message": "route not found", + }) + }) + + // Graceful shutdown context. + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + // Self-register with orchestrator if configured. + var workerID string + if cfg.OrchestratorURL != "" && cfg.APIKey != "" { + var err error + workerID, err = registerWorker(cfg) + if err != nil { + log.Fatalf("worker registration failed: %v", err) + } + log.Printf("registered with orchestrator as worker %s", workerID) + } else { + log.Println("no orchestrator configured, running standalone") + } + + srv := &http.Server{Addr: cfg.Addr, Handler: r} + + go func() { + log.Printf("worker listening on %s (host_ip: %s)", cfg.Addr, cfg.HostIP) + if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed { + log.Fatalf("worker listen: %v", err) + } + }() + + <-ctx.Done() + log.Println("shutting down worker...") + + shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + // Deregister from orchestrator. + if workerID != "" && cfg.OrchestratorURL != "" { + if err := deregisterWorker(cfg, workerID); err != nil { + log.Printf("worker deregistration failed: %v", err) + } else { + log.Printf("deregistered worker %s", workerID) + } + } + + dc.Shutdown(shutdownCtx) + if err := srv.Shutdown(shutdownCtx); err != nil { + log.Fatalf("worker shutdown: %v", err) + } + + log.Println("worker stopped") +} + +// registerWorker calls POST {orchestratorURL}/internal/v1/workers/register. +func registerWorker(cfg *config.WorkerConfig) (string, error) { + body, _ := json.Marshal(map[string]string{ + "url": workerURL(cfg), + }) + + req, err := http.NewRequest(http.MethodPost, cfg.OrchestratorURL+"/internal/v1/workers/register", bytes.NewReader(body)) + if err != nil { + return "", err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Worker-Key", cfg.APIKey) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return "", fmt.Errorf("connect to orchestrator: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated { + return "", fmt.Errorf("orchestrator returned %d", resp.StatusCode) + } + + var result struct { + WorkerID string `json:"worker_id"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return "", fmt.Errorf("decode response: %w", err) + } + return result.WorkerID, nil +} + +// deregisterWorker calls DELETE {orchestratorURL}/internal/v1/workers/{id}. +func deregisterWorker(cfg *config.WorkerConfig, workerID string) error { + req, err := http.NewRequest(http.MethodDelete, cfg.OrchestratorURL+"/internal/v1/workers/"+workerID, nil) + if err != nil { + return err + } + req.Header.Set("X-Worker-Key", cfg.APIKey) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNoContent { + return fmt.Errorf("orchestrator returned %d", resp.StatusCode) + } + return nil +} + +// workerURL builds the worker's externally reachable URL from its config. +func workerURL(cfg *config.WorkerConfig) string { + // If addr starts with ":", prepend the hostname or default to the host IP. + addr := cfg.Addr + if addr[0] == ':' { + addr = cfg.HostIP + addr + } + return "http://" + addr +} diff --git a/internal/config/config.go b/internal/config/config.go index da95887..792815e 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -6,7 +6,7 @@ import ( "strings" ) -// Config holds all application configuration. +// Config holds all-in-one application configuration. type Config struct { Addr string // HTTP listen address, e.g. ":8080" APIKey string // API key for authentication (env API_KEY). Empty = auth disabled. @@ -22,7 +22,7 @@ func (c *Config) PrimaryProxyAddr() string { return c.ProxyAddrs[0] } -// Load parses flags and env vars. Flags take precedence over env vars. +// Load parses flags and env vars for the all-in-one binary. func Load() *Config { addr := flag.String("addr", envOrDefault("ADDR", ":8080"), "HTTP listen address") proxyAddr := flag.String("proxy-addr", envOrDefault("PROXY_ADDR", ":80,:3000"), "Comma-separated proxy listen addresses (first is used for URL generation)") @@ -37,6 +37,29 @@ func Load() *Config { } } +// WorkerConfig holds configuration for the worker binary. +type WorkerConfig struct { + Addr string // Worker HTTP listen address (default ":9090") + APIKey string // Shared API key for orchestrator ↔ worker auth + OrchestratorURL string // Orchestrator URL for self-registration + HostIP string // Bind IP for container ports ("0.0.0.0" or "127.0.0.1") +} + +// LoadWorker parses flags and env vars for the worker binary. +func LoadWorker() *WorkerConfig { + addr := flag.String("addr", envOrDefault("WORKER_ADDR", ":9090"), "Worker HTTP listen address") + orchestratorURL := flag.String("orchestrator-url", envOrDefault("ORCHESTRATOR_URL", ""), "Orchestrator URL for self-registration") + hostIP := flag.String("host-ip", envOrDefault("HOST_IP", "0.0.0.0"), "Bind IP for container ports") + flag.Parse() + + return &WorkerConfig{ + Addr: *addr, + APIKey: os.Getenv("WORKER_API_KEY"), + OrchestratorURL: *orchestratorURL, + HostIP: *hostIP, + } +} + // parseAddrs splits a comma-separated list of addresses and trims whitespace. func parseAddrs(raw string) []string { parts := strings.Split(raw, ",") diff --git a/internal/worker/handler.go b/internal/worker/handler.go new file mode 100644 index 0000000..d6fe1ee --- /dev/null +++ b/internal/worker/handler.go @@ -0,0 +1,447 @@ +package worker + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "sync" + + "github.com/gin-gonic/gin" + "open-sandbox/internal/docker" + "open-sandbox/models" +) + +// DockerClient defines the sandbox operations used by the worker handlers. +// Identical to api.DockerClient — duplicated to avoid import cycles. +type DockerClient interface { + Ping(ctx context.Context) error + List(ctx context.Context) ([]models.SandboxSummary, error) + Create(ctx context.Context, req models.CreateSandboxRequest) (models.CreateSandboxResponse, error) + Inspect(ctx context.Context, id string) (models.SandboxDetail, error) + Start(ctx context.Context, id string) (models.RestartResponse, error) + Stop(ctx context.Context, id string) error + Restart(ctx context.Context, id string) (models.RestartResponse, error) + Remove(ctx context.Context, id string) error + Pause(ctx context.Context, id string) error + Resume(ctx context.Context, id string) error + RenewExpiration(ctx context.Context, id string, timeout int) error + ExecCommand(ctx context.Context, sandboxID string, req models.ExecCommandRequest) (models.CommandDetail, error) + GetCommand(ctx context.Context, sandboxID, cmdID string) (models.CommandDetail, error) + ListCommands(ctx context.Context, sandboxID string) ([]models.CommandDetail, error) + KillCommand(ctx context.Context, sandboxID, cmdID string, signal int) (models.CommandDetail, error) + StreamCommandLogs(ctx context.Context, sandboxID, cmdID string) (io.ReadCloser, io.ReadCloser, error) + GetCommandLogs(ctx context.Context, sandboxID, cmdID string) (models.CommandLogsResponse, error) + WaitCommand(ctx context.Context, sandboxID, cmdID string) (models.CommandDetail, error) + Stats(ctx context.Context, id string) (models.SandboxStats, error) + ReadFile(ctx context.Context, id, path string) (string, error) + WriteFile(ctx context.Context, id, path, content string) error + DeleteFile(ctx context.Context, id, path string) error + ListDir(ctx context.Context, id, path string) (string, error) + PullImage(ctx context.Context, image string) error + RemoveImage(ctx context.Context, id string, force bool) error + InspectImage(ctx context.Context, id string) (models.ImageDetail, error) + ListImages(ctx context.Context) ([]models.ImageSummary, error) +} + +// Handler holds worker handler dependencies. +type Handler struct { + docker DockerClient +} + +// NewHandler creates a worker Handler. +func NewHandler(d DockerClient) *Handler { + return &Handler{docker: d} +} + +// --- Error helpers --- + +type errorResponse struct { + Code string `json:"code"` + Message string `json:"message"` +} + +func mapError(c *gin.Context, err error) { + if errors.Is(err, docker.ErrNotFound) { + c.JSON(http.StatusNotFound, errorResponse{Code: "NOT_FOUND", Message: "sandbox not found"}) + return + } + if errors.Is(err, docker.ErrImageNotFound) { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: "image not found locally"}) + return + } + if errors.Is(err, docker.ErrAlreadyRunning) || errors.Is(err, docker.ErrAlreadyStopped) || + errors.Is(err, docker.ErrAlreadyPaused) || errors.Is(err, docker.ErrNotPaused) || + errors.Is(err, docker.ErrNotRunning) || errors.Is(err, docker.ErrCommandFinished) { + c.JSON(http.StatusConflict, errorResponse{Code: "CONFLICT", Message: err.Error()}) + return + } + if errors.Is(err, docker.ErrCommandNotFound) { + c.JSON(http.StatusNotFound, errorResponse{Code: "NOT_FOUND", Message: "command not found"}) + return + } + if errors.Is(err, context.DeadlineExceeded) { + c.JSON(http.StatusRequestTimeout, errorResponse{Code: "TIMEOUT", Message: "operation timed out"}) + return + } + c.JSON(http.StatusInternalServerError, errorResponse{Code: "INTERNAL_ERROR", Message: err.Error()}) +} + +// --- Sandbox handlers --- + +func (h *Handler) health(c *gin.Context) { + if err := h.docker.Ping(c.Request.Context()); err != nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"status": "unhealthy", "error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"status": "healthy"}) +} + +func (h *Handler) listSandboxes(c *gin.Context) { + items, err := h.docker.List(c.Request.Context()) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"sandboxes": items}) +} + +func (h *Handler) createSandbox(c *gin.Context) { + var req models.CreateSandboxRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: err.Error()}) + return + } + result, err := h.docker.Create(c.Request.Context(), req) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusCreated, result) +} + +func (h *Handler) inspectSandbox(c *gin.Context) { + info, err := h.docker.Inspect(c.Request.Context(), c.Param("id")) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, info) +} + +func (h *Handler) deleteSandbox(c *gin.Context) { + if err := h.docker.Remove(c.Request.Context(), c.Param("id")); err != nil { + mapError(c, err) + return + } + c.Status(http.StatusNoContent) +} + +func (h *Handler) startSandbox(c *gin.Context) { + result, err := h.docker.Start(c.Request.Context(), c.Param("id")) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, result) +} + +func (h *Handler) stopSandbox(c *gin.Context) { + if err := h.docker.Stop(c.Request.Context(), c.Param("id")); err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"status": "stopped"}) +} + +func (h *Handler) restartSandbox(c *gin.Context) { + result, err := h.docker.Restart(c.Request.Context(), c.Param("id")) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, result) +} + +func (h *Handler) pauseSandbox(c *gin.Context) { + if err := h.docker.Pause(c.Request.Context(), c.Param("id")); err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"status": "paused"}) +} + +func (h *Handler) resumeSandbox(c *gin.Context) { + if err := h.docker.Resume(c.Request.Context(), c.Param("id")); err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"status": "resumed"}) +} + +func (h *Handler) renewExpiration(c *gin.Context) { + var req models.RenewExpirationRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: err.Error()}) + return + } + if err := h.docker.RenewExpiration(c.Request.Context(), c.Param("id"), req.Timeout); err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, models.RenewExpirationResponse{Status: "renewed", Timeout: req.Timeout}) +} + +func (h *Handler) getStats(c *gin.Context) { + stats, err := h.docker.Stats(c.Request.Context(), c.Param("id")) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, stats) +} + +// --- Command handlers --- + +func (h *Handler) execCommand(c *gin.Context) { + var req models.ExecCommandRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: err.Error()}) + return + } + cmd, err := h.docker.ExecCommand(c.Request.Context(), c.Param("id"), req) + if err != nil { + mapError(c, err) + return + } + + if c.Query("wait") == "true" { + h.streamWait(c, c.Param("id"), cmd.ID) + return + } + c.JSON(http.StatusOK, models.CommandResponse{Command: cmd}) +} + +func (h *Handler) listCommands(c *gin.Context) { + cmds, err := h.docker.ListCommands(c.Request.Context(), c.Param("id")) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, models.CommandListResponse{Commands: cmds}) +} + +func (h *Handler) getCommand(c *gin.Context) { + cmd, err := h.docker.GetCommand(c.Request.Context(), c.Param("id"), c.Param("cmdId")) + if err != nil { + mapError(c, err) + return + } + if c.Query("wait") == "true" { + h.streamWait(c, c.Param("id"), c.Param("cmdId")) + return + } + c.JSON(http.StatusOK, models.CommandResponse{Command: cmd}) +} + +func (h *Handler) killCommand(c *gin.Context) { + var req models.KillCommandRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: err.Error()}) + return + } + cmd, err := h.docker.KillCommand(c.Request.Context(), c.Param("id"), c.Param("cmdId"), req.Signal) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, models.CommandResponse{Command: cmd}) +} + +func (h *Handler) getCommandLogs(c *gin.Context) { + sandboxID := c.Param("id") + cmdID := c.Param("cmdId") + + if c.Query("stream") == "true" { + h.streamLogs(c, sandboxID, cmdID) + return + } + logs, err := h.docker.GetCommandLogs(c.Request.Context(), sandboxID, cmdID) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, logs) +} + +// --- File handlers --- + +func (h *Handler) readFile(c *gin.Context) { + path := c.Query("path") + if path == "" { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: "path query param is required"}) + return + } + content, err := h.docker.ReadFile(c.Request.Context(), c.Param("id"), path) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, models.FileReadResponse{Path: path, Content: content}) +} + +func (h *Handler) writeFile(c *gin.Context) { + path := c.Query("path") + if path == "" { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: "path query param is required"}) + return + } + var req models.FileWriteRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: err.Error()}) + return + } + if err := h.docker.WriteFile(c.Request.Context(), c.Param("id"), path, req.Content); err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"path": path, "status": "written"}) +} + +func (h *Handler) deleteFile(c *gin.Context) { + path := c.Query("path") + if path == "" { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: "path query param is required"}) + return + } + if err := h.docker.DeleteFile(c.Request.Context(), c.Param("id"), path); err != nil { + mapError(c, err) + return + } + c.Status(http.StatusNoContent) +} + +func (h *Handler) listDir(c *gin.Context) { + path := c.DefaultQuery("path", "/") + output, err := h.docker.ListDir(c.Request.Context(), c.Param("id"), path) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, models.FileListResponse{Path: path, Output: output}) +} + +// --- Image handlers --- + +func (h *Handler) listImages(c *gin.Context) { + images, err := h.docker.ListImages(c.Request.Context()) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"images": images}) +} + +func (h *Handler) inspectImage(c *gin.Context) { + detail, err := h.docker.InspectImage(c.Request.Context(), c.Param("id")) + if err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, detail) +} + +func (h *Handler) pullImage(c *gin.Context) { + var req models.ImagePullRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, errorResponse{Code: "BAD_REQUEST", Message: err.Error()}) + return + } + if err := h.docker.PullImage(c.Request.Context(), req.Image); err != nil { + mapError(c, err) + return + } + c.JSON(http.StatusOK, models.ImagePullResponse{Status: "pulled", Image: req.Image}) +} + +func (h *Handler) deleteImage(c *gin.Context) { + force := c.Query("force") == "true" + if err := h.docker.RemoveImage(c.Request.Context(), c.Param("id"), force); err != nil { + mapError(c, err) + return + } + c.Status(http.StatusNoContent) +} + +// --- Streaming helpers --- + +func (h *Handler) streamLogs(c *gin.Context, sandboxID, cmdID string) { + stdoutR, stderrR, err := h.docker.StreamCommandLogs(c.Request.Context(), sandboxID, cmdID) + if err != nil { + mapError(c, err) + return + } + defer stdoutR.Close() + defer stderrR.Close() + + c.Header("Content-Type", "application/x-ndjson") + c.Status(http.StatusOK) + flusher, _ := c.Writer.(http.Flusher) + enc := json.NewEncoder(c.Writer) + + type logLine struct { + Type string `json:"type"` + Data string `json:"data"` + } + + lines := make(chan logLine, 64) + readStream := func(r io.ReadCloser, streamType string) { + scanner := bufio.NewScanner(r) + for scanner.Scan() { + lines <- logLine{Type: streamType, Data: scanner.Text() + "\n"} + } + } + + var wg sync.WaitGroup + wg.Add(2) + go func() { defer wg.Done(); readStream(stdoutR, "stdout") }() + go func() { defer wg.Done(); readStream(stderrR, "stderr") }() + go func() { wg.Wait(); close(lines) }() + + for line := range lines { + if c.IsAborted() { + return + } + enc.Encode(line) + if flusher != nil { + flusher.Flush() + } + } +} + +func (h *Handler) streamWait(c *gin.Context, sandboxID, cmdID string) { + c.Header("Content-Type", "application/x-ndjson") + c.Status(http.StatusOK) + flusher, _ := c.Writer.(http.Flusher) + enc := json.NewEncoder(c.Writer) + + cmd, err := h.docker.GetCommand(c.Request.Context(), sandboxID, cmdID) + if err != nil { + return + } + enc.Encode(models.CommandResponse{Command: cmd}) + if flusher != nil { + flusher.Flush() + } + + cmd, err = h.docker.WaitCommand(c.Request.Context(), sandboxID, cmdID) + if err != nil { + return + } + enc.Encode(models.CommandResponse{Command: cmd}) + if flusher != nil { + flusher.Flush() + } +} diff --git a/internal/worker/middleware.go b/internal/worker/middleware.go new file mode 100644 index 0000000..36b3909 --- /dev/null +++ b/internal/worker/middleware.go @@ -0,0 +1,24 @@ +package worker + +import ( + "crypto/subtle" + "net/http" + "strings" + + "github.com/gin-gonic/gin" +) + +// APIKeyAuth returns a middleware that validates the X-Worker-Key header. +func APIKeyAuth(key string) gin.HandlerFunc { + return func(c *gin.Context) { + token := strings.TrimSpace(c.GetHeader("X-Worker-Key")) + if token == "" || subtle.ConstantTimeCompare([]byte(token), []byte(key)) != 1 { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ + "code": "UNAUTHORIZED", + "message": "invalid or missing worker api key", + }) + return + } + c.Next() + } +} diff --git a/internal/worker/router.go b/internal/worker/router.go new file mode 100644 index 0000000..1d05ed4 --- /dev/null +++ b/internal/worker/router.go @@ -0,0 +1,44 @@ +package worker + +import "github.com/gin-gonic/gin" + +// RegisterRoutes attaches all internal worker API routes to the given router group. +// All routes are prefixed with /internal/v1. +func (h *Handler) RegisterRoutes(g *gin.RouterGroup) { + // Sandboxes + sb := g.Group("/sandboxes") + sb.GET("", h.listSandboxes) + sb.POST("", h.createSandbox) + sb.GET("/:id", h.inspectSandbox) + sb.DELETE("/:id", h.deleteSandbox) + sb.POST("/:id/start", h.startSandbox) + sb.POST("/:id/stop", h.stopSandbox) + sb.POST("/:id/restart", h.restartSandbox) + sb.POST("/:id/pause", h.pauseSandbox) + sb.POST("/:id/resume", h.resumeSandbox) + sb.POST("/:id/renew-expiration", h.renewExpiration) + sb.GET("/:id/stats", h.getStats) + + // Commands + sb.POST("/:id/cmd", h.execCommand) + sb.GET("/:id/cmd", h.listCommands) + sb.GET("/:id/cmd/:cmdId", h.getCommand) + sb.POST("/:id/cmd/:cmdId/kill", h.killCommand) + sb.GET("/:id/cmd/:cmdId/logs", h.getCommandLogs) + + // Files + sb.GET("/:id/files", h.readFile) + sb.PUT("/:id/files", h.writeFile) + sb.DELETE("/:id/files", h.deleteFile) + sb.GET("/:id/files/list", h.listDir) + + // Images + img := g.Group("/images") + img.GET("", h.listImages) + img.GET("/:id", h.inspectImage) + img.POST("/pull", h.pullImage) + img.DELETE("/:id", h.deleteImage) + + // Health (no auth — registered outside the auth group in main) + g.GET("/health", h.health) +} From 76d9d6b7c82188bda59c7db299c3c6319cfd2b50 Mon Sep 17 00:00:00 2001 From: Nicolas Vargas Date: Mon, 23 Feb 2026 22:42:01 -0500 Subject: [PATCH 3/6] Feat(core): Add worker model and orchestrator support --- cmd/orchestrator/main.go | 109 ++++++ internal/config/config.go | 33 ++ internal/database/database.go | 2 +- internal/database/models.go | 20 +- internal/database/repository.go | 73 ++++ internal/docker/client.go | 4 +- internal/docker/names.go | 4 +- internal/docker/names_test.go | 4 +- internal/orchestrator/client.go | 632 ++++++++++++++++++++++++++++++ internal/orchestrator/handler.go | 69 ++++ internal/orchestrator/registry.go | 145 +++++++ internal/proxy/resolver.go | 34 +- 12 files changed, 1115 insertions(+), 14 deletions(-) create mode 100644 cmd/orchestrator/main.go create mode 100644 internal/orchestrator/client.go create mode 100644 internal/orchestrator/handler.go create mode 100644 internal/orchestrator/registry.go diff --git a/cmd/orchestrator/main.go b/cmd/orchestrator/main.go new file mode 100644 index 0000000..085dca3 --- /dev/null +++ b/cmd/orchestrator/main.go @@ -0,0 +1,109 @@ +package main + +import ( + "context" + "log" + "net/http" + "os/signal" + "strings" + "syscall" + "time" + + "github.com/gin-gonic/gin" + "open-sandbox/internal/api" + "open-sandbox/internal/config" + "open-sandbox/internal/database" + "open-sandbox/internal/orchestrator" + "open-sandbox/internal/proxy" + "open-sandbox/internal/worker" +) + +func main() { + cfg := config.LoadOrchestrator() + + db := database.New("sandbox.db") + repo := database.NewRepository(db) + + // Worker registry (loads active workers from DB). + registry := orchestrator.NewRegistry(repo) + + // RemoteDockerClient — implements DockerClient via HTTP to workers. + dc := orchestrator.NewRemoteClient(registry, repo) + + // --- Reverse proxy (multi-listen) --- + proxyServer := proxy.New(cfg.BaseDomain, repo) + proxyHandler := proxyServer.Handler() + + var proxySrvs []*http.Server + for _, addr := range cfg.ProxyAddrs { + srv := &http.Server{Addr: addr, Handler: proxyHandler} + proxySrvs = append(proxySrvs, srv) + go func(a string) { + log.Printf("proxy listening on %s (domain: *.%s)", a, cfg.BaseDomain) + if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed { + log.Fatalf("proxy listen %s: %v", a, err) + } + }(addr) + } + log.Printf("proxy URLs via %s", strings.Join(cfg.ProxyAddrs, ", ")) + + // --- API server --- + r := gin.New() + r.Use(gin.Logger(), gin.Recovery()) + + // Public API (same as all-in-one, but backed by RemoteDockerClient). + v1 := r.Group("/v1") + if cfg.APIKey != "" { + v1.Use(api.APIKeyAuth(cfg.APIKey)) + } + + h := api.New(dc, cfg.BaseDomain, cfg.PrimaryProxyAddr()) + h.RegisterHealthCheck(r) + h.RegisterRoutes(v1) + + // Internal API for worker registration. + internal := r.Group("/internal/v1") + if cfg.WorkerKey != "" { + internal.Use(worker.APIKeyAuth(cfg.WorkerKey)) + } + + oh := orchestrator.NewHandler(registry) + oh.RegisterRoutes(internal) + + r.NoRoute(func(c *gin.Context) { + c.JSON(http.StatusNotFound, gin.H{ + "code": "NOT_FOUND", + "message": "route not found", + }) + }) + + // Graceful shutdown. + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + srv := &http.Server{Addr: cfg.Addr, Handler: r} + + go func() { + log.Printf("orchestrator API listening on %s", cfg.Addr) + if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed { + log.Fatalf("api listen: %v", err) + } + }() + + <-ctx.Done() + log.Println("shutting down orchestrator...") + + shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + for _, ps := range proxySrvs { + if err := ps.Shutdown(shutdownCtx); err != nil { + log.Printf("proxy shutdown %s: %v", ps.Addr, err) + } + } + if err := srv.Shutdown(shutdownCtx); err != nil { + log.Fatalf("api shutdown: %v", err) + } + + log.Println("orchestrator stopped") +} diff --git a/internal/config/config.go b/internal/config/config.go index 792815e..529b70d 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -60,6 +60,39 @@ func LoadWorker() *WorkerConfig { } } +// OrchestratorConfig holds configuration for the orchestrator binary. +type OrchestratorConfig struct { + Addr string // API HTTP listen address (default ":8080") + APIKey string // Public API key + WorkerKey string // Shared key for worker ↔ orchestrator auth + ProxyAddrs []string // Proxy listen addresses + BaseDomain string // Base domain for proxy subdomains +} + +// PrimaryProxyAddr returns the first proxy address, used for generating URLs. +func (c *OrchestratorConfig) PrimaryProxyAddr() string { + if len(c.ProxyAddrs) == 0 { + return ":80" + } + return c.ProxyAddrs[0] +} + +// LoadOrchestrator parses flags and env vars for the orchestrator binary. +func LoadOrchestrator() *OrchestratorConfig { + addr := flag.String("addr", envOrDefault("ADDR", ":8080"), "API HTTP listen address") + proxyAddr := flag.String("proxy-addr", envOrDefault("PROXY_ADDR", ":80,:3000"), "Comma-separated proxy listen addresses") + baseDomain := flag.String("base-domain", envOrDefault("BASE_DOMAIN", "localhost"), "Base domain for subdomain routing") + flag.Parse() + + return &OrchestratorConfig{ + Addr: *addr, + APIKey: os.Getenv("API_KEY"), + WorkerKey: os.Getenv("WORKER_API_KEY"), + ProxyAddrs: parseAddrs(*proxyAddr), + BaseDomain: *baseDomain, + } +} + // parseAddrs splits a comma-separated list of addresses and trims whitespace. func parseAddrs(raw string) []string { parts := strings.Split(raw, ",") diff --git a/internal/database/database.go b/internal/database/database.go index 39cc2a4..05f34a4 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -18,7 +18,7 @@ func New(path string) *gorm.DB { log.Fatalf("database: failed to open %s: %v", path, err) } - if err := db.AutoMigrate(&Sandbox{}, &Command{}); err != nil { + if err := db.AutoMigrate(&Sandbox{}, &Command{}, &Worker{}); err != nil { log.Fatalf("database: migration failed: %v", err) } diff --git a/internal/database/models.go b/internal/database/models.go index 312f0a6..9fa9eff 100644 --- a/internal/database/models.go +++ b/internal/database/models.go @@ -36,11 +36,21 @@ func (j *JSONMap) Scan(src any) error { // Sandbox persists the container ID, metadata, and its assigned host ports. type Sandbox struct { - ID string `gorm:"primaryKey"` // Docker container ID - Name string - Image string - Ports JSONMap `gorm:"type:json"` // e.g. {"3000/tcp": "32768"} - Port string // container port exposed, e.g. "3000/tcp" + ID string `gorm:"primaryKey"` // Docker container ID + Name string + Image string + Ports JSONMap `gorm:"type:json"` // e.g. {"3000/tcp": "32768"} + Port string // container port exposed, e.g. "3000/tcp" + WorkerID string // FK to Worker.ID — empty in all-in-one mode +} + +// Worker represents a registered worker node. +type Worker struct { + ID string `gorm:"primaryKey"` + URL string `gorm:"uniqueIndex"` // e.g. "http://10.0.0.2:9090" + APIKey string // API key for authenticating with this worker + Status string // "active" / "inactive" + CreatedAt int64 // unix milliseconds } // Command persists an executed command's metadata and result. diff --git a/internal/database/repository.go b/internal/database/repository.go index f8804f5..bb651b7 100644 --- a/internal/database/repository.go +++ b/internal/database/repository.go @@ -100,3 +100,76 @@ func (r *Repository) UpdateCommandFinished(id string, exitCode int, finishedAt i func (r *Repository) DeleteCommandsBySandbox(sandboxID string) error { return r.db.Where("sandbox_id = ?", sandboxID).Delete(&Command{}).Error } + +// --- Worker operations --- + +// SaveWorker creates or updates a worker record. +func (r *Repository) SaveWorker(w Worker) error { + return r.db.Save(&w).Error +} + +// FindWorkerByID returns a worker by ID, or nil if not found. +func (r *Repository) FindWorkerByID(id string) (*Worker, error) { + var w Worker + if err := r.db.First(&w, "id = ?", id).Error; err != nil { + if err == gorm.ErrRecordNotFound { + return nil, nil + } + return nil, err + } + return &w, nil +} + +// FindWorkerByURL returns a worker by URL, or nil if not found. +func (r *Repository) FindWorkerByURL(url string) (*Worker, error) { + var w Worker + if err := r.db.First(&w, "url = ?", url).Error; err != nil { + if err == gorm.ErrRecordNotFound { + return nil, nil + } + return nil, err + } + return &w, nil +} + +// FindActiveWorkers returns all workers with status "active". +func (r *Repository) FindActiveWorkers() ([]Worker, error) { + var workers []Worker + if err := r.db.Where("status = ?", "active").Find(&workers).Error; err != nil { + return nil, err + } + return workers, nil +} + +// DeleteWorker removes a worker record by ID. +func (r *Repository) DeleteWorker(id string) error { + return r.db.Delete(&Worker{}, "id = ?", id).Error +} + +// UpdateWorkerStatus changes a worker's status. +func (r *Repository) UpdateWorkerStatus(id, status string) error { + return r.db.Model(&Worker{}).Where("id = ?", id).Update("status", status).Error +} + +// FindSandboxWorker returns the sandbox joined with its worker URL. +// Returns the sandbox and the worker URL, or nil if not found. +func (r *Repository) FindSandboxWorker(sandboxID string) (*Sandbox, string, error) { + var sb Sandbox + if err := r.db.First(&sb, "id = ?", sandboxID).Error; err != nil { + if err == gorm.ErrRecordNotFound { + return nil, "", nil + } + return nil, "", err + } + if sb.WorkerID == "" { + return &sb, "", nil + } + w, err := r.FindWorkerByID(sb.WorkerID) + if err != nil { + return nil, "", err + } + if w == nil { + return &sb, "", nil + } + return &sb, w.URL, nil +} diff --git a/internal/docker/client.go b/internal/docker/client.go index 87c9042..4c5d9a5 100644 --- a/internal/docker/client.go +++ b/internal/docker/client.go @@ -287,12 +287,12 @@ func (c *Client) Create(ctx context.Context, req models.CreateSandboxRequest) (m name := req.Name if name == "" { if c.repo != nil { - name = generateUniqueName(func(n string) bool { + name = GenerateUniqueName(func(n string) bool { sb, _ := c.repo.FindByName(n) return sb != nil }) } else { - name = generateUniqueName(nil) + name = GenerateUniqueName(nil) } } diff --git a/internal/docker/names.go b/internal/docker/names.go index 065bc4d..ab2d58e 100644 --- a/internal/docker/names.go +++ b/internal/docker/names.go @@ -91,10 +91,10 @@ func generateName() string { } } -// generateUniqueName returns a name that does not collide with existing names. +// GenerateUniqueName returns a name that does not collide with existing names. // If exists is nil (no DB), returns the first generated name without collision checks. // After 10 attempts, it appends a random 4-digit suffix. -func generateUniqueName(exists func(string) bool) string { +func GenerateUniqueName(exists func(string) bool) string { if exists == nil { return generateName() } diff --git a/internal/docker/names_test.go b/internal/docker/names_test.go index 089a379..31217af 100644 --- a/internal/docker/names_test.go +++ b/internal/docker/names_test.go @@ -48,7 +48,7 @@ func TestGenerateUniqueName_SkipsExisting(t *testing.T) { for range 20 { taken[generateName()] = true } - name := generateUniqueName(func(n string) bool { return taken[n] }) + name := GenerateUniqueName(func(n string) bool { return taken[n] }) if taken[name] { t.Fatalf("returned existing name: %q", name) } @@ -56,7 +56,7 @@ func TestGenerateUniqueName_SkipsExisting(t *testing.T) { func TestGenerateUniqueName_FallbackSuffix(t *testing.T) { calls := 0 - name := generateUniqueName(func(n string) bool { + name := GenerateUniqueName(func(n string) bool { calls++ return calls <= 10 }) diff --git a/internal/orchestrator/client.go b/internal/orchestrator/client.go new file mode 100644 index 0000000..6a0a7e8 --- /dev/null +++ b/internal/orchestrator/client.go @@ -0,0 +1,632 @@ +package orchestrator + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "sync" + + "open-sandbox/internal/database" + "open-sandbox/internal/docker" + "open-sandbox/models" +) + +// RemoteDockerClient implements the DockerClient interface by forwarding +// operations to worker nodes over HTTP. +type RemoteDockerClient struct { + registry *WorkerRegistry + repo *database.Repository + http *http.Client +} + +// NewRemoteClient creates a RemoteDockerClient. +func NewRemoteClient(registry *WorkerRegistry, repo *database.Repository) *RemoteDockerClient { + return &RemoteDockerClient{ + registry: registry, + repo: repo, + http: &http.Client{}, + } +} + +// --- Helpers --- + +// workerForSandbox looks up which worker owns a sandbox. +func (c *RemoteDockerClient) workerForSandbox(sandboxID string) (WorkerInfo, error) { + sb, err := c.repo.FindByID(sandboxID) + if err != nil { + return WorkerInfo{}, err + } + if sb == nil { + return WorkerInfo{}, docker.ErrNotFound + } + if sb.WorkerID == "" { + return WorkerInfo{}, fmt.Errorf("sandbox %s has no worker assigned", sandboxID) + } + return c.registry.Lookup(sb.WorkerID) +} + +// doJSON sends a JSON request to a worker and decodes the response. +func (c *RemoteDockerClient) doJSON(ctx context.Context, w WorkerInfo, method, path string, body any, out any) error { + var bodyReader io.Reader + if body != nil { + b, err := json.Marshal(body) + if err != nil { + return err + } + bodyReader = bytes.NewReader(b) + } + + req, err := http.NewRequestWithContext(ctx, method, w.URL+"/internal/v1"+path, bodyReader) + if err != nil { + return err + } + req.Header.Set("X-Worker-Key", w.APIKey) + if body != nil { + req.Header.Set("Content-Type", "application/json") + } + + resp, err := c.http.Do(req) + if err != nil { + return fmt.Errorf("worker %s: %w", w.ID, err) + } + defer resp.Body.Close() + + if resp.StatusCode >= 400 { + return c.mapWorkerError(resp) + } + + if out != nil { + return json.NewDecoder(resp.Body).Decode(out) + } + return nil +} + +// doNoContent sends a request expecting 204 No Content. +func (c *RemoteDockerClient) doNoContent(ctx context.Context, w WorkerInfo, method, path string, body any) error { + var bodyReader io.Reader + if body != nil { + b, err := json.Marshal(body) + if err != nil { + return err + } + bodyReader = bytes.NewReader(b) + } + + req, err := http.NewRequestWithContext(ctx, method, w.URL+"/internal/v1"+path, bodyReader) + if err != nil { + return err + } + req.Header.Set("X-Worker-Key", w.APIKey) + if body != nil { + req.Header.Set("Content-Type", "application/json") + } + + resp, err := c.http.Do(req) + if err != nil { + return fmt.Errorf("worker %s: %w", w.ID, err) + } + defer resp.Body.Close() + + if resp.StatusCode >= 400 { + return c.mapWorkerError(resp) + } + return nil +} + +// doStream opens a streaming connection to a worker and returns the response body. +func (c *RemoteDockerClient) doStream(ctx context.Context, w WorkerInfo, method, path string) (*http.Response, error) { + req, err := http.NewRequestWithContext(ctx, method, w.URL+"/internal/v1"+path, nil) + if err != nil { + return nil, err + } + req.Header.Set("X-Worker-Key", w.APIKey) + + resp, err := c.http.Do(req) + if err != nil { + return nil, fmt.Errorf("worker %s: %w", w.ID, err) + } + if resp.StatusCode >= 400 { + defer resp.Body.Close() + return nil, c.mapWorkerError(resp) + } + return resp, nil +} + +// mapWorkerError converts a worker error response to a sentinel error. +func (c *RemoteDockerClient) mapWorkerError(resp *http.Response) error { + var errResp struct { + Code string `json:"code"` + Message string `json:"message"` + } + json.NewDecoder(resp.Body).Decode(&errResp) + + switch resp.StatusCode { + case http.StatusNotFound: + if errResp.Message == "command not found" { + return docker.ErrCommandNotFound + } + return docker.ErrNotFound + case http.StatusConflict: + // Map back to sentinel errors based on message. + switch errResp.Message { + case docker.ErrAlreadyRunning.Error(): + return docker.ErrAlreadyRunning + case docker.ErrAlreadyStopped.Error(): + return docker.ErrAlreadyStopped + case docker.ErrAlreadyPaused.Error(): + return docker.ErrAlreadyPaused + case docker.ErrNotPaused.Error(): + return docker.ErrNotPaused + case docker.ErrNotRunning.Error(): + return docker.ErrNotRunning + case docker.ErrCommandFinished.Error(): + return docker.ErrCommandFinished + } + return fmt.Errorf("conflict: %s", errResp.Message) + case http.StatusBadRequest: + if errResp.Message == "image not found locally" { + return docker.ErrImageNotFound + } + return fmt.Errorf("bad request: %s", errResp.Message) + default: + return fmt.Errorf("worker error %d: %s", resp.StatusCode, errResp.Message) + } +} + +// --- DockerClient implementation --- + +func (c *RemoteDockerClient) Ping(ctx context.Context) error { + workers := c.registry.All() + if len(workers) == 0 { + return ErrNoWorkers + } + // Ping first worker as health indicator. + return c.doJSON(ctx, workers[0], http.MethodGet, "/health", nil, nil) +} + +func (c *RemoteDockerClient) List(ctx context.Context) ([]models.SandboxSummary, error) { + workers := c.registry.All() + if len(workers) == 0 { + return []models.SandboxSummary{}, nil + } + + type result struct { + items []models.SandboxSummary + err error + } + + ch := make(chan result, len(workers)) + for _, w := range workers { + go func(w WorkerInfo) { + var resp struct { + Sandboxes []models.SandboxSummary `json:"sandboxes"` + } + err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes", nil, &resp) + ch <- result{items: resp.Sandboxes, err: err} + }(w) + } + + var all []models.SandboxSummary + for range workers { + r := <-ch + if r.err != nil { + log.Printf("orchestrator: list from worker failed: %v", r.err) + continue + } + all = append(all, r.items...) + } + return all, nil +} + +func (c *RemoteDockerClient) Create(ctx context.Context, req models.CreateSandboxRequest) (models.CreateSandboxResponse, error) { + // Pick a worker via round-robin. + w, err := c.registry.Next() + if err != nil { + return models.CreateSandboxResponse{}, err + } + + // Pre-generate a unique name. + name := docker.GenerateUniqueName(func(n string) bool { + sb, _ := c.repo.FindByName(n) + return sb != nil + }) + req.Name = name + + var resp models.CreateSandboxResponse + if err := c.doJSON(ctx, w, http.MethodPost, "/sandboxes", req, &resp); err != nil { + return models.CreateSandboxResponse{}, err + } + + // Persist sandbox in orchestrator DB with worker_id. + if err := c.repo.Save(database.Sandbox{ + ID: resp.ID, + Name: resp.Name, + Image: req.Image, + Ports: database.JSONMap(portsToMap(resp.Ports)), + Port: mainPort(req.Ports), + WorkerID: w.ID, + }); err != nil { + log.Printf("orchestrator: failed to persist sandbox %s: %v", resp.ID, err) + } + + return resp, nil +} + +func (c *RemoteDockerClient) Inspect(ctx context.Context, id string) (models.SandboxDetail, error) { + w, err := c.workerForSandbox(id) + if err != nil { + return models.SandboxDetail{}, err + } + var resp models.SandboxDetail + if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+id, nil, &resp); err != nil { + return models.SandboxDetail{}, err + } + return resp, nil +} + +func (c *RemoteDockerClient) Start(ctx context.Context, id string) (models.RestartResponse, error) { + w, err := c.workerForSandbox(id) + if err != nil { + return models.RestartResponse{}, err + } + var resp models.RestartResponse + if err := c.doJSON(ctx, w, http.MethodPost, "/sandboxes/"+id+"/start", nil, &resp); err != nil { + return models.RestartResponse{}, err + } + // Update ports in DB. + c.repo.UpdatePorts(id, database.JSONMap(portsToMap(resp.Ports))) + return resp, nil +} + +func (c *RemoteDockerClient) Stop(ctx context.Context, id string) error { + w, err := c.workerForSandbox(id) + if err != nil { + return err + } + return c.doJSON(ctx, w, http.MethodPost, "/sandboxes/"+id+"/stop", nil, nil) +} + +func (c *RemoteDockerClient) Restart(ctx context.Context, id string) (models.RestartResponse, error) { + w, err := c.workerForSandbox(id) + if err != nil { + return models.RestartResponse{}, err + } + var resp models.RestartResponse + if err := c.doJSON(ctx, w, http.MethodPost, "/sandboxes/"+id+"/restart", nil, &resp); err != nil { + return models.RestartResponse{}, err + } + // Update ports in DB. + c.repo.UpdatePorts(id, database.JSONMap(portsToMap(resp.Ports))) + return resp, nil +} + +func (c *RemoteDockerClient) Remove(ctx context.Context, id string) error { + w, err := c.workerForSandbox(id) + if err != nil { + return err + } + if err := c.doNoContent(ctx, w, http.MethodDelete, "/sandboxes/"+id, nil); err != nil { + return err + } + // Clean up DB. + c.repo.DeleteCommandsBySandbox(id) + c.repo.Delete(id) + return nil +} + +func (c *RemoteDockerClient) Pause(ctx context.Context, id string) error { + w, err := c.workerForSandbox(id) + if err != nil { + return err + } + return c.doJSON(ctx, w, http.MethodPost, "/sandboxes/"+id+"/pause", nil, nil) +} + +func (c *RemoteDockerClient) Resume(ctx context.Context, id string) error { + w, err := c.workerForSandbox(id) + if err != nil { + return err + } + return c.doJSON(ctx, w, http.MethodPost, "/sandboxes/"+id+"/resume", nil, nil) +} + +func (c *RemoteDockerClient) RenewExpiration(ctx context.Context, id string, timeout int) error { + w, err := c.workerForSandbox(id) + if err != nil { + return err + } + return c.doJSON(ctx, w, http.MethodPost, "/sandboxes/"+id+"/renew-expiration", + models.RenewExpirationRequest{Timeout: timeout}, nil) +} + +func (c *RemoteDockerClient) Stats(ctx context.Context, id string) (models.SandboxStats, error) { + w, err := c.workerForSandbox(id) + if err != nil { + return models.SandboxStats{}, err + } + var resp models.SandboxStats + if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+id+"/stats", nil, &resp); err != nil { + return models.SandboxStats{}, err + } + return resp, nil +} + +// --- Commands --- + +func (c *RemoteDockerClient) ExecCommand(ctx context.Context, sandboxID string, req models.ExecCommandRequest) (models.CommandDetail, error) { + w, err := c.workerForSandbox(sandboxID) + if err != nil { + return models.CommandDetail{}, err + } + var resp models.CommandResponse + if err := c.doJSON(ctx, w, http.MethodPost, "/sandboxes/"+sandboxID+"/cmd", req, &resp); err != nil { + return models.CommandDetail{}, err + } + return resp.Command, nil +} + +func (c *RemoteDockerClient) GetCommand(ctx context.Context, sandboxID, cmdID string) (models.CommandDetail, error) { + w, err := c.workerForSandbox(sandboxID) + if err != nil { + return models.CommandDetail{}, err + } + var resp models.CommandResponse + if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+sandboxID+"/cmd/"+cmdID, nil, &resp); err != nil { + return models.CommandDetail{}, err + } + return resp.Command, nil +} + +func (c *RemoteDockerClient) ListCommands(ctx context.Context, sandboxID string) ([]models.CommandDetail, error) { + w, err := c.workerForSandbox(sandboxID) + if err != nil { + return nil, err + } + var resp models.CommandListResponse + if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+sandboxID+"/cmd", nil, &resp); err != nil { + return nil, err + } + return resp.Commands, nil +} + +func (c *RemoteDockerClient) KillCommand(ctx context.Context, sandboxID, cmdID string, signal int) (models.CommandDetail, error) { + w, err := c.workerForSandbox(sandboxID) + if err != nil { + return models.CommandDetail{}, err + } + var resp models.CommandResponse + if err := c.doJSON(ctx, w, http.MethodPost, "/sandboxes/"+sandboxID+"/cmd/"+cmdID+"/kill", + models.KillCommandRequest{Signal: signal}, &resp); err != nil { + return models.CommandDetail{}, err + } + return resp.Command, nil +} + +func (c *RemoteDockerClient) StreamCommandLogs(ctx context.Context, sandboxID, cmdID string) (io.ReadCloser, io.ReadCloser, error) { + w, err := c.workerForSandbox(sandboxID) + if err != nil { + return nil, nil, err + } + resp, err := c.doStream(ctx, w, http.MethodGet, "/sandboxes/"+sandboxID+"/cmd/"+cmdID+"/logs?stream=true") + if err != nil { + return nil, nil, err + } + // Worker streams interleaved ND-JSON. Return the body as a single reader for both. + // The public API handler will pipe it directly to the client. + return resp.Body, io.NopCloser(bytes.NewReader(nil)), nil +} + +func (c *RemoteDockerClient) GetCommandLogs(ctx context.Context, sandboxID, cmdID string) (models.CommandLogsResponse, error) { + w, err := c.workerForSandbox(sandboxID) + if err != nil { + return models.CommandLogsResponse{}, err + } + var resp models.CommandLogsResponse + if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+sandboxID+"/cmd/"+cmdID+"/logs", nil, &resp); err != nil { + return models.CommandLogsResponse{}, err + } + return resp, nil +} + +func (c *RemoteDockerClient) WaitCommand(ctx context.Context, sandboxID, cmdID string) (models.CommandDetail, error) { + w, err := c.workerForSandbox(sandboxID) + if err != nil { + return models.CommandDetail{}, err + } + // Use ?wait=true to block on worker side. + resp, err := c.doStream(ctx, w, http.MethodGet, "/sandboxes/"+sandboxID+"/cmd/"+cmdID+"?wait=true") + if err != nil { + return models.CommandDetail{}, err + } + defer resp.Body.Close() + + // Read ND-JSON lines; the last one has the final state. + var last models.CommandResponse + dec := json.NewDecoder(resp.Body) + for dec.More() { + if err := dec.Decode(&last); err != nil { + break + } + } + return last.Command, nil +} + +// --- Files --- + +func (c *RemoteDockerClient) ReadFile(ctx context.Context, id, path string) (string, error) { + w, err := c.workerForSandbox(id) + if err != nil { + return "", err + } + var resp models.FileReadResponse + if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+id+"/files?path="+path, nil, &resp); err != nil { + return "", err + } + return resp.Content, nil +} + +func (c *RemoteDockerClient) WriteFile(ctx context.Context, id, path, content string) error { + w, err := c.workerForSandbox(id) + if err != nil { + return err + } + return c.doJSON(ctx, w, http.MethodPut, "/sandboxes/"+id+"/files?path="+path, + models.FileWriteRequest{Content: content}, nil) +} + +func (c *RemoteDockerClient) DeleteFile(ctx context.Context, id, path string) error { + w, err := c.workerForSandbox(id) + if err != nil { + return err + } + return c.doNoContent(ctx, w, http.MethodDelete, "/sandboxes/"+id+"/files?path="+path, nil) +} + +func (c *RemoteDockerClient) ListDir(ctx context.Context, id, path string) (string, error) { + w, err := c.workerForSandbox(id) + if err != nil { + return "", err + } + var resp models.FileListResponse + if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+id+"/files/list?path="+path, nil, &resp); err != nil { + return "", err + } + return resp.Output, nil +} + +// --- Images --- + +func (c *RemoteDockerClient) PullImage(ctx context.Context, image string) error { + workers := c.registry.All() + if len(workers) == 0 { + return ErrNoWorkers + } + + // Pull on all workers in parallel. + var wg sync.WaitGroup + errs := make(chan error, len(workers)) + for _, w := range workers { + wg.Add(1) + go func(w WorkerInfo) { + defer wg.Done() + errs <- c.doJSON(ctx, w, http.MethodPost, "/images/pull", + models.ImagePullRequest{Image: image}, nil) + }(w) + } + wg.Wait() + close(errs) + + // Return first error if any. + for err := range errs { + if err != nil { + return err + } + } + return nil +} + +func (c *RemoteDockerClient) RemoveImage(ctx context.Context, id string, force bool) error { + workers := c.registry.All() + if len(workers) == 0 { + return ErrNoWorkers + } + + forceParam := "" + if force { + forceParam = "?force=true" + } + + var wg sync.WaitGroup + errs := make(chan error, len(workers)) + for _, w := range workers { + wg.Add(1) + go func(w WorkerInfo) { + defer wg.Done() + errs <- c.doNoContent(ctx, w, http.MethodDelete, "/images/"+id+forceParam, nil) + }(w) + } + wg.Wait() + close(errs) + + for err := range errs { + if err != nil { + return err + } + } + return nil +} + +func (c *RemoteDockerClient) InspectImage(ctx context.Context, id string) (models.ImageDetail, error) { + workers := c.registry.All() + for _, w := range workers { + var resp models.ImageDetail + if err := c.doJSON(ctx, w, http.MethodGet, "/images/"+id, nil, &resp); err == nil { + return resp, nil + } + } + return models.ImageDetail{}, docker.ErrImageNotFound +} + +func (c *RemoteDockerClient) ListImages(ctx context.Context) ([]models.ImageSummary, error) { + workers := c.registry.All() + if len(workers) == 0 { + return []models.ImageSummary{}, nil + } + + type result struct { + items []models.ImageSummary + err error + } + ch := make(chan result, len(workers)) + for _, w := range workers { + go func(w WorkerInfo) { + var resp struct { + Images []models.ImageSummary `json:"images"` + } + err := c.doJSON(ctx, w, http.MethodGet, "/images", nil, &resp) + ch <- result{items: resp.Images, err: err} + }(w) + } + + // Merge and deduplicate by ID. + seen := make(map[string]bool) + var all []models.ImageSummary + for range workers { + r := <-ch + if r.err != nil { + log.Printf("orchestrator: list images from worker failed: %v", r.err) + continue + } + for _, img := range r.items { + if !seen[img.ID] { + seen[img.ID] = true + all = append(all, img) + } + } + } + return all, nil +} + +// --- Helpers --- + +// portsToMap converts ["3000/tcp", "8080/tcp"] to a map for DB storage. +// The orchestrator doesn't know host ports at this point; just stores the keys. +func portsToMap(ports []string) map[string]string { + m := make(map[string]string, len(ports)) + for _, p := range ports { + m[p] = "" + } + return m +} + +// mainPort returns the first port from a list, or empty. +func mainPort(ports []string) string { + if len(ports) > 0 { + return ports[0] + } + return "" +} diff --git a/internal/orchestrator/handler.go b/internal/orchestrator/handler.go new file mode 100644 index 0000000..dc9b2c5 --- /dev/null +++ b/internal/orchestrator/handler.go @@ -0,0 +1,69 @@ +package orchestrator + +import ( + "net/http" + + "github.com/gin-gonic/gin" +) + +// Handler exposes orchestrator-specific HTTP endpoints (worker registration). +type Handler struct { + registry *WorkerRegistry +} + +// NewHandler creates an orchestrator Handler. +func NewHandler(registry *WorkerRegistry) *Handler { + return &Handler{registry: registry} +} + +type registerRequest struct { + URL string `json:"url" binding:"required"` +} + +type registerResponse struct { + WorkerID string `json:"worker_id"` +} + +// RegisterWorker handles POST /internal/v1/workers/register. +func (h *Handler) RegisterWorker(c *gin.Context) { + var req registerRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"code": "BAD_REQUEST", "message": err.Error()}) + return + } + + // Use the API key from the request header as the worker's auth key. + apiKey := c.GetHeader("X-Worker-Key") + + id, err := h.registry.Register(req.URL, apiKey) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"code": "INTERNAL_ERROR", "message": err.Error()}) + return + } + + c.JSON(http.StatusCreated, registerResponse{WorkerID: id}) +} + +// DeregisterWorker handles DELETE /internal/v1/workers/:id. +func (h *Handler) DeregisterWorker(c *gin.Context) { + id := c.Param("id") + if err := h.registry.Deregister(id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"code": "INTERNAL_ERROR", "message": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"status": "deregistered"}) +} + +// ListWorkers handles GET /internal/v1/workers. +func (h *Handler) ListWorkers(c *gin.Context) { + workers := h.registry.All() + c.JSON(http.StatusOK, gin.H{"workers": workers}) +} + +// RegisterRoutes attaches worker management endpoints. +func (h *Handler) RegisterRoutes(g *gin.RouterGroup) { + w := g.Group("/workers") + w.POST("/register", h.RegisterWorker) + w.GET("", h.ListWorkers) + w.DELETE("/:id", h.DeregisterWorker) +} diff --git a/internal/orchestrator/registry.go b/internal/orchestrator/registry.go new file mode 100644 index 0000000..c14a7b5 --- /dev/null +++ b/internal/orchestrator/registry.go @@ -0,0 +1,145 @@ +package orchestrator + +import ( + "crypto/rand" + "encoding/hex" + "errors" + "sync" + "sync/atomic" + "time" + + "open-sandbox/internal/database" +) + +var ErrNoWorkers = errors.New("no active workers available") + +// WorkerInfo holds in-memory worker metadata for scheduling. +type WorkerInfo struct { + ID string + URL string + APIKey string +} + +// WorkerRegistry manages worker registration and round-robin scheduling. +type WorkerRegistry struct { + repo *database.Repository + mu sync.RWMutex + workers []WorkerInfo // cached active workers + counter atomic.Uint64 +} + +// NewRegistry creates a WorkerRegistry and loads active workers from DB. +func NewRegistry(repo *database.Repository) *WorkerRegistry { + r := &WorkerRegistry{repo: repo} + r.reload() + return r +} + +// Register adds a new worker or re-activates an existing one. +// Returns the worker ID. +func (r *WorkerRegistry) Register(url, apiKey string) (string, error) { + // Check if worker with this URL already exists. + existing, err := r.repo.FindWorkerByURL(url) + if err != nil { + return "", err + } + + if existing != nil { + // Re-activate existing worker. + existing.Status = "active" + existing.APIKey = apiKey + if err := r.repo.SaveWorker(*existing); err != nil { + return "", err + } + r.reload() + return existing.ID, nil + } + + // Create new worker. + id := generateWorkerID() + w := database.Worker{ + ID: id, + URL: url, + APIKey: apiKey, + Status: "active", + CreatedAt: time.Now().UnixMilli(), + } + if err := r.repo.SaveWorker(w); err != nil { + return "", err + } + r.reload() + return id, nil +} + +// Deregister marks a worker as inactive and reloads the cache. +func (r *WorkerRegistry) Deregister(id string) error { + if err := r.repo.UpdateWorkerStatus(id, "inactive"); err != nil { + return err + } + r.reload() + return nil +} + +// Next returns the next worker via round-robin. +func (r *WorkerRegistry) Next() (WorkerInfo, error) { + r.mu.RLock() + defer r.mu.RUnlock() + + if len(r.workers) == 0 { + return WorkerInfo{}, ErrNoWorkers + } + + idx := r.counter.Add(1) - 1 + w := r.workers[idx%uint64(len(r.workers))] + return w, nil +} + +// Lookup returns the worker info for a given worker ID. +func (r *WorkerRegistry) Lookup(workerID string) (WorkerInfo, error) { + r.mu.RLock() + defer r.mu.RUnlock() + + for _, w := range r.workers { + if w.ID == workerID { + return w, nil + } + } + return WorkerInfo{}, ErrNoWorkers +} + +// All returns all active workers. +func (r *WorkerRegistry) All() []WorkerInfo { + r.mu.RLock() + defer r.mu.RUnlock() + + out := make([]WorkerInfo, len(r.workers)) + copy(out, r.workers) + return out +} + +// reload refreshes the in-memory worker list from the database. +func (r *WorkerRegistry) reload() { + workers, err := r.repo.FindActiveWorkers() + if err != nil { + return + } + + infos := make([]WorkerInfo, 0, len(workers)) + for _, w := range workers { + infos = append(infos, WorkerInfo{ + ID: w.ID, + URL: w.URL, + APIKey: w.APIKey, + }) + } + + r.mu.Lock() + r.workers = infos + r.mu.Unlock() +} + +func generateWorkerID() string { + b := make([]byte, 8) + rand.Read(b) + return "wrk_" + hex.EncodeToString(b) +} diff --git a/internal/proxy/resolver.go b/internal/proxy/resolver.go index beefc31..a916a6b 100644 --- a/internal/proxy/resolver.go +++ b/internal/proxy/resolver.go @@ -7,7 +7,9 @@ import ( "open-sandbox/internal/database" ) -// resolve looks up the sandbox by name and returns the target URL (http://127.0.0.1:{hostPort}). +// resolve looks up the sandbox by name and returns the target URL. +// In all-in-one mode: http://127.0.0.1:{hostPort} +// In orchestrator mode: http://{workerIP}:{hostPort} func (s *Server) resolve(name string) (*url.URL, error) { // Check cache first. if target, ok := s.cache.get(name); ok { @@ -29,15 +31,43 @@ func (s *Server) resolve(name string) (*url.URL, error) { return nil, err } + // Determine the host: worker IP if assigned, otherwise 127.0.0.1. + host := "127.0.0.1" + if sb.WorkerID != "" { + w, wErr := s.repo.FindWorkerByID(sb.WorkerID) + if wErr == nil && w != nil { + if ip := extractHost(w.URL); ip != "" { + host = ip + } + } + } + target := &url.URL{ Scheme: "http", - Host: "127.0.0.1:" + hostPort, + Host: host + ":" + hostPort, } s.cache.set(name, target) return target, nil } +// extractHost extracts the host (IP or hostname) from a URL string. +// "http://10.0.0.2:9090" → "10.0.0.2" +func extractHost(rawURL string) string { + u, err := url.Parse(rawURL) + if err != nil { + return "" + } + h := u.Hostname() + // Don't use 0.0.0.0 as a target — it means "all interfaces" on the worker, + // but we need the actual worker IP. In practice the orchestrator should have + // the real IP from the worker's registration URL. + if h == "0.0.0.0" { + return "" + } + return h +} + // resolveHostPort returns the Docker-assigned host port for the sandbox's port. // If Port is not set but there is exactly one port in the map, it uses that. func resolveHostPort(sb *database.Sandbox) (string, error) { From 406912eea85027c3fe0f318cdfc315fd40608e9e Mon Sep 17 00:00:00 2001 From: Nicolas Vargas Date: Mon, 23 Feb 2026 22:46:31 -0500 Subject: [PATCH 4/6] Feat(core): Add worker-scoped image listing and resilient client --- internal/api/handler.go | 24 ++++++- internal/orchestrator/client.go | 113 ++++++++++++++++++++++---------- 2 files changed, 100 insertions(+), 37 deletions(-) diff --git a/internal/api/handler.go b/internal/api/handler.go index 9f86355..dc946aa 100644 --- a/internal/api/handler.go +++ b/internal/api/handler.go @@ -2,6 +2,7 @@ package api import ( "bufio" + "context" "encoding/json" "fmt" "io" @@ -748,17 +749,36 @@ func (h *Handler) getImage(c *gin.Context) { c.JSON(http.StatusOK, detail) } +// WorkerImageLister is an optional interface for listing images from a specific worker. +type WorkerImageLister interface { + ListImagesFromWorkers(ctx context.Context, workerID string) ([]models.ImageSummary, error) +} + // listImages handles GET /v1/images. // @Summary List local images -// @Description Returns all Docker images available locally. +// @Description Returns all Docker images available locally. In distributed mode, use ?worker_id= to filter by worker. // @Tags images // @Produce json +// @Param worker_id query string false "Filter images by worker ID (distributed mode only)" // @Success 200 {object} map[string]interface{} "List of images" // @Failure 500 {object} ErrorResponse // @Security ApiKeyAuth // @Router /images [get] func (h *Handler) listImages(c *gin.Context) { - images, err := h.docker.ListImages(c.Request.Context()) + var images []models.ImageSummary + var err error + + // If ?worker_id= is set and the client supports it, filter by worker. + if workerID := c.Query("worker_id"); workerID != "" { + if lister, ok := h.docker.(WorkerImageLister); ok { + images, err = lister.ListImagesFromWorkers(c.Request.Context(), workerID) + } else { + images, err = h.docker.ListImages(c.Request.Context()) + } + } else { + images, err = h.docker.ListImages(c.Request.Context()) + } + if err != nil { internalError(c, err) return diff --git a/internal/orchestrator/client.go b/internal/orchestrator/client.go index 6a0a7e8..c59b43e 100644 --- a/internal/orchestrator/client.go +++ b/internal/orchestrator/client.go @@ -8,7 +8,7 @@ import ( "io" "log" "net/http" - "sync" + "net/url" "open-sandbox/internal/database" "open-sandbox/internal/docker" @@ -463,7 +463,7 @@ func (c *RemoteDockerClient) ReadFile(ctx context.Context, id, path string) (str return "", err } var resp models.FileReadResponse - if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+id+"/files?path="+path, nil, &resp); err != nil { + if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+id+"/files?path="+url.QueryEscape(path), nil, &resp); err != nil { return "", err } return resp.Content, nil @@ -474,7 +474,7 @@ func (c *RemoteDockerClient) WriteFile(ctx context.Context, id, path, content st if err != nil { return err } - return c.doJSON(ctx, w, http.MethodPut, "/sandboxes/"+id+"/files?path="+path, + return c.doJSON(ctx, w, http.MethodPut, "/sandboxes/"+id+"/files?path="+url.QueryEscape(path), models.FileWriteRequest{Content: content}, nil) } @@ -483,7 +483,7 @@ func (c *RemoteDockerClient) DeleteFile(ctx context.Context, id, path string) er if err != nil { return err } - return c.doNoContent(ctx, w, http.MethodDelete, "/sandboxes/"+id+"/files?path="+path, nil) + return c.doNoContent(ctx, w, http.MethodDelete, "/sandboxes/"+id+"/files?path="+url.QueryEscape(path), nil) } func (c *RemoteDockerClient) ListDir(ctx context.Context, id, path string) (string, error) { @@ -492,7 +492,7 @@ func (c *RemoteDockerClient) ListDir(ctx context.Context, id, path string) (stri return "", err } var resp models.FileListResponse - if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+id+"/files/list?path="+path, nil, &resp); err != nil { + if err := c.doJSON(ctx, w, http.MethodGet, "/sandboxes/"+id+"/files/list?path="+url.QueryEscape(path), nil, &resp); err != nil { return "", err } return resp.Output, nil @@ -506,26 +506,40 @@ func (c *RemoteDockerClient) PullImage(ctx context.Context, image string) error return ErrNoWorkers } - // Pull on all workers in parallel. - var wg sync.WaitGroup - errs := make(chan error, len(workers)) + // Pull on all workers in parallel. Collect per-worker errors. + type pullResult struct { + workerID string + err error + } + ch := make(chan pullResult, len(workers)) for _, w := range workers { - wg.Add(1) go func(w WorkerInfo) { - defer wg.Done() - errs <- c.doJSON(ctx, w, http.MethodPost, "/images/pull", + err := c.doJSON(ctx, w, http.MethodPost, "/images/pull", models.ImagePullRequest{Image: image}, nil) + ch <- pullResult{workerID: w.ID, err: err} }(w) } - wg.Wait() - close(errs) - // Return first error if any. - for err := range errs { - if err != nil { - return err + var firstErr error + var failed int + for range workers { + r := <-ch + if r.err != nil { + failed++ + log.Printf("orchestrator: pull %s on worker %s failed: %v", image, r.workerID, r.err) + if firstErr == nil { + firstErr = r.err + } } } + + // Fail only if ALL workers failed. + if failed == len(workers) { + return firstErr + } + if failed > 0 { + log.Printf("orchestrator: pull %s succeeded on %d/%d workers", image, len(workers)-failed, len(workers)) + } return nil } @@ -535,45 +549,74 @@ func (c *RemoteDockerClient) RemoveImage(ctx context.Context, id string, force b return ErrNoWorkers } - forceParam := "" + path := "/images/" + url.PathEscape(id) if force { - forceParam = "?force=true" + path += "?force=true" } - var wg sync.WaitGroup - errs := make(chan error, len(workers)) + type removeResult struct { + workerID string + err error + } + ch := make(chan removeResult, len(workers)) for _, w := range workers { - wg.Add(1) go func(w WorkerInfo) { - defer wg.Done() - errs <- c.doNoContent(ctx, w, http.MethodDelete, "/images/"+id+forceParam, nil) + err := c.doNoContent(ctx, w, http.MethodDelete, path, nil) + ch <- removeResult{workerID: w.ID, err: err} }(w) } - wg.Wait() - close(errs) - for err := range errs { - if err != nil { - return err + var firstErr error + var failed int + for range workers { + r := <-ch + if r.err != nil { + failed++ + log.Printf("orchestrator: remove image %s on worker %s failed: %v", id, r.workerID, r.err) + if firstErr == nil { + firstErr = r.err + } } } + + if failed == len(workers) { + return firstErr + } return nil } func (c *RemoteDockerClient) InspectImage(ctx context.Context, id string) (models.ImageDetail, error) { workers := c.registry.All() + path := "/images/" + url.PathEscape(id) for _, w := range workers { var resp models.ImageDetail - if err := c.doJSON(ctx, w, http.MethodGet, "/images/"+id, nil, &resp); err == nil { + if err := c.doJSON(ctx, w, http.MethodGet, path, nil, &resp); err == nil { return resp, nil } } return models.ImageDetail{}, docker.ErrImageNotFound } +// ListImages returns images from all workers (or a single worker if workerID is specified). +// Results are merged and deduplicated by image ID. func (c *RemoteDockerClient) ListImages(ctx context.Context) ([]models.ImageSummary, error) { - workers := c.registry.All() - if len(workers) == 0 { + return c.ListImagesFromWorkers(ctx, "") +} + +// ListImagesFromWorkers lists images from a specific worker or all workers. +func (c *RemoteDockerClient) ListImagesFromWorkers(ctx context.Context, workerID string) ([]models.ImageSummary, error) { + var targets []WorkerInfo + if workerID != "" { + w, err := c.registry.Lookup(workerID) + if err != nil { + return nil, err + } + targets = []WorkerInfo{w} + } else { + targets = c.registry.All() + } + + if len(targets) == 0 { return []models.ImageSummary{}, nil } @@ -581,8 +624,8 @@ func (c *RemoteDockerClient) ListImages(ctx context.Context) ([]models.ImageSumm items []models.ImageSummary err error } - ch := make(chan result, len(workers)) - for _, w := range workers { + ch := make(chan result, len(targets)) + for _, w := range targets { go func(w WorkerInfo) { var resp struct { Images []models.ImageSummary `json:"images"` @@ -595,7 +638,7 @@ func (c *RemoteDockerClient) ListImages(ctx context.Context) ([]models.ImageSumm // Merge and deduplicate by ID. seen := make(map[string]bool) var all []models.ImageSummary - for range workers { + for range targets { r := <-ch if r.err != nil { log.Printf("orchestrator: list images from worker failed: %v", r.err) From 75257c1a9ca0ca6c3c4db743897c346bc27d1014 Mon Sep 17 00:00:00 2001 From: Nicolas Vargas Date: Mon, 23 Feb 2026 22:53:18 -0500 Subject: [PATCH 5/6] Feat(testing): Add tests for orchestrator and worker packages --- internal/orchestrator/client_test.go | 465 ++++++++++++++++++++ internal/orchestrator/handler_test.go | 96 +++++ internal/orchestrator/registry_test.go | 105 +++++ internal/worker/handler_test.go | 571 +++++++++++++++++++++++++ internal/worker/middleware_test.go | 82 ++++ 5 files changed, 1319 insertions(+) create mode 100644 internal/orchestrator/client_test.go create mode 100644 internal/orchestrator/handler_test.go create mode 100644 internal/orchestrator/registry_test.go create mode 100644 internal/worker/handler_test.go create mode 100644 internal/worker/middleware_test.go diff --git a/internal/orchestrator/client_test.go b/internal/orchestrator/client_test.go new file mode 100644 index 0000000..05efef8 --- /dev/null +++ b/internal/orchestrator/client_test.go @@ -0,0 +1,465 @@ +package orchestrator_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "open-sandbox/internal/database" + "open-sandbox/internal/orchestrator" + "open-sandbox/models" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// mockWorker creates an httptest.Server that mimics the worker internal API. +func mockWorker(t *testing.T) *httptest.Server { + mux := http.NewServeMux() + + mux.HandleFunc("GET /internal/v1/health", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]string{"status": "healthy"}) + }) + + mux.HandleFunc("POST /internal/v1/sandboxes", func(w http.ResponseWriter, r *http.Request) { + var req models.CreateSandboxRequest + json.NewDecoder(r.Body).Decode(&req) + w.WriteHeader(http.StatusCreated) + json.NewEncoder(w).Encode(models.CreateSandboxResponse{ + ID: "abc123", + Name: req.Name, + Ports: []string{"3000/tcp"}, + }) + }) + + mux.HandleFunc("GET /internal/v1/sandboxes/{id}", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.SandboxDetail{ + ID: r.PathValue("id"), + Name: "test-sandbox", + Image: "node:24", + Status: "running", + Running: true, + }) + }) + + mux.HandleFunc("DELETE /internal/v1/sandboxes/{id}", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNoContent) + }) + + mux.HandleFunc("POST /internal/v1/sandboxes/{id}/stop", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]string{"status": "stopped"}) + }) + + mux.HandleFunc("POST /internal/v1/sandboxes/{id}/start", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.RestartResponse{ + Status: "started", + Ports: []string{"3000/tcp"}, + }) + }) + + mux.HandleFunc("POST /internal/v1/sandboxes/{id}/restart", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.RestartResponse{ + Status: "restarted", + Ports: []string{"3000/tcp"}, + }) + }) + + mux.HandleFunc("POST /internal/v1/sandboxes/{id}/pause", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]string{"status": "paused"}) + }) + + mux.HandleFunc("POST /internal/v1/sandboxes/{id}/resume", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]string{"status": "resumed"}) + }) + + mux.HandleFunc("GET /internal/v1/sandboxes/{id}/stats", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.SandboxStats{CPU: 1.5, PIDs: 10}) + }) + + mux.HandleFunc("GET /internal/v1/sandboxes/{id}/cmd", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.CommandListResponse{Commands: []models.CommandDetail{}}) + }) + + mux.HandleFunc("POST /internal/v1/sandboxes/{id}/cmd", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.CommandResponse{ + Command: models.CommandDetail{ID: "cmd_abc", Name: "ls", SandboxID: r.PathValue("id")}, + }) + }) + + mux.HandleFunc("GET /internal/v1/sandboxes/{id}/files", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.FileReadResponse{ + Path: r.URL.Query().Get("path"), Content: "hello", + }) + }) + + mux.HandleFunc("PUT /internal/v1/sandboxes/{id}/files", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]string{"status": "written"}) + }) + + mux.HandleFunc("DELETE /internal/v1/sandboxes/{id}/files", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNoContent) + }) + + mux.HandleFunc("GET /internal/v1/sandboxes/{id}/files/list", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.FileListResponse{ + Path: r.URL.Query().Get("path"), Output: "drwxr-xr-x 2 root root 4096 Jan 1 00:00 .", + }) + }) + + mux.HandleFunc("GET /internal/v1/images", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "images": []models.ImageSummary{ + {ID: "sha256:abc", Tags: []string{"node:24"}, Size: 100}, + }, + }) + }) + + mux.HandleFunc("GET /internal/v1/images/{id}", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.ImageDetail{ + ID: "sha256:abc", Tags: []string{"node:24"}, Size: 100, + }) + }) + + mux.HandleFunc("POST /internal/v1/images/pull", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(models.ImagePullResponse{Status: "pulled"}) + }) + + mux.HandleFunc("DELETE /internal/v1/images/{id}", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNoContent) + }) + + return httptest.NewServer(mux) +} + +// setupClient creates a RemoteDockerClient with a mock worker registered. +func setupClient(t *testing.T, workerURL string) (*orchestrator.RemoteDockerClient, *database.Repository) { + db := database.New(":memory:") + repo := database.NewRepository(db) + reg := orchestrator.NewRegistry(repo) + + // Register mock worker. + workerID, err := reg.Register(workerURL, "test-key") + require.NoError(t, err) + + // Pre-create a sandbox record so workerForSandbox lookups work. + repo.Save(database.Sandbox{ + ID: "abc123", + Name: "test-sandbox", + Image: "node:24", + WorkerID: workerID, + }) + + client := orchestrator.NewRemoteClient(reg, repo) + return client, repo +} + +func TestRemoteClient_Ping(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + err := client.Ping(context.Background()) + assert.NoError(t, err) +} + +func TestRemoteClient_Create(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, repo := setupClient(t, srv.URL) + + resp, err := client.Create(context.Background(), models.CreateSandboxRequest{ + Image: "node:24", + Ports: []string{"3000"}, + }) + require.NoError(t, err) + assert.Equal(t, "abc123", resp.ID) + assert.NotEmpty(t, resp.Name) + + // Verify persisted in orchestrator DB. + sb, err := repo.FindByID("abc123") + require.NoError(t, err) + require.NotNil(t, sb) + assert.NotEmpty(t, sb.WorkerID) +} + +func TestRemoteClient_Inspect(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + detail, err := client.Inspect(context.Background(), "abc123") + require.NoError(t, err) + assert.Equal(t, "abc123", detail.ID) + assert.True(t, detail.Running) +} + +func TestRemoteClient_InspectNotFound(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + _, err := client.Inspect(context.Background(), "nonexistent") + assert.Error(t, err) +} + +func TestRemoteClient_Stop(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + err := client.Stop(context.Background(), "abc123") + assert.NoError(t, err) +} + +func TestRemoteClient_Start(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + resp, err := client.Start(context.Background(), "abc123") + require.NoError(t, err) + assert.Equal(t, "started", resp.Status) +} + +func TestRemoteClient_Restart(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + resp, err := client.Restart(context.Background(), "abc123") + require.NoError(t, err) + assert.Equal(t, "restarted", resp.Status) +} + +func TestRemoteClient_Remove(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, repo := setupClient(t, srv.URL) + + err := client.Remove(context.Background(), "abc123") + require.NoError(t, err) + + // Verify cleaned up from DB. + sb, _ := repo.FindByID("abc123") + assert.Nil(t, sb) +} + +func TestRemoteClient_Pause(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + err := client.Pause(context.Background(), "abc123") + assert.NoError(t, err) +} + +func TestRemoteClient_Resume(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + err := client.Resume(context.Background(), "abc123") + assert.NoError(t, err) +} + +func TestRemoteClient_Stats(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + stats, err := client.Stats(context.Background(), "abc123") + require.NoError(t, err) + assert.Equal(t, 1.5, stats.CPU) + assert.Equal(t, uint64(10), stats.PIDs) +} + +func TestRemoteClient_ExecCommand(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + cmd, err := client.ExecCommand(context.Background(), "abc123", models.ExecCommandRequest{ + Command: "ls", + }) + require.NoError(t, err) + assert.Equal(t, "cmd_abc", cmd.ID) +} + +func TestRemoteClient_ListCommands(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + cmds, err := client.ListCommands(context.Background(), "abc123") + require.NoError(t, err) + assert.NotNil(t, cmds) +} + +func TestRemoteClient_ReadFile(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + content, err := client.ReadFile(context.Background(), "abc123", "/app/index.js") + require.NoError(t, err) + assert.Equal(t, "hello", content) +} + +func TestRemoteClient_WriteFile(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + err := client.WriteFile(context.Background(), "abc123", "/app/index.js", "console.log('hi')") + assert.NoError(t, err) +} + +func TestRemoteClient_DeleteFile(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + err := client.DeleteFile(context.Background(), "abc123", "/app/index.js") + assert.NoError(t, err) +} + +func TestRemoteClient_ListDir(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + output, err := client.ListDir(context.Background(), "abc123", "/app") + require.NoError(t, err) + assert.Contains(t, output, "drwxr-xr-x") +} + +func TestRemoteClient_PullImage(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + err := client.PullImage(context.Background(), "node:24") + assert.NoError(t, err) +} + +func TestRemoteClient_RemoveImage(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + err := client.RemoveImage(context.Background(), "sha256:abc", false) + assert.NoError(t, err) +} + +func TestRemoteClient_InspectImage(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + detail, err := client.InspectImage(context.Background(), "sha256:abc") + require.NoError(t, err) + assert.Equal(t, "sha256:abc", detail.ID) +} + +func TestRemoteClient_ListImages(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + client, _ := setupClient(t, srv.URL) + + images, err := client.ListImages(context.Background()) + require.NoError(t, err) + assert.Len(t, images, 1) + assert.Equal(t, "sha256:abc", images[0].ID) +} + +func TestRemoteClient_ListImagesFromWorker(t *testing.T) { + srv := mockWorker(t) + defer srv.Close() + db := database.New(":memory:") + repo := database.NewRepository(db) + reg := orchestrator.NewRegistry(repo) + workerID, _ := reg.Register(srv.URL, "key") + client := orchestrator.NewRemoteClient(reg, repo) + + images, err := client.ListImagesFromWorkers(context.Background(), workerID) + require.NoError(t, err) + assert.Len(t, images, 1) +} + +func TestRemoteClient_ListImagesDeduplicates(t *testing.T) { + // Two workers returning the same image. + srv1 := mockWorker(t) + defer srv1.Close() + srv2 := mockWorker(t) + defer srv2.Close() + + db := database.New(":memory:") + repo := database.NewRepository(db) + reg := orchestrator.NewRegistry(repo) + reg.Register(srv1.URL, "k1") + reg.Register(srv2.URL, "k2") + client := orchestrator.NewRemoteClient(reg, repo) + + images, err := client.ListImages(context.Background()) + require.NoError(t, err) + // Both workers return sha256:abc, should be deduplicated. + assert.Len(t, images, 1) +} + +func TestRemoteClient_PullImage_PartialFailure(t *testing.T) { + // One working worker, one broken worker. + srv := mockWorker(t) + defer srv.Close() + + brokenSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + json.NewEncoder(w).Encode(map[string]string{"code": "INTERNAL_ERROR", "message": "disk full"}) + })) + defer brokenSrv.Close() + + db := database.New(":memory:") + repo := database.NewRepository(db) + reg := orchestrator.NewRegistry(repo) + reg.Register(srv.URL, "k1") + reg.Register(brokenSrv.URL, "k2") + client := orchestrator.NewRemoteClient(reg, repo) + + // Should succeed because at least one worker succeeded. + err := client.PullImage(context.Background(), "node:24") + assert.NoError(t, err) +} + +func TestRemoteClient_PullImage_AllFail(t *testing.T) { + brokenSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + json.NewEncoder(w).Encode(map[string]string{"code": "INTERNAL_ERROR", "message": "disk full"}) + })) + defer brokenSrv.Close() + + db := database.New(":memory:") + repo := database.NewRepository(db) + reg := orchestrator.NewRegistry(repo) + reg.Register(brokenSrv.URL, "k1") + client := orchestrator.NewRemoteClient(reg, repo) + + err := client.PullImage(context.Background(), "node:24") + assert.Error(t, err) +} + +func TestRemoteClient_NoWorkers(t *testing.T) { + db := database.New(":memory:") + repo := database.NewRepository(db) + reg := orchestrator.NewRegistry(repo) + client := orchestrator.NewRemoteClient(reg, repo) + + err := client.Ping(context.Background()) + assert.ErrorIs(t, err, orchestrator.ErrNoWorkers) + + _, err = client.Create(context.Background(), models.CreateSandboxRequest{Image: "node:24"}) + assert.ErrorIs(t, err, orchestrator.ErrNoWorkers) + + err = client.PullImage(context.Background(), "node:24") + assert.ErrorIs(t, err, orchestrator.ErrNoWorkers) +} diff --git a/internal/orchestrator/handler_test.go b/internal/orchestrator/handler_test.go new file mode 100644 index 0000000..c2cb2da --- /dev/null +++ b/internal/orchestrator/handler_test.go @@ -0,0 +1,96 @@ +package orchestrator_test + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "open-sandbox/internal/database" + "open-sandbox/internal/orchestrator" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func orchestratorRouter() (*gin.Engine, *orchestrator.WorkerRegistry) { + gin.SetMode(gin.TestMode) + db := database.New(":memory:") + repo := database.NewRepository(db) + reg := orchestrator.NewRegistry(repo) + h := orchestrator.NewHandler(reg) + + r := gin.New() + g := r.Group("/internal/v1") + h.RegisterRoutes(g) + return r, reg +} + +func TestHandler_RegisterWorker(t *testing.T) { + r, _ := orchestratorRouter() + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/internal/v1/workers/register", + strings.NewReader(`{"url":"http://10.0.0.1:9090"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Worker-Key", "secret") + r.ServeHTTP(w, req) + + assert.Equal(t, http.StatusCreated, w.Code) + + var resp struct { + WorkerID string `json:"worker_id"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.NotEmpty(t, resp.WorkerID) + assert.Contains(t, resp.WorkerID, "wrk_") +} + +func TestHandler_RegisterWorker_MissingURL(t *testing.T) { + r, _ := orchestratorRouter() + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/internal/v1/workers/register", + strings.NewReader(`{}`)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) +} + +func TestHandler_ListWorkers(t *testing.T) { + r, reg := orchestratorRouter() + + reg.Register("http://w1:9090", "k1") + reg.Register("http://w2:9090", "k2") + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/internal/v1/workers", nil) + r.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + + var resp struct { + Workers []struct { + ID string `json:"ID"` + URL string `json:"URL"` + } `json:"workers"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.Len(t, resp.Workers, 2) +} + +func TestHandler_DeregisterWorker(t *testing.T) { + r, reg := orchestratorRouter() + + id, _ := reg.Register("http://w1:9090", "k1") + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodDelete, "/internal/v1/workers/"+id, nil) + r.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Len(t, reg.All(), 0) +} diff --git a/internal/orchestrator/registry_test.go b/internal/orchestrator/registry_test.go new file mode 100644 index 0000000..9ec7b45 --- /dev/null +++ b/internal/orchestrator/registry_test.go @@ -0,0 +1,105 @@ +package orchestrator_test + +import ( + "testing" + + "open-sandbox/internal/database" + "open-sandbox/internal/orchestrator" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newTestRegistry() *orchestrator.WorkerRegistry { + db := database.New(":memory:") + repo := database.NewRepository(db) + return orchestrator.NewRegistry(repo) +} + +func TestRegistry_RegisterAndAll(t *testing.T) { + reg := newTestRegistry() + + id1, err := reg.Register("http://10.0.0.1:9090", "key1") + require.NoError(t, err) + assert.NotEmpty(t, id1) + assert.Contains(t, id1, "wrk_") + + id2, err := reg.Register("http://10.0.0.2:9090", "key2") + require.NoError(t, err) + assert.NotEqual(t, id1, id2) + + all := reg.All() + assert.Len(t, all, 2) +} + +func TestRegistry_Deregister(t *testing.T) { + reg := newTestRegistry() + + id, err := reg.Register("http://10.0.0.1:9090", "key") + require.NoError(t, err) + + assert.Len(t, reg.All(), 1) + + require.NoError(t, reg.Deregister(id)) + assert.Len(t, reg.All(), 0) +} + +func TestRegistry_ReActivate(t *testing.T) { + reg := newTestRegistry() + + id1, err := reg.Register("http://10.0.0.1:9090", "key") + require.NoError(t, err) + + require.NoError(t, reg.Deregister(id1)) + assert.Len(t, reg.All(), 0) + + // Re-register same URL should re-activate, return same ID. + id2, err := reg.Register("http://10.0.0.1:9090", "new-key") + require.NoError(t, err) + assert.Equal(t, id1, id2) + assert.Len(t, reg.All(), 1) +} + +func TestRegistry_RoundRobin(t *testing.T) { + reg := newTestRegistry() + + reg.Register("http://w1:9090", "k1") + reg.Register("http://w2:9090", "k2") + reg.Register("http://w3:9090", "k3") + + // Collect 9 calls — should distribute evenly. + counts := map[string]int{} + for range 9 { + w, err := reg.Next() + require.NoError(t, err) + counts[w.URL]++ + } + + assert.Equal(t, 3, counts["http://w1:9090"]) + assert.Equal(t, 3, counts["http://w2:9090"]) + assert.Equal(t, 3, counts["http://w3:9090"]) +} + +func TestRegistry_NextNoWorkers(t *testing.T) { + reg := newTestRegistry() + + _, err := reg.Next() + assert.ErrorIs(t, err, orchestrator.ErrNoWorkers) +} + +func TestRegistry_Lookup(t *testing.T) { + reg := newTestRegistry() + + id, _ := reg.Register("http://10.0.0.1:9090", "key") + + w, err := reg.Lookup(id) + require.NoError(t, err) + assert.Equal(t, "http://10.0.0.1:9090", w.URL) +} + +func TestRegistry_LookupNotFound(t *testing.T) { + reg := newTestRegistry() + + _, err := reg.Lookup("wrk_nonexistent") + assert.ErrorIs(t, err, orchestrator.ErrNoWorkers) +} diff --git a/internal/worker/handler_test.go b/internal/worker/handler_test.go new file mode 100644 index 0000000..1bd963b --- /dev/null +++ b/internal/worker/handler_test.go @@ -0,0 +1,571 @@ +package worker + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "open-sandbox/internal/docker" + "open-sandbox/models" +) + +// stubDocker implements DockerClient for testing. +type stubDocker struct { + pingErr error + listResult []models.SandboxSummary + createResult models.CreateSandboxResponse + createErr error + inspectResult models.SandboxDetail + inspectErr error + startResult models.RestartResponse + startErr error + stopErr error + restartResult models.RestartResponse + restartErr error + removeErr error + pauseErr error + resumeErr error + renewErr error + statsResult models.SandboxStats + statsErr error + execResult models.CommandDetail + execErr error + getCmd models.CommandDetail + getCmdErr error + listCmds []models.CommandDetail + listCmdsErr error + killResult models.CommandDetail + killErr error + logsResult models.CommandLogsResponse + logsErr error + waitResult models.CommandDetail + waitErr error + readResult string + readErr error + writeErr error + deleteFileErr error + listDirResult string + listDirErr error + pullErr error + removeImgErr error + inspectImg models.ImageDetail + inspectImgErr error + listImgs []models.ImageSummary + listImgsErr error +} + +func (s *stubDocker) Ping(ctx context.Context) error { return s.pingErr } +func (s *stubDocker) List(ctx context.Context) ([]models.SandboxSummary, error) { + return s.listResult, nil +} +func (s *stubDocker) Create(ctx context.Context, req models.CreateSandboxRequest) (models.CreateSandboxResponse, error) { + return s.createResult, s.createErr +} +func (s *stubDocker) Inspect(ctx context.Context, id string) (models.SandboxDetail, error) { + return s.inspectResult, s.inspectErr +} +func (s *stubDocker) Start(ctx context.Context, id string) (models.RestartResponse, error) { + return s.startResult, s.startErr +} +func (s *stubDocker) Stop(ctx context.Context, id string) error { return s.stopErr } +func (s *stubDocker) Restart(ctx context.Context, id string) (models.RestartResponse, error) { + return s.restartResult, s.restartErr +} +func (s *stubDocker) Remove(ctx context.Context, id string) error { return s.removeErr } +func (s *stubDocker) Pause(ctx context.Context, id string) error { return s.pauseErr } +func (s *stubDocker) Resume(ctx context.Context, id string) error { return s.resumeErr } +func (s *stubDocker) RenewExpiration(ctx context.Context, id string, timeout int) error { + return s.renewErr +} +func (s *stubDocker) ExecCommand(ctx context.Context, sandboxID string, req models.ExecCommandRequest) (models.CommandDetail, error) { + return s.execResult, s.execErr +} +func (s *stubDocker) GetCommand(ctx context.Context, sandboxID, cmdID string) (models.CommandDetail, error) { + return s.getCmd, s.getCmdErr +} +func (s *stubDocker) ListCommands(ctx context.Context, sandboxID string) ([]models.CommandDetail, error) { + return s.listCmds, s.listCmdsErr +} +func (s *stubDocker) KillCommand(ctx context.Context, sandboxID, cmdID string, signal int) (models.CommandDetail, error) { + return s.killResult, s.killErr +} +func (s *stubDocker) StreamCommandLogs(ctx context.Context, sandboxID, cmdID string) (io.ReadCloser, io.ReadCloser, error) { + stdout := io.NopCloser(strings.NewReader("hello\n")) + stderr := io.NopCloser(strings.NewReader("")) + return stdout, stderr, s.logsErr +} +func (s *stubDocker) GetCommandLogs(ctx context.Context, sandboxID, cmdID string) (models.CommandLogsResponse, error) { + return s.logsResult, s.logsErr +} +func (s *stubDocker) WaitCommand(ctx context.Context, sandboxID, cmdID string) (models.CommandDetail, error) { + return s.waitResult, s.waitErr +} +func (s *stubDocker) Stats(ctx context.Context, id string) (models.SandboxStats, error) { + return s.statsResult, s.statsErr +} +func (s *stubDocker) ReadFile(ctx context.Context, id, path string) (string, error) { + return s.readResult, s.readErr +} +func (s *stubDocker) WriteFile(ctx context.Context, id, path, content string) error { + return s.writeErr +} +func (s *stubDocker) DeleteFile(ctx context.Context, id, path string) error { + return s.deleteFileErr +} +func (s *stubDocker) ListDir(ctx context.Context, id, path string) (string, error) { + return s.listDirResult, s.listDirErr +} +func (s *stubDocker) PullImage(ctx context.Context, image string) error { return s.pullErr } +func (s *stubDocker) RemoveImage(ctx context.Context, id string, force bool) error { + return s.removeImgErr +} +func (s *stubDocker) InspectImage(ctx context.Context, id string) (models.ImageDetail, error) { + return s.inspectImg, s.inspectImgErr +} +func (s *stubDocker) ListImages(ctx context.Context) ([]models.ImageSummary, error) { + return s.listImgs, s.listImgsErr +} + +func setupRouter(stub *stubDocker) *gin.Engine { + gin.SetMode(gin.TestMode) + r := gin.New() + h := NewHandler(stub) + g := r.Group("/internal/v1") + h.RegisterRoutes(g) + return r +} + +// --- Health --- + +func TestHealth_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/health", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestHealth_Unhealthy(t *testing.T) { + r := setupRouter(&stubDocker{pingErr: docker.ErrNotRunning}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/health", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusServiceUnavailable { + t.Fatalf("expected 503, got %d", w.Code) + } +} + +// --- Sandbox CRUD --- + +func TestCreateSandbox_OK(t *testing.T) { + stub := &stubDocker{ + createResult: models.CreateSandboxResponse{ID: "abc123", Name: "test-box"}, + } + r := setupRouter(stub) + body, _ := json.Marshal(models.CreateSandboxRequest{Image: "node:24"}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusCreated { + t.Fatalf("expected 201, got %d: %s", w.Code, w.Body.String()) + } +} + +func TestCreateSandbox_BadRequest(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes", strings.NewReader("{}")) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", w.Code) + } +} + +func TestInspectSandbox_OK(t *testing.T) { + stub := &stubDocker{ + inspectResult: models.SandboxDetail{ID: "abc", Name: "test"}, + } + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/abc", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestInspectSandbox_NotFound(t *testing.T) { + stub := &stubDocker{inspectErr: docker.ErrNotFound} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/missing", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusNotFound { + t.Fatalf("expected 404, got %d", w.Code) + } +} + +func TestDeleteSandbox_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("DELETE", "/internal/v1/sandboxes/abc", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusNoContent { + t.Fatalf("expected 204, got %d", w.Code) + } +} + +func TestDeleteSandbox_NotFound(t *testing.T) { + stub := &stubDocker{removeErr: docker.ErrNotFound} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("DELETE", "/internal/v1/sandboxes/missing", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusNotFound { + t.Fatalf("expected 404, got %d", w.Code) + } +} + +// --- Lifecycle --- + +func TestStartSandbox_OK(t *testing.T) { + r := setupRouter(&stubDocker{startResult: models.RestartResponse{Status: "started"}}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/start", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestStartSandbox_AlreadyRunning(t *testing.T) { + r := setupRouter(&stubDocker{startErr: docker.ErrAlreadyRunning}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/start", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusConflict { + t.Fatalf("expected 409, got %d", w.Code) + } +} + +func TestStopSandbox_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/stop", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestStopSandbox_AlreadyStopped(t *testing.T) { + r := setupRouter(&stubDocker{stopErr: docker.ErrAlreadyStopped}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/stop", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusConflict { + t.Fatalf("expected 409, got %d", w.Code) + } +} + +func TestPauseSandbox_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/pause", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestResumeSandbox_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/resume", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestRestartSandbox_OK(t *testing.T) { + r := setupRouter(&stubDocker{restartResult: models.RestartResponse{Status: "restarted"}}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/restart", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +// --- Stats --- + +func TestStats_OK(t *testing.T) { + stub := &stubDocker{statsResult: models.SandboxStats{CPU: 12.5}} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/abc/stats", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +// --- Commands --- + +func TestExecCommand_OK(t *testing.T) { + stub := &stubDocker{execResult: models.CommandDetail{ID: "cmd_1"}} + r := setupRouter(stub) + body, _ := json.Marshal(models.ExecCommandRequest{Command: "ls"}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/cmd", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } +} + +func TestExecCommand_NotRunning(t *testing.T) { + stub := &stubDocker{execErr: docker.ErrNotRunning} + r := setupRouter(stub) + body, _ := json.Marshal(models.ExecCommandRequest{Command: "ls"}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/cmd", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusConflict { + t.Fatalf("expected 409, got %d", w.Code) + } +} + +func TestListCommands_OK(t *testing.T) { + stub := &stubDocker{listCmds: []models.CommandDetail{{ID: "cmd_1"}}} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/abc/cmd", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestKillCommand_OK(t *testing.T) { + stub := &stubDocker{killResult: models.CommandDetail{ID: "cmd_1"}} + r := setupRouter(stub) + body, _ := json.Marshal(models.KillCommandRequest{Signal: 15}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/cmd/cmd_1/kill", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestKillCommand_AlreadyFinished(t *testing.T) { + stub := &stubDocker{killErr: docker.ErrCommandFinished} + r := setupRouter(stub) + body, _ := json.Marshal(models.KillCommandRequest{Signal: 9}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/cmd/cmd_1/kill", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusConflict { + t.Fatalf("expected 409, got %d", w.Code) + } +} + +func TestGetCommandLogs_OK(t *testing.T) { + stub := &stubDocker{logsResult: models.CommandLogsResponse{Stdout: "output"}} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/abc/cmd/cmd_1/logs", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +// --- Files --- + +func TestReadFile_OK(t *testing.T) { + stub := &stubDocker{readResult: "file content"} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/abc/files?path=/app/main.go", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestReadFile_MissingPath(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/abc/files", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", w.Code) + } +} + +func TestWriteFile_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + body, _ := json.Marshal(models.FileWriteRequest{Content: "hello"}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("PUT", "/internal/v1/sandboxes/abc/files?path=/app/main.go", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestWriteFile_MissingPath(t *testing.T) { + r := setupRouter(&stubDocker{}) + body, _ := json.Marshal(models.FileWriteRequest{Content: "hello"}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("PUT", "/internal/v1/sandboxes/abc/files", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", w.Code) + } +} + +func TestDeleteFile_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("DELETE", "/internal/v1/sandboxes/abc/files?path=/app/tmp.txt", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusNoContent { + t.Fatalf("expected 204, got %d", w.Code) + } +} + +func TestDeleteFile_MissingPath(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("DELETE", "/internal/v1/sandboxes/abc/files", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", w.Code) + } +} + +func TestListDir_OK(t *testing.T) { + stub := &stubDocker{listDirResult: "file1\nfile2\n"} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/abc/files/list", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +// --- Images --- + +func TestListImages_OK(t *testing.T) { + stub := &stubDocker{listImgs: []models.ImageSummary{{ID: "sha256:abc"}}} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/images", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestInspectImage_OK(t *testing.T) { + stub := &stubDocker{inspectImg: models.ImageDetail{ID: "sha256:abc"}} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/images/sha256:abc", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestInspectImage_NotFound(t *testing.T) { + stub := &stubDocker{inspectImgErr: docker.ErrImageNotFound} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/images/missing", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", w.Code) + } +} + +func TestPullImage_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + body, _ := json.Marshal(models.ImagePullRequest{Image: "node:24"}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/images/pull", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } +} + +func TestDeleteImage_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("DELETE", "/internal/v1/images/sha256:abc", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusNoContent { + t.Fatalf("expected 204, got %d", w.Code) + } +} + +// --- Error mapping --- + +func TestMapError_Timeout(t *testing.T) { + stub := &stubDocker{inspectErr: context.DeadlineExceeded} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/abc", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusRequestTimeout { + t.Fatalf("expected 408, got %d", w.Code) + } +} + +func TestMapError_CommandNotFound(t *testing.T) { + stub := &stubDocker{getCmdErr: docker.ErrCommandNotFound} + r := setupRouter(stub) + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/internal/v1/sandboxes/abc/cmd/missing", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusNotFound { + t.Fatalf("expected 404, got %d", w.Code) + } +} + +// --- Renew Expiration --- + +func TestRenewExpiration_OK(t *testing.T) { + r := setupRouter(&stubDocker{}) + body, _ := json.Marshal(models.RenewExpirationRequest{Timeout: 600}) + w := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/internal/v1/sandboxes/abc/renew-expiration", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } +} diff --git a/internal/worker/middleware_test.go b/internal/worker/middleware_test.go new file mode 100644 index 0000000..78bd47a --- /dev/null +++ b/internal/worker/middleware_test.go @@ -0,0 +1,82 @@ +package worker + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func init() { + gin.SetMode(gin.TestMode) +} + +func setupAuthRouter(key string) *gin.Engine { + r := gin.New() + r.Use(APIKeyAuth(key)) + r.GET("/test", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + return r +} + +func TestAPIKeyAuth_ValidKey(t *testing.T) { + r := setupAuthRouter("secret-key") + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/test", nil) + req.Header.Set("X-Worker-Key", "secret-key") + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } +} + +func TestAPIKeyAuth_MissingHeader(t *testing.T) { + r := setupAuthRouter("secret-key") + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/test", nil) + r.ServeHTTP(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", w.Code) + } +} + +func TestAPIKeyAuth_WrongKey(t *testing.T) { + r := setupAuthRouter("secret-key") + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/test", nil) + req.Header.Set("X-Worker-Key", "wrong-key") + r.ServeHTTP(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", w.Code) + } +} + +func TestAPIKeyAuth_WhitespaceKey(t *testing.T) { + r := setupAuthRouter("secret-key") + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/test", nil) + req.Header.Set("X-Worker-Key", " secret-key ") + r.ServeHTTP(w, req) + + // trimmed key should match + if w.Code != http.StatusOK { + t.Fatalf("expected 200 (trimmed), got %d", w.Code) + } +} + +func TestAPIKeyAuth_EmptyKey(t *testing.T) { + r := setupAuthRouter("secret-key") + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/test", nil) + req.Header.Set("X-Worker-Key", "") + r.ServeHTTP(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d", w.Code) + } +} From 9afc109faf30b70d0ce698d526e75725cdf21f96 Mon Sep 17 00:00:00 2001 From: Nicolas Vargas Date: Mon, 23 Feb 2026 23:21:11 -0500 Subject: [PATCH 6/6] Feat(docs) --- docs/docs.go | 14 +++++++++++++- docs/swagger.json | 14 +++++++++++++- docs/swagger.yaml | 11 ++++++++++- 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/docs/docs.go b/docs/docs.go index bb28aa3..3508c9e 100644 --- a/docs/docs.go +++ b/docs/docs.go @@ -54,7 +54,7 @@ const docTemplate = `{ "ApiKeyAuth": [] } ], - "description": "Returns all Docker images available locally.", + "description": "Returns all Docker images available locally. In distributed mode, use ?worker_id= to filter by worker.", "produces": [ "application/json" ], @@ -62,6 +62,14 @@ const docTemplate = `{ "images" ], "summary": "List local images", + "parameters": [ + { + "type": "string", + "description": "Filter images by worker ID (distributed mode only)", + "name": "worker_id", + "in": "query" + } + ], "responses": { "200": { "description": "List of images", @@ -1392,6 +1400,10 @@ const docTemplate = `{ "type": "string", "example": "node:24" }, + "name": { + "description": "pre-generated name (used by orchestrator → worker). Empty = auto-generate.", + "type": "string" + }, "ports": { "description": "container ports to expose, e.g. [\"3000\", \"8080/tcp\"]. First port is the default for proxy routing.", "type": "array", diff --git a/docs/swagger.json b/docs/swagger.json index a944e0e..8e99a72 100644 --- a/docs/swagger.json +++ b/docs/swagger.json @@ -48,7 +48,7 @@ "ApiKeyAuth": [] } ], - "description": "Returns all Docker images available locally.", + "description": "Returns all Docker images available locally. In distributed mode, use ?worker_id= to filter by worker.", "produces": [ "application/json" ], @@ -56,6 +56,14 @@ "images" ], "summary": "List local images", + "parameters": [ + { + "type": "string", + "description": "Filter images by worker ID (distributed mode only)", + "name": "worker_id", + "in": "query" + } + ], "responses": { "200": { "description": "List of images", @@ -1386,6 +1394,10 @@ "type": "string", "example": "node:24" }, + "name": { + "description": "pre-generated name (used by orchestrator → worker). Empty = auto-generate.", + "type": "string" + }, "ports": { "description": "container ports to expose, e.g. [\"3000\", \"8080/tcp\"]. First port is the default for proxy routing.", "type": "array", diff --git a/docs/swagger.yaml b/docs/swagger.yaml index a201f2c..c2d004c 100644 --- a/docs/swagger.yaml +++ b/docs/swagger.yaml @@ -72,6 +72,9 @@ definitions: image: example: node:24 type: string + name: + description: pre-generated name (used by orchestrator → worker). Empty = auto-generate. + type: string ports: description: container ports to expose, e.g. ["3000", "8080/tcp"]. First port is the default for proxy routing. @@ -323,7 +326,13 @@ paths: - system /images: get: - description: Returns all Docker images available locally. + description: Returns all Docker images available locally. In distributed mode, + use ?worker_id= to filter by worker. + parameters: + - description: Filter images by worker ID (distributed mode only) + in: query + name: worker_id + type: string produces: - application/json responses: