@xyit/wellspring 0.0.0-stage → 1.1.0

package/src/sandbox/manager.ts

@@ -0,0 +1,406 @@
+/*
+ * SandboxManager：选择 backend、准备 workspace/scratch、执行 agent-code、收集 artifact。
+ */
+
+import { createHash, randomUUID } from "node:crypto"
+import { mkdir, readFile, readdir, rm, stat, writeFile } from "node:fs/promises"
+import { basename, dirname, join, relative, resolve, sep } from "node:path"
+
+import type { SandboxAccess } from "../config/types.ts"
+
+import { auditSandbox } from "./audit.ts"
+import { BwrapBackend } from "./backends/bwrap.ts"
+import { DirectBackend } from "./backends/direct.ts"
+import { createSandboxMemoryLimit, type MemoryLimitState } from "./memory-limit.ts"
+import { startSandboxProcess, type SandboxProcessHandle } from "./supervisor.ts"
+import {
+    SandboxExecutionError,
+    SandboxUnavailableError,
+    type AgentCodeRequest,
+    type AgentCodeResult,
+    type SandboxArtifact,
+    type SandboxBackend,
+    type SandboxBackendStatus,
+    type AgentCodeTaskHandle,
+    type SandboxLimits,
+    type SandboxManagerOptions,
+    type StopMode,
+} from "./types.ts"
+
+const MAX_ARTIFACT_FILES = 32
+const MAX_ARTIFACT_BYTES = 64 * 1024 * 1024
+
+interface ResolvedLimits extends SandboxLimits {
+    readonly timeoutMs: number
+    readonly memoryMB: number
+    readonly pids: number
+    readonly outputLimitBytes: number
+}
+
+function clampInteger(value: number, min: number, max: number): number {
+    return Math.min(max, Math.max(min, Math.trunc(value)))
+}
+
+function baseEnv(): Record<string, string> {
+    return {
+        PATH: "/usr/bin:/bin",
+        LANG: "C.UTF-8",
+        LC_ALL: "C.UTF-8",
+        TERM: "dumb",
+    }
+}
+
+async function ensureDirectory(path: string): Promise<void> {
+    await mkdir(path, { recursive: true })
+}
+
+function isInside(root: string, target: string): boolean {
+    return target === root || target.startsWith(`${root}${sep}`)
+}
+
+async function collectArtifacts(
+    scratchDir: string,
+    sandboxDir: string,
+): Promise<{ artifacts: SandboxArtifact[]; artifactDir: string | null }> {
+    const source = join(scratchDir, "artifacts")
+    const info = await stat(source).catch(() => null)
+    if (info === null || !info.isDirectory()) {
+        return { artifacts: [], artifactDir: null }
+    }
+
+    const destination = join(sandboxDir, "artifacts")
+    await ensureDirectory(destination)
+    const collected: SandboxArtifact[] = []
+    let total = 0
+
+    const entries = await readdir(source, { recursive: true, withFileTypes: true })
+    for (const entry of entries) {
+        if (!entry.isFile() || collected.length >= MAX_ARTIFACT_FILES) {
+            continue
+        }
+        const parent = entry.parentPath ?? source
+        const absolute = join(parent, entry.name)
+        const rel = relative(source, absolute)
+        if (rel.startsWith("..") || resolve(absolute) === resolve(source)) {
+            continue
+        }
+        const fileInfo = await stat(absolute)
+        if (fileInfo.size > MAX_ARTIFACT_BYTES || total + fileInfo.size > MAX_ARTIFACT_BYTES) {
+            continue
+        }
+        total += fileInfo.size
+        const data = await readFile(absolute)
+        const sha256 = createHash("sha256").update(data).digest("hex")
+        const target = join(destination, rel)
+        await ensureDirectory(dirname(target))
+        await writeFile(target, data)
+        collected.push({ path: rel.split(sep).join("/"), size: fileInfo.size, sha256 })
+    }
+
+    return { artifacts: collected, artifactDir: destination }
+}
+
+export class SandboxManager {
+    readonly #options: SandboxManagerOptions
+    #backend: SandboxBackend | null
+    #resolved: Promise<SandboxBackend> | null = null
+    #memoryLimitWarned = false
+    #active = 0
+    readonly #waiters: Array<() => void> = []
+
+    constructor(options: SandboxManagerOptions) {
+        this.#options = options
+        this.#backend = options.backend ?? null
+    }
+
+    async probe(): Promise<SandboxBackendStatus> {
+        if (this.#backend !== null) {
+            return await this.#backend.probe()
+        }
+        if (this.#options.settings.backend === "nsjail") {
+            return {
+                available: false,
+                backend: "nsjail",
+                reason: "nsjail 后端尚未实现",
+                version: null,
+            }
+        }
+        return await new BwrapBackend().probe()
+    }
+
+    async #resolveBackend(access: SandboxAccess): Promise<SandboxBackend> {
+        if (access === "danger-full-access") {
+            return new DirectBackend()
+        }
+        if (this.#backend !== null) {
+            return this.#backend
+        }
+        if (this.#resolved !== null) {
+            return await this.#resolved
+        }
+        this.#resolved = (async (): Promise<SandboxBackend> => {
+            const { settings, logger } = this.#options
+            if (settings.mode === "off") {
+                throw new SandboxUnavailableError("sandbox.mode = off，Agent 代码执行需要 sandbox")
+            }
+
+            const preferred = settings.backend
+            const candidates: SandboxBackend[] = []
+            if (preferred === "auto" || preferred === "bwrap") {
+                candidates.push(new BwrapBackend())
+            }
+            if (preferred === "nsjail") {
+                // nsjail 后端尚未实现，显式报告不可用
+                throw new SandboxUnavailableError("nsjail 后端尚未实现")
+            }
+
+            for (const candidate of candidates) {
+                const status = await candidate.probe()
+                if (status.available) {
+                    return candidate
+                }
+                logger?.warn(
+                    { backend: candidate.name, reason: status.reason },
+                    "sandbox backend 不可用",
+                )
+            }
+
+            if (settings.mode === "best-effort") {
+                logger?.warn(
+                    { backend: "direct" },
+                    "sandbox 降级为 direct（无隔离），仅限开发/受信环境",
+                )
+                return new DirectBackend()
+            }
+            throw new SandboxUnavailableError("没有可用的 OS 级 sandbox 后端")
+        })()
+        return await this.#resolved
+    }
+
+    async #acquireSlot(): Promise<void> {
+        const max = this.#options.settings.maxConcurrent
+        if (this.#active < max) {
+            this.#active += 1
+            return
+        }
+        await new Promise<void>((resolve) => {
+            this.#waiters.push(resolve)
+        })
+        this.#active += 1
+    }
+
+    #releaseSlot(): void {
+        this.#active -= 1
+        const next = this.#waiters.shift()
+        next?.()
+    }
+
+    #limits(request: AgentCodeRequest): ResolvedLimits {
+        const settings = this.#options.settings
+        return {
+            timeoutMs: clampInteger(request.timeoutMs ?? settings.defaultTimeoutMs, 100, 600_000),
+            memoryMB: clampInteger(request.memoryMB ?? settings.defaultMemoryMB, 32, 8_192),
+            pids: clampInteger(64, 1, 512),
+            outputLimitBytes: clampInteger(
+                request.outputLimitBytes ?? settings.outputLimitBytes,
+                1024,
+                64 * 1024 * 1024,
+            ),
+        }
+    }
+
+    async startAgentCode(request: AgentCodeRequest): Promise<AgentCodeTaskHandle> {
+        if (typeof request.command !== "string" || request.command.trim() === "") {
+            throw new SandboxExecutionError("命令不能为空")
+        }
+        if (request.argv.some((item) => typeof item !== "string")) {
+            throw new SandboxExecutionError("argv 必须是字符串数组")
+        }
+
+        const access = request.access ?? this.#options.settings.access
+        if (
+            access !== "read-only" &&
+            access !== "workspace-write" &&
+            access !== "danger-full-access"
+        ) {
+            throw new SandboxExecutionError(`未知 sandbox 权限：${String(access)}`)
+        }
+        const workspaceDir = resolve(request.workspaceDir ?? process.cwd())
+        const cwd = resolve(request.cwd ?? workspaceDir)
+        if (access !== "danger-full-access" && !isInside(workspaceDir, cwd)) {
+            throw new SandboxExecutionError(`cwd 不在 workspace 内：${cwd}`)
+        }
+        const workspaceInfo = await stat(workspaceDir).catch(() => null)
+        if (workspaceInfo === null || !workspaceInfo.isDirectory()) {
+            throw new SandboxExecutionError(`workspace 不存在：${workspaceDir}`)
+        }
+        const cwdInfo = await stat(cwd).catch(() => null)
+        if (cwdInfo === null || !cwdInfo.isDirectory()) {
+            throw new SandboxExecutionError(`cwd 不存在：${cwd}`)
+        }
+
+        const backend = await this.#resolveBackend(access)
+        await this.#acquireSlot()
+        try {
+            return await this.#startAgentCodeLocked(request, backend, workspaceDir, cwd, access)
+        } catch (error) {
+            this.#releaseSlot()
+            throw error
+        }
+    }
+
+    async runAgentCode(request: AgentCodeRequest): Promise<AgentCodeResult> {
+        const task = await this.startAgentCode(request)
+        return await task.wait()
+    }
+
+    async #startAgentCodeLocked(
+        request: AgentCodeRequest,
+        backend: SandboxBackend,
+        workspaceDir: string,
+        cwd: string,
+        access: SandboxAccess,
+    ): Promise<AgentCodeTaskHandle> {
+        const limits = this.#limits(request)
+        const sandboxId = randomUUID()
+        const sandboxDir = join(this.#options.paths.dataDir, "sandbox", sandboxId)
+        const scratchDir = join(sandboxDir, "scratch")
+        await ensureDirectory(join(scratchDir, "tmp"))
+        await ensureDirectory(join(scratchDir, "artifacts"))
+
+        const env = baseEnv()
+        if (backend.name === "bwrap") {
+            env["HOME"] = "/scratch"
+            env["TMPDIR"] = "/scratch/tmp"
+            env["SCRATCH"] = "/scratch"
+            env["WORKSPACE"] = "/workspace"
+        } else {
+            env["HOME"] = scratchDir
+            env["TMPDIR"] = join(scratchDir, "tmp")
+            env["SCRATCH"] = scratchDir
+            env["WORKSPACE"] = workspaceDir
+        }
+
+        const memoryLimit: MemoryLimitState | null =
+            backend.name === "bwrap"
+                ? await createSandboxMemoryLimit({
+                      bytes: limits.memoryMB * 1024 * 1024,
+                      timeoutMs: limits.timeoutMs,
+                  })
+                : null
+        if (memoryLimit !== null && !memoryLimit.active && !this.#memoryLimitWarned) {
+            this.#memoryLimitWarned = true
+            this.#options.logger?.warn(
+                { reason: memoryLimit.reason },
+                "沙箱物理内存限制不可用，仅保留 V8 堆上限与超时",
+            )
+        }
+
+        let processHandle: SandboxProcessHandle
+        try {
+            const baseSpec = backend.buildAgentLaunch({
+                workspaceDir,
+                scratchDir,
+                cwd,
+                access,
+                network: request.network,
+                command: request.command,
+                argv: request.argv,
+                env,
+                limits,
+            })
+            const spec =
+                memoryLimit !== null && memoryLimit.active ? memoryLimit.wrap(baseSpec) : baseSpec
+            processHandle = startSandboxProcess(spec, limits, {
+                stdin: request.stdin ?? "",
+            })
+        } catch (error) {
+            await memoryLimit?.finish()
+            this.#releaseSlot()
+            throw error
+        }
+
+        let finalize: Promise<AgentCodeResult> | null = null
+        const wait = (): Promise<AgentCodeResult> => {
+            finalize ??= (async (): Promise<AgentCodeResult> => {
+                try {
+                    const processResult = await processHandle.done
+                    const memory = await memoryLimit?.finish()
+                    const oomKilled = memory?.oomKilled === true
+                    const { artifacts, artifactDir } = await collectArtifacts(
+                        scratchDir,
+                        sandboxDir,
+                    )
+                    await rm(scratchDir, { recursive: true, force: true })
+                    auditSandbox(this.#options.paths, {
+                        type: "sandbox.agent-code",
+                        sandboxId,
+                        backend: backend.name,
+                        access,
+                        exitCode: processResult.exitCode,
+                        timedOut: processResult.timedOut,
+                        durationMs: processResult.durationMs,
+                        artifacts: artifacts.length,
+                    })
+                    this.#options.logger?.info(
+                        {
+                            sandboxId,
+                            backend: backend.name,
+                            command: basename(request.command),
+                            exitCode: processResult.exitCode,
+                            timedOut: processResult.timedOut,
+                            durationMs: processResult.durationMs,
+                        },
+                        "sandbox agent-code 执行完成",
+                    )
+                    return {
+                        sandboxId,
+                        backend: backend.name,
+                        access,
+                        exitCode: processResult.exitCode,
+                        signal: processResult.signal,
+                        stdout: processResult.stdout,
+                        stderr: processResult.stderr,
+                        truncated: processResult.truncated,
+                        timedOut: processResult.timedOut,
+                        memoryLimit: memoryLimit?.active === true ? "cgroup" : "none",
+                        memoryLimitReason:
+                            memoryLimit !== null && !memoryLimit.active ? memoryLimit.reason : null,
+                        oomKilled,
+                        durationMs: processResult.durationMs,
+                        artifacts,
+                        artifactDir:
+                            artifactDir === null
+                                ? null
+                                : backend.name === "bwrap"
+                                  ? "/scratch/artifacts"
+                                  : artifactDir,
+                        sandboxDir: backend.name === "bwrap" ? "/scratch" : sandboxDir,
+                        hostArtifactDir: artifactDir,
+                        hostSandboxDir: sandboxDir,
+                    }
+                } finally {
+                    this.#releaseSlot()
+                }
+            })()
+            void finalize.catch(() => undefined)
+            return finalize
+        }
+
+        return {
+            sandboxId,
+            backend: backend.name,
+            access,
+            get stdout() {
+                return processHandle.stdout
+            },
+            get stderr() {
+                return processHandle.stderr
+            },
+            stop(mode: StopMode) {
+                processHandle.stop(mode)
+            },
+            wait,
+        }
+    }
+}

package/src/sandbox/memory-limit.ts

@@ -0,0 +1,245 @@
+/*
+ * 沙箱物理内存限制。
+ *
+ * RLIMIT_AS 限制的是虚拟地址空间，会直接打死 V8 / JVM / Go 这类需要预留大块地址空间的
+ * 运行时，且并不能约束真实物理内存。这里改用 cgroup v2 的 memory.max：
+ *
+ *   1. 用 `systemd-run --user --scope -p MemoryMax=<bytes> -p MemorySwapMax=0` 起一个
+ *      临时 scope（用户 session 下无需 root）；
+ *   2. 读 systemctl 的 ControlGroup 拿到 cgroup 路径；
+ *   3. 用一层 /bin/sh 包住真正的 sandbox 命令：先把 shell 自己写进 cgroup.procs，
+ *      再 exec bwrap，保证 bwrap 及其后代一出生就在受限 cgroup 里；
+ *   4. 结束后读 memory.events 判断是否发生过 oom_kill，再停掉 scope。
+ *
+ * 宿主没有 systemd user manager（例如容器、SSH 非登录环境）时返回 inactive，
+ * 调用方退回“只限制 V8 堆 + 超时”的方式，不阻塞命令执行。
+ */
+
+import { spawn } from "node:child_process"
+import { randomUUID } from "node:crypto"
+import { accessSync, constants, readFileSync } from "node:fs"
+import { join } from "node:path"
+
+import type { LaunchSpec } from "./types.ts"
+
+export const CGROUP_V2_ROOT = "/sys/fs/cgroup"
+
+export interface MemoryLimitState {
+    /** 是否真的挂上了 cgroup 限制 */
+    readonly active: boolean
+    /** active=false 时的原因（给结果与日志用） */
+    readonly reason: string | null
+    /** 把启动命令包进受限 cgroup；inactive 时原样返回 */
+    wrap(spec: LaunchSpec): LaunchSpec
+    /** 读 oom 事件并清理 scope；可重复调用 */
+    finish(): Promise<{ oomKilled: boolean }>
+}
+
+export interface MemoryLimitCommandResult {
+    readonly code: number | null
+    readonly stdout: string
+    readonly stderr: string
+}
+
+/** 便于单测替换外部命令、文件与时间 */
+export interface MemoryLimitDeps {
+    readonly platform: NodeJS.Platform
+    readonly env: NodeJS.ProcessEnv
+    readonly run: (command: string, args: readonly string[]) => Promise<MemoryLimitCommandResult>
+    readonly spawnHolder: (command: string, args: readonly string[]) => { readonly pid?: number }
+    readonly readText: (path: string) => string | null
+    readonly writable: (path: string) => boolean
+    readonly sleep: (ms: number) => Promise<void>
+    readonly randomSuffix: () => string
+}
+
+export interface MemoryLimitOptions {
+    /** 物理内存上限，字节 */
+    readonly bytes: number
+    /** 沙箱超时；用来给 scope 的 RuntimeMaxSec 留余量，避免泄漏 */
+    readonly timeoutMs: number
+    readonly deps?: Partial<MemoryLimitDeps>
+}
+
+const CGROUP_POLL_TIMES = 10
+const CGROUP_POLL_INTERVAL_MS = 50
+const SCOPE_RUNTIME_GRACE_MS = 120_000
+
+function runCommand(command: string, args: readonly string[]): Promise<MemoryLimitCommandResult> {
+    return new Promise((resolveCommand) => {
+        const child = spawn(command, [...args], {
+            stdio: ["ignore", "pipe", "pipe"],
+            env: process.env,
+        })
+        const out: Buffer[] = []
+        const err: Buffer[] = []
+        child.stdout?.on("data", (chunk: Buffer) => out.push(chunk))
+        child.stderr?.on("data", (chunk: Buffer) => err.push(chunk))
+        child.once("error", (error) => {
+            resolveCommand({ code: null, stdout: "", stderr: error.message })
+        })
+        child.once("close", (code) => {
+            resolveCommand({
+                code,
+                stdout: Buffer.concat(out).toString("utf8"),
+                stderr: Buffer.concat(err).toString("utf8"),
+            })
+        })
+    })
+}
+
+const defaultDeps: MemoryLimitDeps = {
+    platform: process.platform,
+    env: process.env,
+    run: runCommand,
+    spawnHolder: (command, args) => {
+        const child = spawn(command, [...args], {
+            stdio: "ignore",
+            env: process.env,
+            detached: true,
+        })
+        child.unref()
+        return { pid: child.pid }
+    },
+    readText: (path) => {
+        try {
+            return readFileSync(path, "utf8")
+        } catch {
+            return null
+        }
+    },
+    writable: (path) => {
+        try {
+            accessSync(path, constants.W_OK)
+            return true
+        } catch {
+            return false
+        }
+    },
+    sleep: (ms) => new Promise((resolveSleep) => setTimeout(resolveSleep, ms)),
+    randomSuffix: () => randomUUID().slice(0, 8),
+}
+
+function inactive(reason: string): MemoryLimitState {
+    return {
+        active: false,
+        reason,
+        wrap: (spec) => spec,
+        finish: async () => ({ oomKilled: false }),
+    }
+}
+
+/** 从 memory.events 里读 oom_kill 计数 */
+export function parseOomKill(memoryEvents: string | null): boolean {
+    if (memoryEvents === null) {
+        return false
+    }
+    const match = /^oom_kill\s+(\d+)\s*$/m.exec(memoryEvents)
+    return match !== null && Number(match[1]) > 0
+}
+
+/** sandbox 命令外层包一层 sh：先把自己放进 cgroup，再 exec 原命令 */
+export function wrapSpecInCgroup(spec: LaunchSpec, cgroupPath: string): LaunchSpec {
+    return {
+        ...spec,
+        command: "/bin/sh",
+        args: [
+            "-c",
+            'echo $$ > "$1/cgroup.procs" || exit 126; shift; exec "$@"',
+            "sh",
+            cgroupPath,
+            spec.command,
+            ...spec.args,
+        ],
+    }
+}
+
+export async function createSandboxMemoryLimit(
+    options: MemoryLimitOptions,
+): Promise<MemoryLimitState> {
+    const deps: MemoryLimitDeps = { ...defaultDeps, ...options.deps }
+    const bytes = Math.trunc(options.bytes)
+    if (deps.platform !== "linux") {
+        return inactive("非 Linux 平台不使用 cgroup 内存限制")
+    }
+    if (!Number.isFinite(bytes) || bytes < 16 * 1024 * 1024) {
+        return inactive("内存上限过小，跳过 cgroup 限制")
+    }
+
+    const unit = `wellspring-sandbox-${process.pid}-${deps.randomSuffix()}`
+    const runtimeSec = Math.ceil((Math.max(1, options.timeoutMs) + SCOPE_RUNTIME_GRACE_MS) / 1000)
+    const holderArgs = [
+        "--user",
+        "--scope",
+        "--quiet",
+        "--collect",
+        "--unit",
+        unit,
+        "--property",
+        `MemoryMax=${bytes}`,
+        // 不把物理内存挤进 swap，否则 memory.max 会被换页绕过
+        "--property",
+        "MemorySwapMax=0",
+        "--property",
+        `RuntimeMaxSec=${runtimeSec}`,
+        "--",
+        "/bin/sleep",
+        "infinity",
+    ]
+    try {
+        deps.spawnHolder("systemd-run", holderArgs)
+    } catch (error) {
+        return inactive(
+            `systemd-run 不可用：${error instanceof Error ? error.message : String(error)}`,
+        )
+    }
+
+    const unitName = `${unit}.scope`
+    let cgroup: string | null = null
+    let lastError = ""
+    for (let attempt = 0; attempt < CGROUP_POLL_TIMES; attempt += 1) {
+        const shown = await deps.run("systemctl", [
+            "--user",
+            "show",
+            "-p",
+            "ControlGroup",
+            "--value",
+            unitName,
+        ])
+        const value = shown.stdout.trim()
+        if (value.startsWith("/")) {
+            cgroup = value
+            break
+        }
+        lastError = shown.stderr.trim()
+        await deps.sleep(CGROUP_POLL_INTERVAL_MS)
+    }
+    const stop = async (): Promise<void> => {
+        await deps.run("systemctl", ["--user", "stop", unitName])
+    }
+    if (cgroup === null) {
+        await stop()
+        return inactive(`systemd user scope 不可用${lastError === "" ? "" : `：${lastError}`}`)
+    }
+
+    const cgroupPath = join(CGROUP_V2_ROOT, cgroup)
+    if (!deps.writable(join(cgroupPath, "cgroup.procs"))) {
+        await stop()
+        return inactive(`无法写入 cgroup：${cgroupPath}`)
+    }
+
+    let finished: Promise<{ oomKilled: boolean }> | null = null
+    return {
+        active: true,
+        reason: null,
+        wrap: (spec) => wrapSpecInCgroup(spec, cgroupPath),
+        finish: () => {
+            finished ??= (async () => {
+                const oomKilled = parseOomKill(deps.readText(join(cgroupPath, "memory.events")))
+                await stop()
+                return { oomKilled }
+            })()
+            return finished
+        },
+    }
+}