TJWaterAgent/src/runtime/conversationWorkspace.ts

import { randomUUID } from "node:crypto";
import {
  chown,
  chmod,
  lstat,
  mkdir,
  realpath,
  rename,
  rm,
  stat,
  writeFile,
} from "node:fs/promises";
import { isAbsolute, join, relative, resolve, sep } from "node:path";

export const SANDBOX_UID = 10_001;
export const SANDBOX_GID = 10_001;

const isOutsideRoot = (root: string, candidate: string) => {
  const relativePath = relative(root, candidate);
  return (
    relativePath === ".." ||
    relativePath.startsWith(`..${sep}`) ||
    isAbsolute(relativePath)
  );
};

export const setSandboxOwnership = async (path: string) => {
  if (typeof process.getuid !== "function" || process.getuid() !== 0) {
    return;
  }
  await chown(path, SANDBOX_UID, SANDBOX_GID);
};

export const resolveConversationWorkspace = async (
  workspaceDirectory: string,
  importRoot: string,
) => {
  const workspaceLinkStat = await lstat(workspaceDirectory);
  if (workspaceLinkStat.isSymbolicLink()) {
    throw new Error("conversation workspace must not be a symbolic link");
  }
  const [resolvedImportRoot, resolvedWorkspaceDirectory] = await Promise.all([
    realpath(importRoot),
    realpath(workspaceDirectory),
  ]);
  if (isOutsideRoot(resolvedImportRoot, resolvedWorkspaceDirectory)) {
    throw new Error("conversation workspace must be inside RESULT_REF_IMPORT_DIR");
  }
  const relativeWorkspace = relative(
    resolvedImportRoot,
    resolvedWorkspaceDirectory,
  );
  if (
    !relativeWorkspace ||
    relativeWorkspace.includes("/") ||
    relativeWorkspace.includes("\\")
  ) {
    throw new Error(
      "conversation workspace must be a direct child of RESULT_REF_IMPORT_DIR",
    );
  }
  const workspaceStat = await stat(resolvedWorkspaceDirectory);
  if (!workspaceStat.isDirectory()) {
    throw new Error("conversation workspace must point to a directory");
  }
  return resolvedWorkspaceDirectory;
};

export const resolveExistingPathInsideRoot = async (
  filePath: string,
  rootPath: string,
  outsideMessage = "path must be inside the current conversation workspace",
) => {
  if (!isAbsolute(filePath)) {
    throw new Error("file_path must be absolute");
  }
  const [resolvedFilePath, resolvedRootPath] = await Promise.all([
    realpath(filePath),
    realpath(rootPath),
  ]);
  if (isOutsideRoot(resolvedRootPath, resolvedFilePath)) {
    throw new Error(outsideMessage);
  }
  return resolvedFilePath;
};

export type StagedConversationFile = {
  bytes: number;
  contentType: string;
  filePath: string;
};

export const stageConversationText = async (
  workspaceDirectory: string,
  importRoot: string,
  content: string,
  options: {
    contentType: string;
    extension: string;
    prefix: string;
  },
): Promise<StagedConversationFile> => {
  const workspace = await resolveConversationWorkspace(
    workspaceDirectory,
    importRoot,
  );
  const outputDirectory = resolve(workspace, "tool-data");
  if (isOutsideRoot(workspace, outputDirectory)) {
    throw new Error("tool data directory escaped the conversation workspace");
  }
  await mkdir(outputDirectory, { mode: 0o700, recursive: true });
  const resolvedOutputDirectory = await realpath(outputDirectory);
  if (isOutsideRoot(workspace, resolvedOutputDirectory)) {
    throw new Error("tool data directory escaped the conversation workspace");
  }
  await chmod(resolvedOutputDirectory, 0o700);
  await setSandboxOwnership(resolvedOutputDirectory);

  const fileName = `${options.prefix}-${randomUUID()}.${options.extension}`;
  const filePath = join(resolvedOutputDirectory, fileName);
  const temporaryPath = `${filePath}.${randomUUID()}.tmp`;
  try {
    await writeFile(temporaryPath, content, { encoding: "utf8", mode: 0o600 });
    await setSandboxOwnership(temporaryPath);
    await rename(temporaryPath, filePath);
  } catch (error) {
    await rm(temporaryPath, { force: true }).catch(() => undefined);
    throw error;
  }
  await chmod(filePath, 0o600);
  await setSandboxOwnership(filePath);
  return {
    bytes: Buffer.byteLength(content),
    contentType: options.contentType,
    filePath,
  };
};