kuber/tests/server/materialize.test.ts

import { afterEach, describe, expect, test } from "bun:test";
import { createHash } from "node:crypto";
import { lstat, mkdtemp, readFile, readlink, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { FilesystemCas } from "../../server/cas";
import {
  materializeWorkspace,
  parseWorkspaceManifest,
} from "../../server/materialize";
import {
  BUILD_PROTOCOL_VERSION,
  type Sha256Digest,
  type WorkspaceManifest,
} from "../../shared/build-protocol";

const roots: string[] = [];
afterEach(async () => {
  await Promise.all(
    roots.splice(0).map((root) => rm(root, { recursive: true, force: true })),
  );
});

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

describe("source materialization", () => {
  test("collects aggregate CAS and filesystem timing with manifest counts", async () => {
    const root = await mkdtemp(join(tmpdir(), "kuber-materialize-"));
    roots.push(root);
    const data = Buffer.from("hello");
    const manifest: WorkspaceManifest = {
      version: BUILD_PROTOCOL_VERSION,
      files: [
        {
          path: "example",
          type: "file",
          digest: digest(data),
          size: data.byteLength,
          mode: 0o644,
        },
      ],
    };
    const manifestData = Buffer.from(JSON.stringify(manifest));
    const workspace = digest(manifestData);
    let release!: () => void;
    let entered!: () => void;
    const blocked = new Promise<void>((resolve) => {
      release = resolve;
    });
    const started = new Promise<void>((resolve) => {
      entered = resolve;
    });
    const timing = { casReadMs: 0, fsWriteMs: 0 };
    const submission = materializeWorkspace(
      {
        get: async (requested) => {
          if (requested === workspace) return manifestData;
          entered();
          await blocked;
          return data;
        },
      },
      workspace,
      join(root, "workspace"),
      timing,
    );
    await started;
    await Bun.sleep(25);
    release();
    await submission;
    expect(timing).toMatchObject({
      fileCount: 1,
      manifestBytes: manifestData.byteLength,
      fileBytes: data.byteLength,
    });
    expect(timing.casReadMs).toBeGreaterThan(15);
    expect(timing.fsWriteMs).toBeGreaterThan(0);
  });

  test("bounds concurrent CAS reads while materializing many files", async () => {
    const root = await mkdtemp(join(tmpdir(), "kuber-materialize-"));
    roots.push(root);
    const data = Buffer.from("x");
    const blobDigest = digest(data);
    const manifest: WorkspaceManifest = {
      version: BUILD_PROTOCOL_VERSION,
      files: Array.from({ length: 40 }, (_, index) => ({
        path: `file-${String(index).padStart(2, "0")}`,
        type: "file" as const,
        digest: blobDigest,
        size: data.byteLength,
        mode: 0o644 as const,
      })),
    };
    const manifestBytes = Buffer.from(JSON.stringify(manifest));
    const workspace = digest(manifestBytes);
    let inflight = 0;
    let maxInflight = 0;
    const read = async <T>(result: T): Promise<T> => {
      inflight += 1;
      maxInflight = Math.max(maxInflight, inflight);
      await Bun.sleep(2);
      inflight -= 1;
      return result;
    };
    const cas = {
      has: async () => read(true),
      get: async (requested: Sha256Digest) =>
        requested === workspace ? manifestBytes : read(data),
    };

    await materializeWorkspace(cas, workspace, join(root, "workspace"));

    expect(maxInflight).toBeGreaterThan(1);
    expect(maxInflight).toBeLessThanOrEqual(20);
  });

  test("materializes files and safe symlinks with declared modes", async () => {
    const root = await mkdtemp(join(tmpdir(), "kuber-materialize-"));
    roots.push(root);
    const cas = new FilesystemCas(join(root, "cas"));
    const executable = Buffer.from("#!/bin/sh\necho ok\n");
    const link = Buffer.from("bin/run");
    await cas.put(executable, digest(executable));
    await cas.put(link, digest(link));
    const manifest: WorkspaceManifest = {
      version: BUILD_PROTOCOL_VERSION,
      files: [
        {
          path: "bin/run",
          type: "file",
          digest: digest(executable),
          size: executable.byteLength,
          mode: 0o755,
        },
        {
          path: "run",
          type: "symlink",
          digest: digest(link),
          size: link.byteLength,
          mode: 0o777,
        },
      ],
    };
    const manifestBytes = Buffer.from(JSON.stringify(manifest));
    const workspace = await cas.put(manifestBytes);
    const destination = join(root, "workspaces", "build-1");

    expect(await materializeWorkspace(cas, workspace, destination)).toEqual(
      manifest,
    );
    expect(await readFile(join(destination, "bin/run"), "utf8")).toContain(
      "echo ok",
    );
    expect((await lstat(join(destination, "bin/run"))).mode & 0o777).toBe(
      0o755,
    );
    expect(await readlink(join(destination, "run"))).toBe("bin/run");
  });

  test("rejects traversal, path collisions, unsafe links, and size mismatches atomically", async () => {
    expect(() =>
      parseWorkspaceManifest(
        Buffer.from(
          JSON.stringify({
            version: 1,
            files: [
              {
                path: "../x",
                type: "file",
                digest: `sha256:${"a".repeat(64)}`,
                size: 0,
                mode: 420,
              },
            ],
          }),
        ),
      ),
    ).toThrow("invalid file");
    expect(() =>
      parseWorkspaceManifest(
        Buffer.from(
          JSON.stringify({
            version: 1,
            files: [
              {
                path: "a",
                type: "file",
                digest: `sha256:${"a".repeat(64)}`,
                size: 0,
                mode: 420,
              },
              {
                path: "a/b",
                type: "file",
                digest: `sha256:${"b".repeat(64)}`,
                size: 0,
                mode: 420,
              },
            ],
          }),
        ),
      ),
    ).toThrow("conflicts");
    for (const [files, conflict] of [
      [["a/b/c", "a"], "Workspace path conflicts with a directory: a"],
      [["a/b/c", "a/b"], "Workspace path conflicts with a directory: a/b"],
      [["a", "a/b/c"], "Workspace path conflicts with a file: a/b/c"],
    ] as const) {
      expect(() =>
        parseWorkspaceManifest(
          Buffer.from(
            JSON.stringify({
              version: BUILD_PROTOCOL_VERSION,
              files: files.map((path) => ({
                path,
                type: "file",
                digest: `sha256:${"a".repeat(64)}`,
                size: 0,
                mode: 0o644,
              })),
            }),
          ),
        ),
      ).toThrow(conflict);
    }

    const root = await mkdtemp(join(tmpdir(), "kuber-materialize-"));
    roots.push(root);
    const cas = new FilesystemCas(join(root, "cas"));
    const link = Buffer.from("../../outside");
    const linkDigest = await cas.put(link);
    const manifest = Buffer.from(
      JSON.stringify({
        version: 1,
        files: [
          {
            path: "nested/link",
            type: "symlink",
            digest: linkDigest,
            size: link.byteLength,
            mode: 0o777,
          },
        ],
      }),
    );
    const workspace = await cas.put(manifest);
    const destination = join(root, "workspace");
    await expect(
      materializeWorkspace(cas, workspace, destination),
    ).rejects.toThrow("Unsafe symlink");
    await expect(lstat(destination)).rejects.toMatchObject({ code: "ENOENT" });
  });

  test("accepts a large manifest with distinct nested paths", () => {
    const manifest: WorkspaceManifest = {
      version: BUILD_PROTOCOL_VERSION,
      files: Array.from({ length: 3_000 }, (_, index) => ({
        path: `dir-${index}/nested/file`,
        type: "file" as const,
        digest: `sha256:${"a".repeat(64)}` as Sha256Digest,
        size: 0,
        mode: 0o644 as const,
      })),
    };
    expect(
      parseWorkspaceManifest(Buffer.from(JSON.stringify(manifest))),
    ).toEqual(manifest);
  });
});