rea/src/artifacts/SafeOutputTree.ts

import { createHash } from "node:crypto";
import { constants } from "node:fs";
import {
  lstat,
  mkdir,
  open,
  realpath,
  rm,
  type FileHandle,
} from "node:fs/promises";
import { basename, dirname, isAbsolute, join, resolve } from "node:path";
import type { Readable } from "node:stream";
import { streamChunkToBuffer } from "./StreamBytes.js";

import {
  ArtifactPathRegistry,
  normalizeArtifactPath,
} from "./ArtifactPaths.js";
import { ArtifactReaderFailure } from "./ArtifactReader.js";

/** One file durably written to an operation-owned output tree. */
export interface SafeOutputFile {
  readonly relativePath: string;
  readonly sha256: string;
  readonly bytesWritten: number;
}

/** Cleanup state established while rolling back an uncommitted tree. */
export type SafeOutputCleanup =
  | { readonly status: "not-required" }
  | { readonly status: "complete"; readonly residualPaths: readonly [] }
  | {
      readonly status: "incomplete";
      readonly residualPaths: readonly [string, ...string[]];
    };

/** Symlink-resistant materialization in an exclusively owned, initially absent tree. */
export class SafeOutputTree {
  readonly #registry = new ArtifactPathRegistry();
  readonly #outputRoot: string;
  #published = false;
  #cleanup: SafeOutputCleanup = {
    status: "not-required",
  };

  private constructor(
    outputRoot: string,
    private readonly platform: NodeJS.Platform,
  ) {
    this.#outputRoot = outputRoot;
  }

  /**
   * Exclusively create the absent destination as this operation's owned tree.
   *
   * POSIX mode bits and directory `fchmod`/`fsync` have no Windows equivalent:
   * `mkdir` already applies the requested mode, so the redundant handle
   * `chmod` is skipped there to avoid an `EPERM` on the directory descriptor.
   */
  static async create(
    outputRoot: string,
    platform: NodeJS.Platform = process.platform,
  ): Promise<SafeOutputTree> {
    if (!isAbsolute(outputRoot))
      throw new ArtifactReaderFailure(
        "path",
        "Extraction output root must be absolute",
      );
    const requested = resolve(outputRoot);
    const name = basename(requested);
    if (name === "." || name === "..")
      throw new ArtifactReaderFailure("path", "Invalid extraction output root");
    const parent = await realpath(dirname(requested)).catch(
      (cause: unknown) => {
        throw new ArtifactReaderFailure(
          "unavailable",
          "Extraction output parent is unavailable",
          { cause },
        );
      },
    );
    const canonicalOutput = join(parent, name);
    await mkdir(canonicalOutput, { mode: 0o700 }).catch((cause: unknown) => {
      if (isAlreadyExists(cause))
        throw new ArtifactReaderFailure(
          "path",
          "Extraction output root already exists",
          { cause },
        );
      throw cause;
    });
    try {
      if (platform !== "win32") {
        const stagingHandle = await open(
          canonicalOutput,
          constants.O_RDONLY | constants.O_DIRECTORY,
        );
        try {
          await stagingHandle.chmod(0o700);
        } finally {
          await stagingHandle.close();
        }
      }
      return new SafeOutputTree(canonicalOutput, platform);
    } catch (cause: unknown) {
      await rm(canonicalOutput, { recursive: true, force: true });
      throw cause;
    }
  }

  get outputRoot(): string {
    return this.#outputRoot;
  }

  get cleanup(): SafeOutputCleanup {
    return structuredClone(this.#cleanup);
  }

  /** Stream one regular file with exact byte and digest verification. */
  async write(
    relativePath: string,
    source: Readable,
    expectedSha256: string,
    signal?: AbortSignal,
  ): Promise<SafeOutputFile> {
    this.#assertWritable();
    const path = normalizeArtifactPath(relativePath);
    this.#registry.add(path, "file");
    const destination = await this.#prepareParent(path);
    const handle = await open(
      destination,
      constants.O_CREAT |
        constants.O_EXCL |
        constants.O_WRONLY |
        constants.O_NOFOLLOW,
      0o600,
    ).catch((cause: unknown) => {
      throw new ArtifactReaderFailure(
        "path",
        `Could not exclusively create extraction path: ${path}`,
        { cause },
      );
    });
    const hash = createHash("sha256");
    let bytes = 0;
    try {
      for await (const raw of source) {
        abortIfNeeded(signal);
        const chunk = streamChunkToBuffer(raw);
        bytes += chunk.length;
        hash.update(chunk);
        await writeAll(handle, chunk);
      }
      const sha256 = hash.digest("hex");
      if (sha256 !== expectedSha256)
        throw new ArtifactReaderFailure(
          "integrity",
          `Extracted content disagrees with inventory: ${path}`,
        );
      await handle.sync();
      await handle.close();
      const readback = await hashFile(destination, bytes, signal);
      if (readback.sha256 !== sha256 || readback.bytes !== bytes)
        throw new ArtifactReaderFailure(
          "integrity",
          `Durable readback verification failed: ${path}`,
        );
      return { relativePath: path, sha256, bytesWritten: bytes };
    } catch (cause: unknown) {
      // best-effort cleanup: file-handle close must not mask the write failure.
      await handle.close().catch(() => undefined);
      throw cause;
    }
  }

  /** Sync the owned output tree and prevent further writes through this instance. */
  async commit(): Promise<void> {
    this.#assertWritable();
    // Windows has no directory fsync; file contents are already synced in write().
    if (this.platform === "win32") {
      this.#published = true;
      return;
    }
    const parent = await open(
      dirname(this.#outputRoot),
      constants.O_RDONLY | constants.O_DIRECTORY,
    );
    let output: FileHandle | undefined;
    try {
      output = await open(
        this.#outputRoot,
        constants.O_RDONLY | constants.O_DIRECTORY,
      );
      await output.sync();
      await parent.sync();
      this.#published = true;
    } catch (cause: unknown) {
      throw new ArtifactReaderFailure(
        "path",
        "Could not durably sync extraction output",
        { cause },
      );
    } finally {
      await Promise.allSettled([parent.close(), output?.close()]);
    }
  }

  /** Remove only this operation's unsealed tree and verify absence. */
  async rollback(): Promise<SafeOutputCleanup> {
    if (this.#published) return structuredClone(this.#cleanup);
    await rm(this.#outputRoot, { recursive: true, force: true });
    const absent = await isAbsent(this.#outputRoot);
    this.#cleanup = absent
      ? { status: "complete", residualPaths: [] }
      : {
          status: "incomplete",
          residualPaths: [basename(this.#outputRoot)],
        };
    if (!absent)
      throw new ArtifactReaderFailure(
        "integrity",
        "Extraction output cleanup could not be verified",
      );
    return structuredClone(this.#cleanup);
  }

  async #prepareParent(relativePath: string): Promise<string> {
    const parts = relativePath.split("/");
    const fileName = parts.pop();
    if (fileName === undefined)
      throw new ArtifactReaderFailure("path", "Invalid extraction path");
    let current = this.#outputRoot;
    for (const part of parts) {
      current = join(current, part);
      await mkdir(current, { mode: 0o700 }).catch((cause: unknown) => {
        if (!isAlreadyExists(cause)) throw cause;
      });
      const metadata = await lstat(current);
      if (!metadata.isDirectory() || metadata.isSymbolicLink())
        throw new ArtifactReaderFailure(
          "path",
          `Unsafe extraction parent: ${relativePath}`,
        );
    }
    return join(current, fileName);
  }

  #assertWritable(): void {
    if (this.#published)
      throw new ArtifactReaderFailure(
        "integrity",
        "Extraction tree is already committed",
      );
  }
}

const hashFile = async (
  path: string,
  maximum: number,
  signal?: AbortSignal,
): Promise<{ readonly sha256: string; readonly bytes: number }> => {
  const handle = await open(path, constants.O_RDONLY | constants.O_NOFOLLOW);
  const hash = createHash("sha256");
  let bytes = 0;
  try {
    for await (const raw of handle.createReadStream({ autoClose: false })) {
      abortIfNeeded(signal);
      const chunk = streamChunkToBuffer(raw);
      bytes += chunk.length;
      if (bytes > maximum)
        throw new ArtifactReaderFailure(
          "integrity",
          "Readback exceeded the bytes written",
        );
      hash.update(chunk);
    }
  } finally {
    await handle.close();
  }
  return { sha256: hash.digest("hex"), bytes };
};

const writeAll = async (
  handle: Awaited<ReturnType<typeof open>>,
  chunk: Buffer,
): Promise<void> => {
  let offset = 0;
  while (offset < chunk.length) {
    const { bytesWritten } = await handle.write(
      chunk,
      offset,
      chunk.length - offset,
    );
    if (bytesWritten === 0)
      throw new ArtifactReaderFailure(
        "unavailable",
        "Extraction output stopped accepting bytes",
      );
    offset += bytesWritten;
  }
};

const isAbsent = async (path: string): Promise<boolean> =>
  lstat(path).then(
    () => false,
    (cause: unknown) => {
      if (isNotFound(cause)) return true;
      throw cause;
    },
  );

const isNotFound = (cause: unknown): boolean =>
  cause instanceof Error && "code" in cause && cause.code === "ENOENT";

const isAlreadyExists = (cause: unknown): boolean =>
  cause instanceof Error && "code" in cause && cause.code === "EEXIST";

const abortIfNeeded = (signal?: AbortSignal): void => {
  if (signal?.aborted === true)
    throw new ArtifactReaderFailure(
      "cancelled",
      "Artifact extraction cancelled",
    );
};