diff --git a/dify-agent-runtime/Makefile b/dify-agent-runtime/Makefile index 64550ef16ce..40cf168c18a 100644 --- a/dify-agent-runtime/Makefile +++ b/dify-agent-runtime/Makefile @@ -4,7 +4,7 @@ PROTO_SRC := ../dify-agent/proto/dify/agent/stub/v1/agent_stub.proto PROTO_OUT := gen AGENT_CLI_HELP_JSON := ../dify-agent/src/dify_agent/layers/_agent_cli_help.json -build: $(BIN_DIR)/shellctl $(BIN_DIR)/shellctl-sanitize-pty $(BIN_DIR)/shellctl-runner-exit $(BIN_DIR)/dify-agent gen-cli-help +build: $(BIN_DIR)/shellctl $(BIN_DIR)/shellctl-sanitize-pty $(BIN_DIR)/shellctl-runner-exit $(BIN_DIR)/shellctl-runner $(BIN_DIR)/dify-agent gen-cli-help $(BIN_DIR)/shellctl: $(shell find cmd/shellctl internal -name '*.go') go build -o $@ ./cmd/shellctl @@ -15,6 +15,10 @@ $(BIN_DIR)/shellctl-sanitize-pty: $(shell find cmd/sanitize-pty internal/sanitiz $(BIN_DIR)/shellctl-runner-exit: $(shell find cmd/runner-exit internal/runner_exit -name '*.go') go build -o $@ ./cmd/runner-exit +$(BIN_DIR)/shellctl-runner: $(shell find cmd/runner internal/landlock internal/cmdutil -name '*.go') + go build -o $@ ./cmd/runner + + $(BIN_DIR)/dify-agent: $(shell find cmd/dify-agent-cli internal/agentcli internal/stubclient -name '*.go') $(shell find gen -name '*.go' 2>/dev/null) go build -o $@ ./cmd/dify-agent-cli @@ -83,12 +87,39 @@ integration-up: fi; \ sleep 2; \ done + $(eval CONTAINER_NAME_NOISO := sandbox-rt-noiso-$(TEST_ID)) + $(eval HOST_PORT_NOISO := $(shell python3 -c 'import socket; s=socket.socket(); s.bind(("",0)); print(s.getsockname()[1]); s.close()')) + $(eval AUTH_TOKEN_NOISO := test-token-noiso-$(TEST_ID)) + @echo "Starting no-isolation container $(CONTAINER_NAME_NOISO) on port $(HOST_PORT_NOISO)..." + docker run -d --name $(CONTAINER_NAME_NOISO) \ + -p $(HOST_PORT_NOISO):5004 \ + -e SHELLCTL_AUTH_TOKEN=$(AUTH_TOKEN_NOISO) \ + -e ENABLE_PATH_ISOLATION=false \ + $(IMAGE_NAME) + @echo 'CONTAINER_NAME_NOISO=$(CONTAINER_NAME_NOISO)' >> $(STATE_FILE) + @echo 'HOST_PORT_NOISO=$(HOST_PORT_NOISO)' >> $(STATE_FILE) + @echo 'AUTH_TOKEN_NOISO=$(AUTH_TOKEN_NOISO)' >> $(STATE_FILE) + @for i in $$(seq 1 30); do \ + if curl -sf http://localhost:$(HOST_PORT_NOISO)/healthz > /dev/null 2>&1; then \ + echo "No-isolation runtime is ready on port $(HOST_PORT_NOISO)"; \ + break; \ + fi; \ + if [ "$$i" -eq 30 ]; then \ + echo "ERROR: no-isolation runtime not ready after 60s" >&2; \ + docker logs $(CONTAINER_NAME_NOISO); \ + docker rm -f $(CONTAINER_NAME_NOISO) 2>/dev/null; \ + exit 1; \ + fi; \ + sleep 2; \ + done integration-test: @test -f $(STATE_FILE) || { echo "ERROR: run 'make integration-up' first" >&2; exit 1; } @. ./$(STATE_FILE); \ SHELLCTL_GO_URL=http://localhost:$$HOST_PORT \ SHELLCTL_TEST_TOKEN=$$AUTH_TOKEN \ + SHELLCTL_GO_URL_NO_ISOLATION=http://localhost:$$HOST_PORT_NOISO \ + SHELLCTL_TEST_TOKEN_NO_ISOLATION=$$AUTH_TOKEN_NOISO \ go test -tags=integration -v -count=1 -timeout=300s ./tests/... integration-logs: @@ -99,6 +130,7 @@ integration-down: @if [ -f $(STATE_FILE) ]; then \ . ./$(STATE_FILE); \ docker rm -f $$CONTAINER_NAME 2>/dev/null || true; \ + docker rm -f $$CONTAINER_NAME_NOISO 2>/dev/null || true; \ rm -f $(STATE_FILE); \ fi diff --git a/dify-agent-runtime/README.md b/dify-agent-runtime/README.md index 6b96ca76435..74e796a1d5c 100644 --- a/dify-agent-runtime/README.md +++ b/dify-agent-runtime/README.md @@ -25,11 +25,12 @@ internal/ make build ``` -Produces three binaries in `bin/`: +Produces binaries in `bin/`: - `shellctl` — the main server (`shellctl serve --listen 0.0.0.0:5004`) - `shellctl-sanitize-pty` — PTY sanitizer for tmux pipe-pane - `shellctl-runner-exit` — exit state writer +- `shellctl-runner` — job runner with integrated Landlock isolation ### Building docker image @@ -62,6 +63,30 @@ make gen-cli-help make test ``` +## Path Isolation + +Each agent job runs inside a Landlock sandbox that restricts filesystem access: + +| Access | Paths (defaults) | +|--------|------------------| +| **Read-Write** | `$HOME` (always, includes `$CWD/.tmp` as `TMPDIR`) | +| **Read-Write (dev)** | `/dev/null`, `/dev/zero`, `/dev/urandom`, `/dev/random`, `/dev/tty` | +| **Read-Only + Exec** | `/usr`, `/bin`, `/sbin`, `/lib`, `/lib64`, `/etc`, `/proc`, `/opt/dify-agent-tools`, `/opt/homebrew`, `/snap` | +| **Denied** | Everything else (`/tmp`, other agents' homes, `/var`, `/srv`, etc.) | + +The runner automatically creates `$CWD/.tmp` and sets `TMPDIR`, `TMP`, `TEMP` to it, so temp files stay isolated per workspace. + +### Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `ENABLE_PATH_ISOLATION` | `true` | Set to `false` to disable Landlock entirely | +| `LANDLOCK_RW_PATHS` | *(empty)* | Comma-separated RW directories (besides `$HOME`) | +| `LANDLOCK_RO_PATHS` | `/usr,/bin,...` | Comma-separated RO+exec directories | +| `LANDLOCK_RW_DEV_PATHS` | `/dev/null,...` | Comma-separated device files with RW access | + +Requires Linux ≥ 5.13. On unsupported kernels, a warning is printed to stderr. + ## Dependencies - Go 1.23+ diff --git a/dify-agent-runtime/cmd/runner-exit/main.go b/dify-agent-runtime/cmd/runner-exit/main.go index b1eb2344a8b..4464dc86a05 100644 --- a/dify-agent-runtime/cmd/runner-exit/main.go +++ b/dify-agent-runtime/cmd/runner-exit/main.go @@ -5,8 +5,8 @@ package main import ( "flag" "fmt" - "os" + "github.com/langgenius/dify/dify-agent-runtime/internal/cmdutil" runnerexit "github.com/langgenius/dify/dify-agent-runtime/internal/runner_exit" ) @@ -19,12 +19,11 @@ func main() { flag.Parse() if *stateDir == "" || *jobID == "" || *endedAt == "" { - fmt.Fprintf(os.Stderr, "shellctl-runner-exit: --state-dir, --job-id, and --ended-at are required\n") - os.Exit(1) + cmdutil.HandleError(fmt.Errorf("missing flags"), 1, "--state-dir, --job-id, and --ended-at are required") } - if err := runnerexit.RecordRunnerExit(*stateDir, *jobID, *exitCode, *endedAt, *busyTimeoutMs); err != nil { - fmt.Fprintf(os.Stderr, "shellctl-runner-exit: %v\n", err) - os.Exit(1) - } + cmdutil.HandleError( + runnerexit.RecordRunnerExit(*stateDir, *jobID, *exitCode, *endedAt, *busyTimeoutMs), + 1, "record runner exit", + ) } diff --git a/dify-agent-runtime/cmd/runner/main.go b/dify-agent-runtime/cmd/runner/main.go new file mode 100644 index 00000000000..910eff8a9d2 --- /dev/null +++ b/dify-agent-runtime/cmd/runner/main.go @@ -0,0 +1,289 @@ +// shellctl-runner is the Go replacement for the previously generated bash+python +// runner script. It is invoked by tmux as: +// +// shellctl-runner +// +// The binary operates in two modes: +// +// 1. Parent mode (default): waits for start-gate, loads env, forks child, +// waits for exit, writes exit artifacts. +// 2. Child mode (--exec flag): applies Landlock restrictions, then exec's +// the user script. +// +// This eliminates the bash/python dependency chain and integrates Landlock +// directly. +package main + +import ( + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "os/signal" + "path/filepath" + "strings" + "syscall" + "time" + + "github.com/langgenius/dify/dify-agent-runtime/internal/cmdutil" + "github.com/langgenius/dify/dify-agent-runtime/internal/envvar" + "github.com/langgenius/dify/dify-agent-runtime/internal/landlock" +) + +func main() { + if len(os.Args) > 1 && os.Args[1] == "--exec" { + childMode() + return + } + parentMode() +} + +// parentMode is the entry point when called by tmux. +// Args: shellctl-runner +func parentMode() { + if len(os.Args) < 4 { + cmdutil.HandleError(fmt.Errorf("bad args"), 125, "usage: shellctl-runner ") + } + + jobDir := os.Args[1] + // jobID := os.Args[2] // unused in parent but passed for compat + cwd := os.Args[3] + + scriptPath := filepath.Join(jobDir, "script") + envPath := filepath.Join(jobDir, ".job-env.json") + startGate := filepath.Join(jobDir, "start-gate") + exitCodePath := filepath.Join(jobDir, "runner-exit-code") + endedAtPath := filepath.Join(jobDir, "runner-ended-at") + + // Wait for start-gate. + for { + if _, err := os.Stat(startGate); err == nil { + break + } + time.Sleep(50 * time.Millisecond) + } + + // Load environment overlay from JSON. + env := os.Environ() + // Remove internal shellctl vars from inherited env. + env = filterEnv(env, []string{ + "TMUX", + "SHELLCTL_STATE_DIR", + "SHELLCTL_RUNTIME_DIR", + "SHELLCTL_TMUX_SOCKET", + "SHELLCTL_RUNNER", + "SHELLCTL_AUTH_TOKEN", + }) + + envOverlay := loadEnvJSON(envPath) + env = mergeEnv(env, envOverlay) + + // Ensure HOME exists. + home := envGet(env, "HOME") + if home != "" { + cmdutil.HandleError(os.MkdirAll(home, 0755), 125, "mkdir HOME %s", home) + } + + // Create a per-workspace temp directory under cwd and inject TMPDIR. + // This avoids granting RW access to the shared /tmp. + agentTmp := filepath.Join(cwd, ".tmp") + cmdutil.HandleError(os.MkdirAll(agentTmp, 0755), 125, "mkdir TMPDIR %s", agentTmp) + env = setEnvIfEmpty(env, "TMPDIR", agentTmp) + env = setEnvIfEmpty(env, "TMP", agentTmp) + env = setEnvIfEmpty(env, "TEMP", agentTmp) + + // Determine if path isolation is enabled. + enableIsolation := envvar.PathIsolationEnabled() + + // Build child command: re-exec self in --exec mode. + self, _ := os.Executable() + childArgs := []string{"--exec", scriptPath, cwd} + if enableIsolation { + childArgs = append([]string{"--exec", "--landlock"}, childArgs[1:]...) + } + + cmd := exec.Command(self, childArgs...) + cmd.Env = env + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + cmd.Dir = cwd + + // Forward signals to child. + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP) + go func() { + for sig := range sigCh { + if cmd.Process != nil { + _ = cmd.Process.Signal(sig) + } + } + }() + + err := cmd.Run() + + exitCode := 0 + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + exitCode = exitErr.ExitCode() + } else { + exitCode = 125 + } + } + + endedAt := time.Now().UTC().Format("2006-01-02T15:04:05Z") + + writeAtomic(exitCodePath, fmt.Sprintf("%d", exitCode)) + writeAtomic(endedAtPath, endedAt) + + os.Exit(exitCode) +} + +// childMode applies Landlock (if --landlock flag) and exec's the user script. +// Args: shellctl-runner --exec [--landlock] +func childMode() { + args := os.Args[2:] // skip program name and "--exec" + applyLandlockFlag := false + + if len(args) > 0 && args[0] == "--landlock" { + applyLandlockFlag = true + args = args[1:] + } + + if len(args) < 2 { + cmdutil.HandleError(fmt.Errorf("bad args"), 125, "--exec: missing script_path and cwd") + } + + scriptPath := args[0] + cwd := args[1] + + // chdir + cmdutil.HandleError(os.Chdir(cwd), 111, "chdir %s", cwd) + + // Determine argv for exec. + argv := buildExecArgv(scriptPath) + + // Resolve binary before Landlock (PATH search). + binary, err := exec.LookPath(argv[0]) + cmdutil.HandleError(err, 127, "%s: not found", argv[0]) + + // Apply Landlock if requested. + if applyLandlockFlag { + home := os.Getenv("HOME") + jobDir := filepath.Dir(scriptPath) + cfg := landlock.ConfigFromEnv(home, cwd, jobDir) + if err := landlock.Restrict(cfg); err != nil { + if errors.Is(err, landlock.ErrNotSupported) { + fmt.Fprintf(os.Stderr, "shellctl-runner: WARNING: %v — running without filesystem isolation\n", err) + } else { + cmdutil.HandleError(err, 125, "landlock restrict") + } + } + } + + // exec replaces the current process. + cmdutil.HandleError(syscall.Exec(binary, argv, os.Environ()), 126, "exec %s", binary) +} + +// buildExecArgv determines whether the script has a shebang or needs sh. +func buildExecArgv(scriptPath string) []string { + f, err := os.Open(scriptPath) + if err != nil { + return []string{"sh", scriptPath} + } + defer func() { _ = f.Close() }() + + buf := make([]byte, 2) + n, _ := f.Read(buf) + if n == 2 && string(buf) == "#!" { + // Has shebang — make executable and run directly. + _ = os.Chmod(scriptPath, 0755) + return []string{scriptPath} + } + return []string{"sh", scriptPath} +} + +// loadEnvJSON reads a JSON object from path and returns key=value pairs. +func loadEnvJSON(path string) map[string]string { + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return nil + } + cmdutil.HandleError(err, 125, "read env json %s", path) + } + var m map[string]string + if err := json.Unmarshal(data, &m); err != nil { + cmdutil.HandleError(err, 125, "parse env json %s", path) + } + return m +} + +// filterEnv removes entries with the given key prefixes from env. +func filterEnv(env []string, remove []string) []string { + result := env[:0] + for _, e := range env { + skip := false + for _, r := range remove { + if strings.HasPrefix(e, r+"=") { + skip = true + break + } + } + if !skip { + result = append(result, e) + } + } + return result +} + +// mergeEnv applies overlay key=value pairs onto the env slice. +func mergeEnv(env []string, overlay map[string]string) []string { + for k, v := range overlay { + found := false + prefix := k + "=" + for i, e := range env { + if strings.HasPrefix(e, prefix) { + env[i] = prefix + v + found = true + break + } + } + if !found { + env = append(env, prefix+v) + } + } + return env +} + +// envGet retrieves a value from a []string env slice. +func envGet(env []string, key string) string { + prefix := key + "=" + for _, e := range env { + if strings.HasPrefix(e, prefix) { + return e[len(prefix):] + } + } + return "" +} + +// setEnvIfEmpty sets key=value in the env slice only if the key is not already present. +func setEnvIfEmpty(env []string, key, value string) []string { + if envGet(env, key) != "" { + return env + } + return append(env, key+"="+value) +} + +// writeAtomic writes value to dest via a temp file + rename. +func writeAtomic(dest, value string) { + tmp := fmt.Sprintf("%s.tmp.%d", dest, os.Getpid()) + if err := os.WriteFile(tmp, []byte(value+"\n"), 0644); err != nil { + fmt.Fprintf(os.Stderr, "shellctl-runner: write %s: %v\n", tmp, err) + return + } + if err := os.Rename(tmp, dest); err != nil { + fmt.Fprintf(os.Stderr, "shellctl-runner: rename %s -> %s: %v\n", tmp, dest, err) + } +} diff --git a/dify-agent-runtime/cmd/sanitize-pty/main.go b/dify-agent-runtime/cmd/sanitize-pty/main.go index 298071f2f19..2a653060de8 100644 --- a/dify-agent-runtime/cmd/sanitize-pty/main.go +++ b/dify-agent-runtime/cmd/sanitize-pty/main.go @@ -5,9 +5,9 @@ package main import ( "flag" - "fmt" "os" + "github.com/langgenius/dify/dify-agent-runtime/internal/cmdutil" "github.com/langgenius/dify/dify-agent-runtime/internal/sanitize" ) @@ -15,8 +15,5 @@ func main() { readyFile := flag.String("ready-file", "", "path to touch before reading stdin") flag.Parse() - if err := sanitize.Run(*readyFile, os.Stdin, os.Stdout); err != nil { - fmt.Fprintf(os.Stderr, "shellctl-sanitize-pty: %v\n", err) - os.Exit(1) - } + cmdutil.HandleError(sanitize.Run(*readyFile, os.Stdin, os.Stdout), 1, "sanitize-pty") } diff --git a/dify-agent-runtime/docker/Dockerfile b/dify-agent-runtime/docker/Dockerfile index 97f8ee81c99..798e36a8d41 100644 --- a/dify-agent-runtime/docker/Dockerfile +++ b/dify-agent-runtime/docker/Dockerfile @@ -14,6 +14,7 @@ COPY . . RUN CGO_ENABLED=0 go build -o /bin/shellctl ./cmd/shellctl && \ CGO_ENABLED=0 go build -o /bin/shellctl-sanitize-pty ./cmd/sanitize-pty && \ CGO_ENABLED=0 go build -o /bin/shellctl-runner-exit ./cmd/runner-exit && \ + CGO_ENABLED=0 go build -o /bin/shellctl-runner ./cmd/runner && \ CGO_ENABLED=0 go build -o /bin/dify-agent ./cmd/dify-agent-cli # ── Runtime stage ──────────────────────────────────────────────────────────── @@ -65,6 +66,7 @@ RUN python -m pip install --no-cache-dir "uv==${UV_VERSION}" COPY --from=go-builder /bin/shellctl /usr/local/bin/shellctl COPY --from=go-builder /bin/shellctl-sanitize-pty /usr/local/bin/shellctl-sanitize-pty COPY --from=go-builder /bin/shellctl-runner-exit /usr/local/bin/shellctl-runner-exit +COPY --from=go-builder /bin/shellctl-runner /usr/local/bin/shellctl-runner COPY --from=go-builder /bin/dify-agent /usr/local/bin/dify-agent RUN useradd --create-home --shell /bin/sh dify \ diff --git a/dify-agent-runtime/go.mod b/dify-agent-runtime/go.mod index 80f71bfb98b..9cf551446c4 100644 --- a/dify-agent-runtime/go.mod +++ b/dify-agent-runtime/go.mod @@ -3,6 +3,7 @@ module github.com/langgenius/dify/dify-agent-runtime go 1.26 require ( + github.com/landlock-lsm/go-landlock v0.9.0 github.com/spf13/cobra v1.10.2 google.golang.org/grpc v1.82.0 google.golang.org/protobuf v1.36.11 @@ -22,6 +23,7 @@ require ( golang.org/x/sys v0.43.0 // indirect golang.org/x/text v0.36.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect + kernel.org/pub/linux/libs/security/libcap/psx v1.2.77 // indirect modernc.org/libc v1.65.7 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.11.0 // indirect diff --git a/dify-agent-runtime/go.sum b/dify-agent-runtime/go.sum index 91aa0612cef..15705eabb41 100644 --- a/dify-agent-runtime/go.sum +++ b/dify-agent-runtime/go.sum @@ -17,6 +17,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/landlock-lsm/go-landlock v0.9.0 h1:2q8G8yx9Hsd5bV+R6PJfgQl0zszNxC8KO+SIqGwfxlw= +github.com/landlock-lsm/go-landlock v0.9.0/go.mod h1:mn5GSi81Jf7yMs5WSi+SUi4sUeNLUGVdbT4Id6wXNQw= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4= @@ -65,6 +67,8 @@ google.golang.org/grpc v1.82.0/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3 google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +kernel.org/pub/linux/libs/security/libcap/psx v1.2.77 h1:Z06sMOzc0GNCwp6efaVrIrz4ywGJ1v+DP0pjVkOfDuA= +kernel.org/pub/linux/libs/security/libcap/psx v1.2.77/go.mod h1:+l6Ee2F59XiJ2I6WR5ObpC1utCQJZ/VLsEbQCD8RG24= modernc.org/cc/v4 v4.26.1 h1:+X5NtzVBn0KgsBCBe+xkDC7twLb/jNVj9FPgiwSQO3s= modernc.org/cc/v4 v4.26.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0= modernc.org/ccgo/v4 v4.28.0 h1:rjznn6WWehKq7dG4JtLRKxb52Ecv8OUGah8+Z/SfpNU= diff --git a/dify-agent-runtime/internal/agentcli/env.go b/dify-agent-runtime/internal/agentcli/env.go index 8e7962f90f8..1db58b057d2 100644 --- a/dify-agent-runtime/internal/agentcli/env.go +++ b/dify-agent-runtime/internal/agentcli/env.go @@ -9,14 +9,16 @@ import ( "net/url" "os" "strings" + + "github.com/langgenius/dify/dify-agent-runtime/internal/envvar" ) const ( - EnvAPIBaseURL = "DIFY_AGENT_STUB_API_BASE_URL" - EnvAuthJWE = "DIFY_AGENT_STUB_AUTH_JWE" - EnvDriveBase = "DIFY_AGENT_STUB_DRIVE_BASE" + EnvAPIBaseURL = envvar.EnvAgentStubAPIBaseURL + EnvAuthJWE = envvar.EnvAgentStubAuthJWE + EnvDriveBase = envvar.EnvAgentStubDriveBase - DefaultDriveBase = "/mnt/drive" + DefaultDriveBase = envvar.DefaultDriveBase ) // Environment holds validated Agent Stub connection parameters. diff --git a/dify-agent-runtime/internal/cmdutil/cmdutil.go b/dify-agent-runtime/internal/cmdutil/cmdutil.go new file mode 100644 index 00000000000..826cf45bc90 --- /dev/null +++ b/dify-agent-runtime/internal/cmdutil/cmdutil.go @@ -0,0 +1,19 @@ +package cmdutil + +import ( + "fmt" + "os" +) + +// HandleError checks err; if non-nil it prints a formatted message to stderr +// and exits with the given code. Callers can use it as a one-liner: +// +// cmdutil.HandleError(os.MkdirAll(dir, 0755), 125, "mkdir %s", dir) +func HandleError(err error, code int, format string, args ...any) { + if err == nil { + return + } + msg := fmt.Sprintf(format, args...) + fmt.Fprintf(os.Stderr, "shellctl: %s: %v\n", msg, err) + os.Exit(code) +} diff --git a/dify-agent-runtime/internal/envvar/envvar.go b/dify-agent-runtime/internal/envvar/envvar.go new file mode 100644 index 00000000000..7ecbd96b503 --- /dev/null +++ b/dify-agent-runtime/internal/envvar/envvar.go @@ -0,0 +1,56 @@ +// Package envvar centralizes all environment variable constants used by the +// dify-agent-runtime binaries (shellctl, shellctl-runner, dify-agent CLI). +package envvar + +import "os" + +// --- Path Isolation (Landlock) --- + +const ( + // EnvEnablePathIsolation controls whether Landlock is applied at all. + EnvEnablePathIsolation = "ENABLE_PATH_ISOLATION" + + // EnvRWPaths overrides the default RW directories (comma-separated). + EnvRWPaths = "LANDLOCK_RW_PATHS" + + // EnvROPaths overrides the default RO+exec directories (comma-separated). + EnvROPaths = "LANDLOCK_RO_PATHS" + + // EnvRWDevPaths overrides the default device files (comma-separated). + EnvRWDevPaths = "LANDLOCK_RW_DEV_PATHS" +) + +// --- Agent Stub --- + +const ( + // EnvAgentStubAPIBaseURL is the Agent Stub HTTP/gRPC endpoint. + EnvAgentStubAPIBaseURL = "DIFY_AGENT_STUB_API_BASE_URL" + + // EnvAgentStubAuthJWE is the per-request JWE token for Agent Stub auth. + EnvAgentStubAuthJWE = "DIFY_AGENT_STUB_AUTH_JWE" + + // EnvAgentStubDriveBase is the sandbox-local drive directory for the agent. + EnvAgentStubDriveBase = "DIFY_AGENT_STUB_DRIVE_BASE" + + // DefaultDriveBase is the default Agent Stub drive mount point. + // currently unused. + DefaultDriveBase = "/mnt/drive" +) + +// PathIsolationEnabled returns whether Landlock filesystem isolation is active. +func PathIsolationEnabled() bool { + v, ok := os.LookupEnv(EnvEnablePathIsolation) + if !ok { + return true + } + return v == "true" +} + +// SetPathIsolation explicitly enables or disables path isolation via env. +func SetPathIsolation(enabled bool) { + if enabled { + _ = os.Setenv(EnvEnablePathIsolation, "true") + } else { + _ = os.Setenv(EnvEnablePathIsolation, "false") + } +} diff --git a/dify-agent-runtime/internal/landlock/config.go b/dify-agent-runtime/internal/landlock/config.go new file mode 100644 index 00000000000..73cd791caa5 --- /dev/null +++ b/dify-agent-runtime/internal/landlock/config.go @@ -0,0 +1,95 @@ +package landlock + +import ( + "os" + "strings" + + "github.com/langgenius/dify/dify-agent-runtime/internal/envvar" +) + +// Config holds the path lists for Landlock filesystem restrictions. +type Config struct { + Home string // Agent's home directory (always RW). + Cwd string // Working directory (always RW, may differ from Home). + JobDir string // Job script directory (RO, auto-set by runner). + RWPaths []string // Additional read-write directories. + ROPaths []string // Read-only + execute directories. + RWDevPaths []string // Device files with read-write access. +} + +var ( + // DefaultRWPaths are directories granted read-write access besides HOME. + // Agent-specific paths (e.g. drive base) are added dynamically by the runner. + DefaultRWPaths = []string{} + + // DefaultROPaths are directories granted read-only + execute access. + DefaultROPaths = []string{ + "/usr", + "/bin", + "/sbin", + "/lib", + "/lib64", + "/etc", + "/proc", + "/opt/dify-agent-tools", + "/opt/homebrew", + "/snap", + } + + // DefaultRWDevPaths are device files granted read-write access. + DefaultRWDevPaths = []string{ + "/dev/null", + "/dev/zero", + "/dev/urandom", + "/dev/random", + "/dev/tty", + } +) + +// DefaultConfig returns a Config with sensible defaults. +func DefaultConfig(home, cwd, jobDir string) *Config { + return &Config{ + Home: home, + Cwd: cwd, + JobDir: jobDir, + RWPaths: DefaultRWPaths, + ROPaths: DefaultROPaths, + RWDevPaths: DefaultRWDevPaths, + } +} + +// ConfigFromEnv returns a Config populated from environment variables, +// falling back to defaults for any unset variable. +// +// If a variable is set (even to empty string), its value replaces the default. +// Set to empty to grant no additional paths beyond $HOME. +// +// LANDLOCK_RW_PATHS — comma-separated RW dirs (default: empty) +// LANDLOCK_RO_PATHS — comma-separated RO dirs (default: system paths) +// LANDLOCK_RW_DEV_PATHS — comma-separated dev files (default: /dev/null,...) +func ConfigFromEnv(home, cwd, jobDir string) *Config { + cfg := DefaultConfig(home, cwd, jobDir) + + if v, ok := os.LookupEnv(envvar.EnvRWPaths); ok { + cfg.RWPaths = splitPaths(v) + } + if v, ok := os.LookupEnv(envvar.EnvROPaths); ok { + cfg.ROPaths = splitPaths(v) + } + if v, ok := os.LookupEnv(envvar.EnvRWDevPaths); ok { + cfg.RWDevPaths = splitPaths(v) + } + + return cfg +} + +func splitPaths(s string) []string { + var paths []string + for _, p := range strings.Split(s, ",") { + p = strings.TrimSpace(p) + if p != "" { + paths = append(paths, p) + } + } + return paths +} diff --git a/dify-agent-runtime/internal/landlock/landlock_linux.go b/dify-agent-runtime/internal/landlock/landlock_linux.go new file mode 100644 index 00000000000..8249df278dd --- /dev/null +++ b/dify-agent-runtime/internal/landlock/landlock_linux.go @@ -0,0 +1,80 @@ +//go:build linux + +package landlock + +import ( + "errors" + "fmt" + "os" + "syscall" + + golandlock "github.com/landlock-lsm/go-landlock/landlock" +) + +// ErrNotSupported is returned when the kernel does not support Landlock. +var ErrNotSupported = errors.New("landlock: not supported by kernel") + +// Restrict applies Landlock filesystem restrictions for the current process. +// +// It first tries strict enforcement. If the kernel lacks Landlock support, +// it falls back to BestEffort (which may silently degrade) and returns +// ErrNotSupported so callers can log a warning. +func Restrict(cfg *Config) error { + rules := buildRules(cfg) + + if err := golandlock.V5.RestrictPaths(rules...); err != nil { + if !isNotSupported(err) { + return fmt.Errorf("landlock: restrict paths: %w", err) + } + if err2 := golandlock.V5.BestEffort().RestrictPaths(rules...); err2 != nil { + return fmt.Errorf("landlock: best-effort restrict paths: %w", err2) + } + return ErrNotSupported + } + return nil +} + +func buildRules(cfg *Config) []golandlock.Rule { + // Collect RO paths that exist on this system. + roPaths := filterExisting(cfg.ROPaths) + if cfg.JobDir != "" { + roPaths = append(roPaths, cfg.JobDir) + } + + // RW dirs: always include HOME and Cwd, plus configured extra paths. + rwDirs := []string{cfg.Home} + if cfg.Cwd != "" && cfg.Cwd != cfg.Home { + rwDirs = append(rwDirs, cfg.Cwd) + } + rwDirs = append(rwDirs, filterExisting(cfg.RWPaths)...) + + // Device files with RW access. + rwDevFiles := filterExisting(cfg.RWDevPaths) + + rules := []golandlock.Rule{ + golandlock.RWDirs(rwDirs...), + } + if len(roPaths) > 0 { + rules = append(rules, golandlock.RODirs(roPaths...)) + } + if len(rwDevFiles) > 0 { + rules = append(rules, golandlock.RWFiles(rwDevFiles...)) + } + return rules +} + +// isNotSupported reports whether err indicates the kernel does not support +// Landlock (ENOSYS = syscall absent, EOPNOTSUPP = feature disabled). +func isNotSupported(err error) bool { + return errors.Is(err, syscall.ENOSYS) || errors.Is(err, syscall.EOPNOTSUPP) +} + +func filterExisting(paths []string) []string { + var out []string + for _, p := range paths { + if _, err := os.Stat(p); err == nil { + out = append(out, p) + } + } + return out +} diff --git a/dify-agent-runtime/internal/landlock/landlock_stub.go b/dify-agent-runtime/internal/landlock/landlock_stub.go new file mode 100644 index 00000000000..d8bb3419a05 --- /dev/null +++ b/dify-agent-runtime/internal/landlock/landlock_stub.go @@ -0,0 +1,13 @@ +//go:build !linux + +package landlock + +import "errors" + +// ErrNotSupported is returned when the kernel does not support Landlock. +var ErrNotSupported = errors.New("landlock: not supported by kernel") + +// Restrict is a no-op on non-Linux platforms; always returns ErrNotSupported. +func Restrict(cfg *Config) error { + return ErrNotSupported +} diff --git a/dify-agent-runtime/internal/server/config.go b/dify-agent-runtime/internal/server/config.go index 6b6258a8f54..226a44b8c43 100644 --- a/dify-agent-runtime/internal/server/config.go +++ b/dify-agent-runtime/internal/server/config.go @@ -8,25 +8,25 @@ import ( ) const ( - DefaultListen = "127.0.0.1:8765" - DefaultTimeoutSeconds = 30.0 - DefaultMaxWaitTimeoutSeconds = 600.0 - DefaultIdleFlushSeconds = 0.5 - DefaultTerminalCols = 200 - DefaultTerminalRows = 50 - DefaultOutputLimitBytes = 16 * 1024 - MaxOutputLimitBytes = 512 * 1024 - DefaultListLimit = 50 - MaxListLimit = 200 - DefaultTerminateGraceSeconds = 10.0 - DefaultGCIntervalSeconds = 60.0 + DefaultListen = "127.0.0.1:8765" + DefaultTimeoutSeconds = 30.0 + DefaultMaxWaitTimeoutSeconds = 600.0 + DefaultIdleFlushSeconds = 0.5 + DefaultTerminalCols = 200 + DefaultTerminalRows = 50 + DefaultOutputLimitBytes = 16 * 1024 + MaxOutputLimitBytes = 512 * 1024 + DefaultListLimit = 50 + MaxListLimit = 200 + DefaultTerminateGraceSeconds = 10.0 + DefaultGCIntervalSeconds = 60.0 DefaultGCFinishedJobRetentionSeconds = 300.0 - DefaultPollInterval = 50 * time.Millisecond - DefaultPipeMonitorInterval = 1 * time.Second - DefaultPipeReadyTimeout = 10 * time.Second - DefaultSQLiteBusyTimeoutMs = 5000 - DefaultAuthTokenEnv = "SHELLCTL_AUTH_TOKEN" - HealthStatus = "ok" + DefaultPollInterval = 50 * time.Millisecond + DefaultPipeMonitorInterval = 1 * time.Second + DefaultPipeReadyTimeout = 10 * time.Second + DefaultSQLiteBusyTimeoutMs = 5000 + DefaultAuthTokenEnv = "SHELLCTL_AUTH_TOKEN" + HealthStatus = "ok" ) // Config holds runtime configuration for the shellctl server. diff --git a/dify-agent-runtime/internal/server/service.go b/dify-agent-runtime/internal/server/service.go index d02e812a3b9..a859aa3474d 100644 --- a/dify-agent-runtime/internal/server/service.go +++ b/dify-agent-runtime/internal/server/service.go @@ -5,6 +5,7 @@ import ( "fmt" "log" "os" + "os/exec" "path/filepath" "strconv" "strings" @@ -826,98 +827,35 @@ func (s *Service) cleanupStarting(jobID, jobDir string) { } func (s *Service) installRunner() { - script := s.runnerScriptSource() - _ = os.WriteFile(s.config.RunnerPath(), []byte(script), 0755) -} + // The runner is now a pre-built Go binary (shellctl-runner). + // Create a symlink at the expected RunnerPath so tmux can invoke it. + dst := s.config.RunnerPath() + _ = os.Remove(dst) // remove stale script/symlink -func (s *Service) runnerScriptSource() string { - // The runner script is a bash wrapper that waits for the start-gate, - // then delegates env loading + exec to a Python bootstrap so that - // JSON env values (including multi-line strings) are applied verbatim. - // The Python helper exec's the target process so SIGINT semantics - // match running the script directly in the tmux pane. - return `#!/usr/bin/env bash -set -uo pipefail + // Look for the binary next to the main shellctl binary first, + // then fall back to PATH lookup. + src := "" + selfExe, _ := os.Executable() + if selfExe != "" { + candidate := filepath.Join(filepath.Dir(selfExe), "shellctl-runner") + if _, err := os.Stat(candidate); err == nil { + src = candidate + } + } + if src == "" { + if p, err := exec.LookPath("shellctl-runner"); err == nil { + src = p + } + } -JOB_DIR="$1" -JOB_ID="$2" -CWD="$3" -SCRIPT_PATH="$JOB_DIR/script" -ENV_PATH="$JOB_DIR/.job-env.json" -START_GATE="$JOB_DIR/start-gate" -RUNNER_EXIT_CODE_PATH="$JOB_DIR/runner-exit-code" -RUNNER_ENDED_AT_PATH="$JOB_DIR/runner-ended-at" - -write_atomic() { - local dest="$1" - local value="$2" - local tmp="${dest}.tmp.$$" - printf '%s\n' "$value" > "$tmp" - mv "$tmp" "$dest" -} - -while [ ! -e "$START_GATE" ]; do - sleep 0.05 -done - -unset TMUX -unset SHELLCTL_STATE_DIR -unset SHELLCTL_RUNTIME_DIR -unset SHELLCTL_TMUX_SOCKET -unset SHELLCTL_RUNNER -unset SHELLCTL_AUTH_TOKEN - -# Use Python bootstrap to load JSON env verbatim and exec the target. -# This avoids depending on jq and correctly handles multi-line values. -python3 -c ' -import json, os, stat, sys -from pathlib import Path - -script_path = Path(sys.argv[1]) -cwd = sys.argv[2] -env_path = Path(sys.argv[3]) - -env = os.environ.copy() -if env_path.exists(): - env.update(json.loads(env_path.read_text(encoding="utf-8"))) - -try: - os.chdir(cwd) -except OSError: - raise SystemExit(111) - -# Ensure HOME directory exists (agent backend may set a per-agent HOME). -home = env.get("HOME", "") -if home and not os.path.isdir(home): - os.makedirs(home, exist_ok=True) - -with script_path.open("r", encoding="utf-8") as handle: - first_line = handle.readline() - -if first_line.startswith("#!"): - script_path.chmod(script_path.stat().st_mode | stat.S_IXUSR) - argv = [str(script_path)] -else: - argv = ["sh", str(script_path)] - -try: - os.execvpe(argv[0], argv, env) -except FileNotFoundError as exc: - print(f"{argv[0]}: {exc.strerror}", file=sys.stderr) - raise SystemExit(127) from exc -except OSError as exc: - print(f"{argv[0]}: {exc.strerror}", file=sys.stderr) - raise SystemExit(126) from exc -' "$SCRIPT_PATH" "$CWD" "$ENV_PATH" -EXIT_CODE=$? - -ENDED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" - -write_atomic "$RUNNER_EXIT_CODE_PATH" "$EXIT_CODE" -write_atomic "$RUNNER_ENDED_AT_PATH" "$ENDED_AT" - -exit "$EXIT_CODE" -` + if src != "" { + // Prefer symlink; fall back to copy path reference. + if err := os.Symlink(src, dst); err != nil { + // If symlink fails (e.g., cross-device), write a tiny shell wrapper. + wrapper := fmt.Sprintf("#!/bin/sh\nexec %q \"$@\"\n", src) + _ = os.WriteFile(dst, []byte(wrapper), 0755) + } + } } // Helpers diff --git a/dify-agent-runtime/tests/acceptance_test.go b/dify-agent-runtime/tests/acceptance_test.go index cfd5bd1f4be..caf3b773791 100644 --- a/dify-agent-runtime/tests/acceptance_test.go +++ b/dify-agent-runtime/tests/acceptance_test.go @@ -26,6 +26,9 @@ var ( goURL = envOrDefault("SHELLCTL_GO_URL", "http://localhost:15005") authToken = envOrDefault("SHELLCTL_TEST_TOKEN", "test-token-123") + goURLNoIsolation = os.Getenv("SHELLCTL_GO_URL_NO_ISOLATION") + authTokenNoIsolation = os.Getenv("SHELLCTL_TEST_TOKEN_NO_ISOLATION") + httpClient = &http.Client{Timeout: 120 * time.Second} ) @@ -41,6 +44,13 @@ func targets() []target { } } +func noIsolationTarget() (target, bool) { + if goURLNoIsolation == "" { + return target{}, false + } + return target{name: "go-no-isolation", baseURL: goURLNoIsolation}, true +} + func TestMain(m *testing.M) { // Warmup: wait for both servers to be ready before running tests for _, tgt := range targets() { @@ -49,10 +59,19 @@ func TestMain(m *testing.M) { os.Exit(1) } } + if tgt, ok := noIsolationTarget(); ok { + if !waitForServer(tgt) { + fmt.Fprintf(os.Stderr, "ERROR: %s server not ready at %s\n", tgt.name, tgt.baseURL) + os.Exit(1) + } + } for _, tgt := range targets() { warmupJob(tgt) } + if tgt, ok := noIsolationTarget(); ok { + warmupJobWithToken(tgt, authTokenNoIsolation) + } os.Exit(m.Run()) } @@ -489,6 +508,166 @@ func TestJobNotFound(t *testing.T) { } } +// --- Landlock Tests --- +// These tests verify that shellctl-run restricts filesystem access +// so each agent job can only write within its own HOME directory while still +// being able to execute system commands. + +func TestLandlockCanWriteHome(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // The job runs with HOME set; writing a file inside HOME should succeed. + result := runJob(t, tgt, map[string]any{ + "script": "touch \"$HOME/landlock-test-file\" && echo ok", + "env": map[string]string{"HOME": "/home/dify"}, + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "ok") { + t.Errorf("expected write to HOME to succeed, got %q", output) + } + }) + } +} + +func TestLandlockCanReadSystemBinaries(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // System binaries should remain executable (RO access to /usr, /bin). + result := runJob(t, tgt, map[string]any{ + "script": "which ls && ls /usr/bin/env && echo ok", + "env": map[string]string{"HOME": "/home/dify"}, + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "ok") { + t.Errorf("expected system binary access to succeed, got %q", output) + } + }) + } +} + +func TestLandlockCanWriteTmpdir(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // TMPDIR ($CWD/.tmp) should be writable; /tmp should be denied. + result := runJob(t, tgt, map[string]any{ + "script": "echo TMPDIR=$TMPDIR && touch $TMPDIR/landlock-tmp-test && echo tmpdir_ok && touch /tmp/landlock-denied 2>&1; echo tmp_exit=$?", + "env": map[string]string{"HOME": "/home/dify"}, + "timeout": 10, + }) + assertJobDone(t, result) + output := result["output"].(string) + if !strings.Contains(output, "tmpdir_ok") { + t.Errorf("expected write to $TMPDIR to succeed, got %q", output) + } + if !strings.Contains(output, "tmp_exit=1") && !strings.Contains(output, "Permission denied") { + t.Errorf("expected write to /tmp to be denied, got %q", output) + } + }) + } +} + +func TestLandlockCannotWriteOutsideHome(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // Writing outside HOME (e.g., /opt) should be denied by Landlock. + result := runJob(t, tgt, map[string]any{ + "script": "touch /opt/landlock-denied 2>&1; echo exit=$?", + "env": map[string]string{"HOME": "/home/dify"}, + "timeout": 10, + }) + assertJobDone(t, result) + output := result["output"].(string) + // The touch should fail with "Permission denied" or similar, + // and exit code should be non-zero. + if strings.Contains(output, "exit=0") { + t.Errorf("expected write to /opt to be denied, but it succeeded: %q", output) + } + }) + } +} + +func TestLandlockCannotReadOtherAgentHome(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // First, create a file in one agent's home. + setup := runJob(t, tgt, map[string]any{ + "script": "mkdir -p /home/agent-a && touch /home/agent-a/secret", + "env": map[string]string{"HOME": "/home/agent-a"}, + "timeout": 10, + }) + assertJobDone(t, setup) + assertExitCode(t, setup, 0) + + // Now run as a different agent and try to read the other's file. + result := runJob(t, tgt, map[string]any{ + "script": "cat /home/agent-a/secret 2>&1; echo exit=$?", + "env": map[string]string{"HOME": "/home/agent-b"}, + "timeout": 10, + }) + assertJobDone(t, result) + output := result["output"].(string) + if strings.Contains(output, "exit=0") { + t.Errorf("expected read of other agent's file to be denied, but it succeeded: %q", output) + } + }) + } +} + +// --- Landlock Disable / Bypass Tests --- + +// TestLandlockDisabledAllowsWriteOutsideHome uses the pre-started no-isolation +// container (ENABLE_PATH_ISOLATION=false) and verifies that isolation is off. +func TestLandlockDisabledAllowsWriteOutsideHome(t *testing.T) { + tgt, ok := noIsolationTarget() + if !ok { + t.Skip("SHELLCTL_GO_URL_NO_ISOLATION not set; no-isolation container not available") + } + + // With isolation disabled, writes to /tmp should succeed. + // /tmp is world-writable (Unix perms) but blocked by Landlock when enabled. + result := runJobWithToken(t, tgt, authTokenNoIsolation, map[string]any{ + "script": "touch /tmp/landlock-disabled-test && echo write_ok", + "env": map[string]string{"HOME": "/home/dify"}, + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "write_ok") { + t.Errorf("expected write to /tmp to succeed with isolation disabled, got %q", output) + } +} + +// TestLandlockEnvBypassBlocked verifies that a caller cannot set +// ENABLE_PATH_ISOLATION=false in job env to escape the sandbox. +func TestLandlockEnvBypassBlocked(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // Attempt to bypass Landlock by passing the disable flag in job env. + result := runJob(t, tgt, map[string]any{ + "script": "touch /opt/landlock-bypass-test 2>&1; echo exit=$?", + "env": map[string]string{ + "HOME": "/home/dify", + "ENABLE_PATH_ISOLATION": "false", + }, + "timeout": 10, + }) + assertJobDone(t, result) + output := result["output"].(string) + // The write should still be denied despite the env override attempt. + if strings.Contains(output, "exit=0") { + t.Errorf("expected write to /opt to be DENIED even with ENABLE_PATH_ISOLATION=false in job env, but it succeeded: %q", output) + } + }) + } +} + // --- Helpers --- func runJob(t *testing.T, tgt target, payload map[string]any) map[string]any { @@ -585,3 +764,46 @@ func envOrDefault(key, defaultVal string) string { } return defaultVal } + +func warmupJobWithToken(tgt target, token string) { + warmupClient := &http.Client{Timeout: 180 * time.Second} + payload, _ := json.Marshal(map[string]any{ + "script": "echo warmup", + "timeout": 10, + }) + for attempt := 0; attempt < 3; attempt++ { + req, _ := http.NewRequest("POST", tgt.baseURL+"/v1/jobs/run", bytes.NewReader(payload)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+token) + resp, err := warmupClient.Do(req) + if err != nil { + time.Sleep(2 * time.Second) + continue + } + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + if resp.StatusCode == 200 { + return + } + time.Sleep(2 * time.Second) + } +} + +func runJobWithToken(t *testing.T, tgt target, token string, payload map[string]any) map[string]any { + t.Helper() + body, _ := json.Marshal(payload) + req, _ := http.NewRequest("POST", tgt.baseURL+"/v1/jobs/run", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+token) + resp, err := httpClient.Do(req) + if err != nil { + t.Fatalf("[%s] POST /v1/jobs/run failed: %v", tgt.name, err) + } + assertStatus(t, resp, 200) + respBody := readBody(t, resp) + var result map[string]any + if err := json.Unmarshal(respBody, &result); err != nil { + t.Fatalf("[%s] failed to parse run response: %v\nbody: %s", tgt.name, err, string(respBody)) + } + return result +}