T-HUB/apps/media-worker/src/supervisor/sandbox.ts

import { AsyncLocalStorage } from "node:async_hooks";
import { createHash } from "node:crypto";
import type * as ChildProcessApi from "node:child_process";
import { chmodSync, mkdirSync, constants, openSync, closeSync, fchmodSync } from "node:fs";
import { resolve } from "node:path";
import { Transform } from "node:stream";
import type { DesiredChannel } from "../types";
import type { RelayJob } from "./index";

const childProcess: { spawn: typeof ChildProcessApi.spawn } = require("node:child_process");

// Resolve the CommonJS package from either src/supervisor or dist/supervisor.
// A small structural boundary keeps worker typechecking independent of package build order.
interface MediaCommand { binary: "ffmpeg" | "ffprobe"; argv: readonly string[] }
interface SandboxModule {
  MediaCommandBuilder: new () => { supervisor(binary: MediaCommand["binary"], argv: readonly string[], scratch: string): MediaCommand };
  SandboxUnavailable: new () => Error;
  spawnSandbox(command: MediaCommand, options: { rootfs: string; cgroupRoot: string; scratchDir: string; limits: { wallTimeMs: number } }, spawnOptions: ChildProcessApi.SpawnOptions | undefined, spawner: typeof ChildProcessApi.spawn): { child: ChildProcessApi.ChildProcess };
}
const sandbox = (): SandboxModule => require(resolve(__dirname, "../../../../packages/media-sandbox/dist")) as SandboxModule;
interface Context { scratch: string }
const context = new AsyncLocalStorage<Context>();
let installed = false;

/**
 * The only common spawn path lives in relay/, outside this change's boundary.
 * A scoped Node spawn adapter covers existing typed FFmpeg builders, including
 * asynchronous swaps/playlist producers, without altering their process API.
 * Non-media children and calls outside a supervisor job retain normal spawning.
 */
function installAdapter(): void {
  if (installed) return;
  const original = childProcess.spawn;
  childProcess.spawn = ((command: string, args?: readonly string[] | ChildProcessApi.SpawnOptions, options?: ChildProcessApi.SpawnOptions) => {
    const job = context.getStore();
    if (!job || !/(?:^|\/)(?:ffmpeg|ffprobe)$/.test(command)) {
      return typeof args === "object" && !Array.isArray(args)
        ? original(command, args as ChildProcessApi.SpawnOptions)
        : original(command, (args as readonly string[] | undefined) ?? [], options ?? {});
    }
    if (!Array.isArray(args)) throw new Error("media argv required");
    const runtime = sandbox();
    const rootfs = process.env.MEDIA_SANDBOX_ROOTFS;
    const cgroupRoot = process.env.MEDIA_SANDBOX_CGROUP_ROOT;
    if (!rootfs || !cgroupRoot) throw new runtime.SandboxUnavailable();
    const media = new runtime.MediaCommandBuilder().supervisor(command.endsWith("ffprobe") ? "ffprobe" : "ffmpeg", args, job.scratch);
    // Make only generated output directories writable by the mapped UID.
    // Traverse through directory fds: a compromised sibling cannot race a
    // symlink replacement into chmod of a host file.
    for (const path of media.argv.filter(value => value.startsWith("file:/scratch/"))) {
      const parts = path.slice("file:/scratch/".length).split("/").slice(0, -1).filter(part => part && part !== ".");
      if (parts.length > 32 || parts.some(part => part === "..")) throw new runtime.SandboxUnavailable();
      let fd = openSync(job.scratch, constants.O_RDONLY | constants.O_DIRECTORY | constants.O_NOFOLLOW);
      try {
        fchmodSync(fd, 0o777);
        for (const part of parts) {
          const next = openSync(`/proc/self/fd/${fd}/${part}`, constants.O_RDONLY | constants.O_DIRECTORY | constants.O_NOFOLLOW);
          closeSync(fd); fd = next; fchmodSync(fd, 0o777);
        }
      } finally { closeSync(fd); }
    }
    const child = runtime.spawnSandbox(media, {
      rootfs, cgroupRoot, scratchDir: job.scratch,
      // Long-running channels still have a finite job lease; uploads use the package defaults.
      limits: { wallTimeMs: Number(process.env.MEDIA_SANDBOX_WALL_TIME_MS ?? 86_400_000) },
    }, options, original).child;
    // Replace the public stderr stream BEFORE spawnManaged attaches its logger.
    // Arbitrary media bytes cannot bypass a caller-supplied passthrough redactor.
    if (child.stderr) {
      const safe = new Transform({ transform(_chunk, _encoding, done) { done(null, Buffer.from("[media stderr suppressed]\n")); } });
      child.stderr.pipe(safe); child.stderr = safe; child.stdio[2] = safe;
    }
    return child;
  }) as typeof childProcess.spawn;
  installed = true;
}
function enabled(): boolean {
  const flag = process.env.MEDIA_SANDBOX;
  if (flag !== undefined && flag !== "on" && flag !== "off") throw new Error("MEDIA_SANDBOX must be on or off");
  if (process.env.NODE_ENV === "production" && flag === "off") throw new Error("production requires media sandbox");
  return flag === "on" || process.env.NODE_ENV === "production";
}
export function supervisorSandboxEnabled(): boolean { return enabled(); }

export function sandboxRelay(job: RelayJob, channel: DesiredChannel): RelayJob {
  if (!enabled()) return job;
  installAdapter();
  const id = createHash("sha256").update(`${channel.workspaceId}:${channel.channelId}`).digest("hex").slice(0, 32);
  const parent = resolve(process.env.MEDIA_SCRATCH_DIR ?? ".media-worker-state/scratch");
  const scratch = resolve(parent, id);
  mkdirSync(parent, { recursive: true }); chmodSync(parent, 0o700);
  mkdirSync(scratch, { recursive: true }); chmodSync(scratch, 0o777);
  const run = <T>(action: () => T): T => context.run({ scratch }, action);
  return {
    start: config => run(() => job.start(config)),
    stop: () => run(() => job.stop()),
    ...(job.reconfigure ? { reconfigure: (config: DesiredChannel) => run(() => job.reconfigure!(config)) } : {}),
  };
}