kuber/server/materialize.ts

import { randomUUID } from "node:crypto";
import { constants } from "node:fs";
import {
  chmod,
  lstat,
  mkdir,
  open,
  rename,
  rm,
  symlink,
} from "node:fs/promises";
import {
  basename,
  dirname,
  isAbsolute,
  join,
  relative,
  resolve,
} from "node:path";
import {
  BUILD_PROTOCOL_VERSION,
  assertSha256Digest,
  type Sha256Digest,
  type WorkspaceFile,
  type WorkspaceManifest,
} from "../shared/build-protocol";

export interface MaterializeCas {
  get(digest: Sha256Digest): Promise<Uint8Array>;
}

const MATERIALIZE_CONCURRENCY = 20;

async function mapConcurrent<T, R>(
  values: T[],
  run: (value: T) => Promise<R>,
): Promise<R[]> {
  const results = new Array<R>(values.length);
  let index = 0;
  const worker = async () => {
    for (;;) {
      const current = index++;
      if (current >= values.length) return;
      results[current] = await run(values[current]!);
    }
  };
  await Promise.all(
    Array.from(
      { length: Math.min(MATERIALIZE_CONCURRENCY, values.length) },
      worker,
    ),
  );
  return results;
}

function safePath(path: string): boolean {
  return (
    path.length > 0 &&
    !isAbsolute(path) &&
    !path.includes("\\") &&
    !path.includes("\0") &&
    path
      .split("/")
      .every((part) => part !== "" && part !== "." && part !== "..")
  );
}

export function parseWorkspaceManifest(data: Uint8Array): WorkspaceManifest {
  let value: unknown;
  try {
    value = JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(data));
  } catch {
    throw new Error("Workspace manifest is not valid UTF-8 JSON");
  }
  const manifest = value as Partial<WorkspaceManifest>;
  if (
    manifest.version !== BUILD_PROTOCOL_VERSION ||
    !Array.isArray(manifest.files)
  ) {
    throw new Error("Unsupported workspace manifest");
  }
  const paths = new Set<string>();
  for (const file of manifest.files as WorkspaceFile[]) {
    if (
      !file ||
      !safePath(file.path) ||
      (file.type !== "file" && file.type !== "symlink") ||
      !Number.isSafeInteger(file.size) ||
      file.size < 0 ||
      ![0o644, 0o755, 0o777].includes(file.mode)
    ) {
      throw new Error("Workspace manifest contains an invalid file");
    }
    assertSha256Digest(file.digest);
    if (paths.has(file.path))
      throw new Error(`Duplicate workspace path: ${file.path}`);
    for (const parent of dirname(file.path).split("/")) {
      if (parent && paths.has(parent)) {
        throw new Error(`Workspace path conflicts with a file: ${file.path}`);
      }
    }
    paths.add(file.path);
  }
  for (const path of paths) {
    if ([...paths].some((other) => other.startsWith(`${path}/`))) {
      throw new Error(`Workspace path conflicts with a directory: ${path}`);
    }
  }
  return manifest as WorkspaceManifest;
}

function safeSymlinkTarget(filePath: string, target: string): boolean {
  if (
    !target ||
    isAbsolute(target) ||
    target.includes("\\") ||
    target.includes("\0")
  )
    return false;
  const resolved = resolve("/workspace", dirname(filePath), target);
  return resolved === "/workspace" || resolved.startsWith("/workspace/");
}

export async function materializeWorkspace(
  cas: MaterializeCas,
  manifestDigest: Sha256Digest,
  destination: string,
): Promise<WorkspaceManifest> {
  assertSha256Digest(manifestDigest);
  const manifest = parseWorkspaceManifest(await cas.get(manifestDigest));

  await mkdir(dirname(destination), { recursive: true });
  try {
    await lstat(destination);
    throw new Error(`Workspace destination already exists: ${destination}`);
  } catch (error) {
    if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
  }
  const temporary = join(
    dirname(destination),
    `.${basename(destination)}.${process.pid}.${randomUUID()}.tmp`,
  );
  await mkdir(temporary, { mode: 0o755 });
  try {
    await mapConcurrent(manifest.files, async (file) => {
      const target = join(temporary, file.path);
      if (relative(temporary, target).startsWith(".."))
        throw new Error("Unsafe workspace path");
      await mkdir(dirname(target), { recursive: true, mode: 0o755 });
      const data = await cas.get(file.digest);
      if (data.byteLength !== file.size) {
        throw new Error(`Workspace blob size mismatch for ${file.path}`);
      }
      if (file.type === "symlink") {
        const link = new TextDecoder("utf-8", { fatal: true }).decode(data);
        if (!safeSymlinkTarget(file.path, link))
          throw new Error(`Unsafe symlink target for ${file.path}`);
        await symlink(link, target);
      } else {
        const handle = await open(
          target,
          constants.O_CREAT | constants.O_EXCL | constants.O_WRONLY,
          file.mode,
        );
        try {
          await handle.writeFile(data);
        } finally {
          await handle.close();
        }
        await chmod(target, file.mode);
      }
    });
    await rename(temporary, destination);
  } catch (error) {
    await rm(temporary, { recursive: true, force: true });
    throw error;
  }
  return manifest;
}