kuber/lib/workspace.ts

import { execFile } from "node:child_process";
import { createHash } from "node:crypto";
import { constants } from "node:fs";
import {
  chmod,
  lstat,
  mkdir,
  open,
  readdir,
  readlink,
  realpath,
  symlink,
} from "node:fs/promises";
import type { Stats } from "node:fs";
import { dirname, isAbsolute, relative, resolve, sep } from "node:path";
import { promisify } from "node:util";
import {
  BUILD_PROTOCOL_VERSION,
  assertSha256Digest,
  type Sha256Digest,
  type WorkspaceFile,
  type WorkspaceManifest,
} from "../shared/build-protocol";

const execFileAsync = promisify(execFile);

export type WorkspaceBlob = {
  digest: Sha256Digest;
  data: Uint8Array;
};

export type WorkspaceSnapshot = {
  manifest: WorkspaceManifest;
  digest: Sha256Digest;
  blobs: WorkspaceBlob[];
};

export type BlobReader =
  | ((digest: Sha256Digest) => Promise<Uint8Array>)
  | { get(digest: Sha256Digest): Promise<Uint8Array> };

function digest(data: Uint8Array | string): Sha256Digest {
  return `sha256:${createHash("sha256").update(data).digest("hex")}`;
}

export function validateWorkspacePath(path: string): void {
  if (
    !path ||
    path.includes("\0") ||
    path.includes("\\") ||
    isAbsolute(path) ||
    path.split("/").some((part) => !part || part === "." || part === "..")
  ) {
    throw new Error(`Unsafe workspace path: ${JSON.stringify(path)}`);
  }
}

function isWithin(root: string, candidate: string): boolean {
  const path = relative(root, candidate);
  return (
    path === "" ||
    (!path.startsWith(`..${sep}`) && path !== ".." && !isAbsolute(path))
  );
}

async function gitFiles(root: string, args: string[]): Promise<string[]> {
  const { stdout } = await execFileAsync(
    "git",
    ["-C", root, "ls-files", "-z", ...args],
    {
      encoding: "buffer",
      maxBuffer: 64 * 1024 * 1024,
    },
  );
  const decoder = new TextDecoder("utf-8", { fatal: true });
  const files: string[] = [];
  let start = 0;
  for (
    let end = stdout.indexOf(0);
    end !== -1;
    end = stdout.indexOf(0, start)
  ) {
    if (end > start) files.push(decoder.decode(stdout.subarray(start, end)));
    start = end + 1;
  }
  return files;
}

async function selectedFiles(root: string): Promise<string[]> {
  const [normal, dotenv] = await Promise.all([
    gitFiles(root, ["--cached", "--others", "--exclude-standard"]),
    gitFiles(root, [
      "--others",
      "--ignored",
      "--exclude-standard",
      "--",
      ".env*",
      "**/.env*",
    ]),
  ]);
  return [...new Set([...normal, ...dotenv])].sort((a, b) =>
    Buffer.from(a).compare(Buffer.from(b)),
  );
}

async function isIgnored(root: string, path: string): Promise<boolean> {
  try {
    await execFileAsync("git", ["-C", root, "check-ignore", "-q", "--", path]);
    return true;
  } catch (error) {
    if ((error as { code?: number }).code === 1) return false;
    throw error;
  }
}

async function rejectSelectedSpecialFiles(
  root: string,
  directory = root,
): Promise<void> {
  for (const entry of await readdir(directory, { withFileTypes: true })) {
    if (directory === root && entry.name === ".git") continue;
    const source = resolve(directory, entry.name);
    const path = relative(root, source).split(sep).join("/");
    if (entry.isDirectory()) {
      if (!(await isIgnored(root, path)))
        await rejectSelectedSpecialFiles(root, source);
      continue;
    }
    if (entry.isFile() || entry.isSymbolicLink()) continue;
    if (!(await isIgnored(root, path)) || entry.name.startsWith(".env"))
      throw new Error(`Special files are not allowed in workspaces: ${path}`);
  }
}

function canonicalManifest(manifest: WorkspaceManifest): string {
  return JSON.stringify({
    version: manifest.version,
    files: manifest.files.map((file) => ({
      path: file.path,
      type: file.type,
      digest: file.digest,
      size: file.size,
      mode: file.mode,
    })),
  });
}

export function serializeWorkspaceManifest(
  manifest: WorkspaceManifest,
): Uint8Array {
  validateWorkspaceManifest(manifest);
  return Buffer.from(canonicalManifest(manifest));
}

export function workspaceManifestDigest(
  manifest: WorkspaceManifest,
): Sha256Digest {
  return digest(serializeWorkspaceManifest(manifest));
}

export function validateWorkspaceManifest(manifest: WorkspaceManifest): void {
  if (
    manifest.version !== BUILD_PROTOCOL_VERSION ||
    !Array.isArray(manifest.files)
  ) {
    throw new Error("Unsupported workspace manifest");
  }

  let previous = "";
  const seen = new Set<string>();
  for (const file of manifest.files) {
    validateWorkspacePath(file.path);
    assertSha256Digest(file.digest);
    if (!Number.isSafeInteger(file.size) || file.size < 0)
      throw new Error(`Invalid size for ${file.path}`);
    if (
      (file.type === "file" && file.mode !== 0o644 && file.mode !== 0o755) ||
      (file.type === "symlink" && file.mode !== 0o777)
    ) {
      throw new Error(`Invalid mode for ${file.path}`);
    }
    if (file.type !== "file" && file.type !== "symlink")
      throw new Error(`Invalid entry type for ${file.path}`);
    if (seen.has(file.path))
      throw new Error(`Duplicate workspace path: ${file.path}`);
    for (const parent of file.path
      .split("/")
      .slice(0, -1)
      .map((_, index, parts) => parts.slice(0, index + 1).join("/"))) {
      if (seen.has(parent))
        throw new Error(`Workspace entry is used as a directory: ${parent}`);
    }
    if (previous && Buffer.from(previous).compare(Buffer.from(file.path)) >= 0)
      throw new Error("Workspace files must be bytewise sorted");
    seen.add(file.path);
    previous = file.path;
  }
}

export async function enumerateWorkspace(
  root: string,
): Promise<WorkspaceSnapshot> {
  const repository = await realpath(root);
  await rejectSelectedSpecialFiles(repository);
  const paths = await selectedFiles(repository);
  const files: WorkspaceFile[] = [];
  const blobs = new Map<Sha256Digest, Uint8Array>();

  for (const path of paths) {
    validateWorkspacePath(path);
    const source = resolve(repository, path);
    if (!isWithin(repository, source))
      throw new Error(`Workspace path escapes root: ${path}`);

    let stat: Stats;
    try {
      stat = await lstat(source);
    } catch (error) {
      if ((error as NodeJS.ErrnoException).code === "ENOENT") continue;
      throw error;
    }

    let data: Uint8Array;
    let entry: WorkspaceFile;
    if (stat.isSymbolicLink()) {
      const target = await readlink(source);
      if (
        isAbsolute(target) ||
        !isWithin(repository, resolve(dirname(source), target))
      )
        throw new Error(`Symlink escapes workspace: ${path} -> ${target}`);
      data = Buffer.from(target);
      entry = {
        path,
        type: "symlink",
        digest: digest(data),
        size: data.byteLength,
        mode: 0o777,
      };
    } else if (stat.isFile()) {
      let mode: 0o644 | 0o755;
      const handle = await open(
        source,
        constants.O_RDONLY | constants.O_NOFOLLOW,
      );
      try {
        const opened = await handle.stat();
        if (!opened.isFile()) throw new Error(`Not a regular file: ${path}`);
        mode = opened.mode & 0o111 ? 0o755 : 0o644;
        data = await handle.readFile();
      } finally {
        await handle.close();
      }
      entry = {
        path,
        type: "file",
        digest: digest(data),
        size: data.byteLength,
        mode,
      };
    } else {
      throw new Error(`Special files are not allowed in workspaces: ${path}`);
    }
    files.push(entry);
    blobs.set(entry.digest, data);
  }

  const manifest = {
    version: BUILD_PROTOCOL_VERSION,
    files,
  } satisfies WorkspaceManifest;
  return {
    manifest,
    digest: workspaceManifestDigest(manifest),
    blobs: [...blobs].map(([blobDigest, data]) => ({
      digest: blobDigest,
      data,
    })),
  };
}

async function ensureParentDirectories(
  root: string,
  path: string,
): Promise<void> {
  let current = root;
  for (const part of path.split("/").slice(0, -1)) {
    current = resolve(current, part);
    try {
      const stat = await lstat(current);
      if (!stat.isDirectory())
        throw new Error(`Workspace parent is not a directory: ${path}`);
    } catch (error) {
      if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
      await mkdir(current, { mode: 0o755 });
    }
  }
}

async function readBlob(
  reader: BlobReader,
  blobDigest: Sha256Digest,
): Promise<Uint8Array> {
  return typeof reader === "function"
    ? reader(blobDigest)
    : reader.get(blobDigest);
}

export async function materializeWorkspace(
  destination: string,
  manifest: WorkspaceManifest,
  reader: BlobReader,
): Promise<void> {
  validateWorkspaceManifest(manifest);
  await mkdir(destination, { recursive: false, mode: 0o755 });
  const root = await realpath(destination);

  for (const file of manifest.files.filter((entry) => entry.type === "file")) {
    await ensureParentDirectories(root, file.path);
    const data = await readBlob(reader, file.digest);
    if (data.byteLength !== file.size || digest(data) !== file.digest)
      throw new Error(`Blob verification failed for ${file.path}`);
    const target = resolve(root, file.path);
    const handle = await open(
      target,
      constants.O_CREAT | constants.O_EXCL | constants.O_WRONLY,
      file.mode,
    );
    try {
      await handle.writeFile(data);
      await handle.sync();
    } finally {
      await handle.close();
    }
    await chmod(target, file.mode);
  }

  for (const file of manifest.files.filter(
    (entry) => entry.type === "symlink",
  )) {
    await ensureParentDirectories(root, file.path);
    const data = await readBlob(reader, file.digest);
    if (data.byteLength !== file.size || digest(data) !== file.digest)
      throw new Error(`Blob verification failed for ${file.path}`);
    const linkTarget = Buffer.from(data).toString("utf8");
    const target = resolve(root, file.path);
    if (
      linkTarget.includes("\0") ||
      isAbsolute(linkTarget) ||
      !isWithin(root, resolve(dirname(target), linkTarget))
    )
      throw new Error(
        `Symlink escapes workspace: ${file.path} -> ${linkTarget}`,
      );
    await symlink(linkTarget, target);
  }
}