zsjz-ai/ai-electron/electron/service/codex/codexArtifactService.ts

import { promises as fs } from 'node:fs';
import { basename, extname, isAbsolute, join, resolve, sep } from 'node:path';
import { dialog, shell } from 'electron';
import { logger } from 'ee-core/log';
import { getCodexDataDir } from './codexHome';

/**
 * 产物（Codex 在会话工作目录里生成的文件）的读写服务。
 *
 * 安全边界：
 *  1. **只允许访问「代码决定的工作目录」** —— 案件工作空间与客户端默认目录。
 *     白名单支持事后补登记（见 requireAllowedRoot），但仅限绝对路径且真实存在的目录；
 *  2. 所有相对路径都做 resolve + 前缀校验，挡住 `..` 穿越与绝对路径注入。
 *
 * 只做「列目录 / 读内容 / 另存 / 打开」，不提供删除与写入：
 * 产物是 Codex 的工作成果，客户端不应替用户改它。
 */

/** 单次列目录返回的条目上限 */
const MAX_ENTRIES = 2000;
/** 文本预览字节上限 */
const MAX_TEXT_BYTES = 2 * 1024 * 1024;
/** 二进制预览字节上限（图片 / PDF） */
const MAX_BINARY_BYTES = 16 * 1024 * 1024;

/** 已注册的会话工作目录（规范化后的小写路径，Windows 下大小写不敏感） */
const allowedRoots = new Set<string>();
/** 白名单落盘位置：客户端重启后历史会话仍要能读产物 */
const ROOTS_FILE = () => join(getCodexDataDir(), 'artifact-roots.json');
/** 白名单条目上限，超出丢最早的 */
const MAX_ROOTS = 50;

let rootsLoaded = false;

function normalizePath(value: string): string {
  const resolved = resolve(value);
  return process.platform === 'win32' ? resolved.toLowerCase() : resolved;
}

/** 首次使用时把落盘的白名单读回内存 */
async function ensureRootsLoaded(): Promise<void> {
  if (rootsLoaded) return;
  rootsLoaded = true;
  try {
    const raw = await fs.readFile(ROOTS_FILE(), 'utf8');
    const parsed = JSON.parse(raw);
    if (Array.isArray(parsed)) {
      for (const item of parsed) {
        if (typeof item === 'string' && item.trim()) allowedRoots.add(normalizePath(item));
      }
    }
  } catch {
    // 首次运行 / 文件损坏：当作空白名单
  }
}

async function persistRoots(): Promise<void> {
  try {
    const list = [...allowedRoots].slice(-MAX_ROOTS);
    await fs.mkdir(getCodexDataDir(), { recursive: true });
    await fs.writeFile(ROOTS_FILE(), JSON.stringify(list, null, 2), 'utf8');
  } catch {
    // 落盘失败只影响下次启动，不阻塞当前会话
  }
}

/** 登记一个可访问的会话工作目录，返回解析后的绝对路径 */
export function registerArtifactRoot(dir: string | null | undefined): string | null {
  if (!dir || !String(dir).trim()) return null;
  const resolved = resolve(String(dir).trim());
  const key = normalizePath(resolved);
  if (!allowedRoots.has(key)) {
    allowedRoots.add(key);
    void persistRoots();
  }
  return resolved;
}

/** 目录名合法性：必须存在且是目录，才允许进入白名单 */
async function isUsableDirectory(dir: string): Promise<boolean> {
  try {
    const stat = await fs.stat(dir);
    return stat.isDirectory();
  } catch {
    return false;
  }
}

/**
 * 解析并校验访问根目录。
 *
 * 白名单**不再要求事先登记**：工作目录是由代码决定的（案件工作空间 / 客户端默认目录），
 * 而产物栏在「还没建会话」时就要列出案件工作空间的数据文件，
 * 那时 threadStart 还没执行过、目录自然没登记过 —— 以前就会在这里被拒。
 *
 * 因此改成：不在白名单但**是绝对路径且确实存在的目录**时补登记（记一条日志便于排查），
 * 不满足这两个条件仍然拒绝。真正防越界的还是下面 resolveInside 的路径穿越校验。
 */
async function requireAllowedRoot(root: string): Promise<string> {
  const raw = String(root ?? '').trim();
  if (!raw) throw new Error('缺少会话工作目录');
  const resolved = resolve(raw);

  await ensureRootsLoaded();
  if (allowedRoots.has(normalizePath(resolved))) return resolved;

  if (!isAbsolute(raw) || !(await isUsableDirectory(resolved))) {
    throw new Error(`工作目录不存在或不可访问：${resolved}`);
  }
  registerArtifactRoot(resolved);
  logger.warn(`[codexArtifact] 补登记工作目录：${resolved}`);
  return resolved;
}

/** 把相对路径解析到 root 之内；越界（.. / 绝对路径）直接抛错 */
function resolveInside(root: string, relativePath: string | undefined): string {
  const base = resolve(root);
  const target = resolve(base, String(relativePath ?? '').replace(/^[/\\]+/, '') || '.');
  const prefix = base.endsWith(sep) ? base : `${base}${sep}`;
  if (target !== base && !target.startsWith(prefix)) {
    throw new Error('路径超出会话工作目录');
  }
  return target;
}

/** 相对路径统一成 POSIX 分隔符（渲染进程侧只认 /） */
function toPosix(value: string): string {
  return value.split(sep).join('/');
}

/* ------------------------------------------------------------------ *
 * 类型判定（与前端 utils/codexArtifact.ts 保持同一套扩展名表）
 * ------------------------------------------------------------------ */

const IMAGE_EXT = ['png', 'jpg', 'jpeg', 'gif', 'webp', 'bmp', 'svg', 'ico'];
const CODE_EXT = [
  'js', 'mjs', 'cjs', 'ts', 'tsx', 'jsx', 'vue', 'py', 'java', 'go', 'rs', 'rb', 'php',
  'c', 'h', 'cpp', 'cs', 'sh', 'ps1', 'bat', 'sql', 'yml', 'yaml', 'toml', 'ini', 'xml',
  'less', 'scss', 'css',
];
const TABLE_EXT = ['csv', 'tsv', 'xlsx', 'xls'];
const TEXT_EXT = ['txt', 'log', 'rst', 'conf', 'env'];

export type ArtifactKind =
  | 'dir' | 'markdown' | 'html' | 'code' | 'json' | 'image' | 'pdf' | 'table' | 'text' | 'binary';

export function artifactKindOf(name: string, isDir = false): ArtifactKind {
  if (isDir) return 'dir';
  const ext = extname(name).replace('.', '').toLowerCase();
  if (ext === 'md' || ext === 'markdown') return 'markdown';
  if (ext === 'html' || ext === 'htm') return 'html';
  if (ext === 'pdf') return 'pdf';
  if (ext === 'json') return 'json';
  if (IMAGE_EXT.includes(ext)) return 'image';
  if (TABLE_EXT.includes(ext)) return 'table';
  if (CODE_EXT.includes(ext)) return 'code';
  if (TEXT_EXT.includes(ext)) return 'text';
  return 'binary';
}

function mimeOf(name: string): string {
  const ext = extname(name).replace('.', '').toLowerCase();
  const map: Record<string, string> = {
    png: 'image/png',
    jpg: 'image/jpeg',
    jpeg: 'image/jpeg',
    gif: 'image/gif',
    webp: 'image/webp',
    bmp: 'image/bmp',
    svg: 'image/svg+xml',
    ico: 'image/x-icon',
    pdf: 'application/pdf',
    json: 'application/json',
    csv: 'text/csv',
    md: 'text/markdown',
    txt: 'text/plain',
  };
  return map[ext] ?? 'application/octet-stream';
}

/* ------------------------------------------------------------------ *
 * 对外能力
 * ------------------------------------------------------------------ */

export interface ArtifactEntry {
  name: string;
  relative: string;
  isDir: boolean;
  size: number;
  mtime: string;
  kind: ArtifactKind;
}

export interface ArtifactListResult {
  root: string;
  relative: string;
  entries: ArtifactEntry[];
  truncated: boolean;
}

/** 列出目录下的条目（不递归）；目录不存在时返回空列表而不是报错 */
export async function listArtifacts(root: string, dir?: string): Promise<ArtifactListResult> {
  const base = await requireAllowedRoot(root);
  const relativeDir = toPosix(String(dir ?? '').replace(/^[/\\]+/, '')).replace(/\/+$/, '');
  const target = resolveInside(base, relativeDir);

  let names: string[];
  try {
    names = await fs.readdir(target);
  } catch (error: any) {
    if (error?.code === 'ENOENT') return { root: base, relative: relativeDir, entries: [], truncated: false };
    throw new Error(`无法读取目录：${error?.message ?? error}`);
  }

  const truncated = names.length > MAX_ENTRIES;
  const picked = truncated ? names.slice(0, MAX_ENTRIES) : names;
  const entries: ArtifactEntry[] = [];

  for (const name of picked) {
    const absolute = resolveInside(base, relativeDir ? `${relativeDir}/${name}` : name);
    try {
      const stat = await fs.stat(absolute);
      const isDir = stat.isDirectory();
      entries.push({
        name,
        relative: toPosix(relativeDir ? `${relativeDir}/${name}` : name),
        isDir,
        size: isDir ? 0 : stat.size,
        mtime: stat.mtime.toISOString(),
        kind: artifactKindOf(name, isDir),
      });
    } catch {
      // 单个条目读不到（被删除 / 无权限）不影响整次列目录
    }
  }

  return { root: base, relative: relativeDir, entries, truncated };
}

export interface ArtifactTextResult {
  content: string;
  size: number;
  truncated: boolean;
}

/** 读文本内容，超出上限只取前一段 */
export async function readArtifactText(
  root: string,
  relativePath: string,
  maxBytes = MAX_TEXT_BYTES,
): Promise<ArtifactTextResult> {
  const base = await requireAllowedRoot(root);
  const target = resolveInside(base, relativePath);
  const stat = await fs.stat(target);
  if (stat.isDirectory()) throw new Error('目标是目录，无法作为文件读取');

  const limit = Math.max(1024, Math.min(Number(maxBytes) || MAX_TEXT_BYTES, MAX_TEXT_BYTES));
  const handle = await fs.open(target, 'r');
  try {
    const length = Math.min(stat.size, limit);
    const buffer = Buffer.alloc(length);
    await handle.read(buffer, 0, length, 0);
    return {
      // 二进制文件按 utf8 读出来会是乱码，但预览侧只对文本类调用本方法
      content: buffer.toString('utf8'),
      size: stat.size,
      truncated: stat.size > length,
    };
  } finally {
    await handle.close();
  }
}

export interface ArtifactBinaryResult {
  base64: string;
  mime: string;
  size: number;
  truncated: boolean;
}

/** 读二进制内容（图片 / PDF），超出上限时截断并标记 */
export async function readArtifactBase64(
  root: string,
  relativePath: string,
  maxBytes = MAX_BINARY_BYTES,
): Promise<ArtifactBinaryResult> {
  const base = await requireAllowedRoot(root);
  const target = resolveInside(base, relativePath);
  const stat = await fs.stat(target);
  if (stat.isDirectory()) throw new Error('目标是目录，无法作为文件读取');

  const limit = Math.max(1024, Math.min(Number(maxBytes) || MAX_BINARY_BYTES, MAX_BINARY_BYTES));
  const length = Math.min(stat.size, limit);
  const handle = await fs.open(target, 'r');
  try {
    const buffer = Buffer.alloc(length);
    await handle.read(buffer, 0, length, 0);
    return {
      base64: buffer.toString('base64'),
      mime: mimeOf(target),
      size: stat.size,
      truncated: stat.size > length,
    };
  } finally {
    await handle.close();
  }
}

/** 另存为：弹系统保存对话框并复制。用户取消时 saved=false */
export async function saveArtifactAs(
  root: string,
  relativePath: string,
  suggestedName?: string,
): Promise<{ saved: boolean; target: string | null }> {
  const base = await requireAllowedRoot(root);
  const source = resolveInside(base, relativePath);
  const stat = await fs.stat(source);
  if (stat.isDirectory()) throw new Error('暂不支持导出目录');

  const result = await dialog.showSaveDialog({
    title: '另存为',
    defaultPath: basename(suggestedName || source),
  });
  if (result.canceled || !result.filePath) return { saved: false, target: null };

  await fs.copyFile(source, result.filePath);
  return { saved: true, target: result.filePath };
}

/** 在系统文件管理器中定位该产物 */
export async function revealArtifact(root: string, relativePath: string): Promise<void> {
  const base = await requireAllowedRoot(root);
  const target = resolveInside(base, relativePath);
  if (!(await pathExists(target))) throw new Error('文件不存在');
  shell.showItemInFolder(target);
}

/** 用系统默认程序打开该产物 */
export async function openArtifact(root: string, relativePath: string): Promise<void> {
  const base = await requireAllowedRoot(root);
  const target = resolveInside(base, relativePath);
  if (!(await pathExists(target))) throw new Error('文件不存在');
  const message = await shell.openPath(target);
  if (message) throw new Error(message);
}

/** 用系统文件管理器打开会话工作目录本身 */
export async function openArtifactRoot(root: string): Promise<void> {
  const base = await requireAllowedRoot(root);
  const message = await shell.openPath(base);
  if (message) throw new Error(message);
}

async function pathExists(target: string): Promise<boolean> {
  try {
    await fs.access(target);
    return true;
  } catch {
    return false;
  }
}