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>;
}

export interface MaterializeTiming {
  fileCount?: number;
  manifestBytes?: number;
  fileBytes?: number;
  // Sum of operation durations; concurrent operations can exceed wall time.
  // CAS get includes its content hash verification.
  casReadMs: number;
  fsWriteMs: number;
}

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);
  }
  const ancestors = new Set<string>();
  for (const path of paths) {
    for (
      let separator = path.indexOf("/");
      separator !== -1;
      separator = path.indexOf("/", separator + 1)
    ) {
      ancestors.add(path.slice(0, separator));
    }
  }
  for (const path of paths) {
    if (ancestors.has(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,
  timing?: MaterializeTiming,
): Promise<WorkspaceManifest> {
  const read = timing
    ? async (digest: Sha256Digest) => {
        const start = performance.now();
        try {
          return await cas.get(digest);
        } finally {
          timing.casReadMs += performance.now() - start;
        }
      }
    : (digest: Sha256Digest) => cas.get(digest);
  const write = timing
    ? async <T>(operation: () => Promise<T>): Promise<T> => {
        const start = performance.now();
        try {
          return await operation();
        } finally {
          timing.fsWriteMs += performance.now() - start;
        }
      }
    : <T>(operation: () => Promise<T>) => operation();
  assertSha256Digest(manifestDigest);
  const manifestData = await read(manifestDigest);
  if (timing) timing.manifestBytes = manifestData.byteLength;
  const manifest = parseWorkspaceManifest(manifestData);
  if (timing) {
    timing.fileCount = manifest.files.length;
    timing.fileBytes = manifest.files.reduce(
      (total, file) => total + file.size,
      0,
    );
  }

  await write(() => mkdir(dirname(destination), { recursive: true }));
  try {
    await write(() => 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 write(() => 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 write(() =>
        mkdir(dirname(target), { recursive: true, mode: 0o755 }),
      );
      const data = await read(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 write(() => symlink(link, target));
      } else {
        const handle = await write(() =>
          open(
            target,
            constants.O_CREAT | constants.O_EXCL | constants.O_WRONLY,
            file.mode,
          ),
        );
        try {
          await write(() => handle.writeFile(data));
        } finally {
          await write(() => handle.close());
        }
        await write(() => chmod(target, file.mode));
      }
    });
    await write(() => rename(temporary, destination));
  } catch (error) {
    await write(() => rm(temporary, { recursive: true, force: true }));
    throw error;
  }
  return manifest;
}