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>;
  has(digest: Sha256Digest): Promise<boolean>;
}

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));
  const missing: Sha256Digest[] = [];
  for (const file of manifest.files)
    if (!(await cas.has(file.digest))) missing.push(file.digest);
  if (missing.length)
    throw new Error(`Workspace blobs are missing: ${missing.join(", ")}`);

  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 {
    for (const file of manifest.files) {
      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;
}