ycode/src/main/sandbox/workspace-files.ts

import { promises as fs, constants } from 'node:fs'
import { dirname, join, relative, resolve } from 'node:path'
import type { SandboxEffectivePolicy } from '@shared/sandbox'
import type { FileChange, DirectoryChange } from './transactions'
import { digest } from './transactions'
import { sandboxConfiguration } from './configuration'
import { assertActive, resolveSandboxPath, within } from './policy'
import { sandboxPathDenied } from './policy-store'

export type Snapshot = Map<string, { hash: string; size: number }>
const excludedDirectories = new Set(['.git', '.ycode', '.tmp', 'node_modules', '.venv', 'venv', '__pycache__', 'out', 'dist', 'release'])
export async function regularBytes(file: string, fileLimit: number): Promise<Buffer> {
  const before = await fs.lstat(file)
  if (!before.isFile() || before.isSymbolicLink() || before.nlink !== 1 || before.size > fileLimit) throw new Error('快照包含链接、特殊文件或超过策略上限的文件')
  const handle = await fs.open(file, constants.O_RDONLY)
  try {
    const opened = await handle.stat()
    if (opened.ino !== before.ino || opened.dev !== before.dev || opened.nlink !== 1 || opened.size !== before.size) throw new Error('文件在读取前发生变化')
    const bytes = await handle.readFile()
    const after = await handle.stat()
    if (bytes.length !== before.size || after.mtimeMs !== opened.mtimeMs || after.size !== opened.size) throw new Error('文件在快照读取期间发生变化')
    return bytes
  } finally { await handle.close() }
}
export async function snapshot(source: string, target: string | undefined, policy: SandboxEffectivePolicy, signal?: AbortSignal): Promise<{ files: Snapshot; directories: Set<string>; excluded: number }> {
  const files: Snapshot = new Map()
  const directories = new Set<string>()
  let totalBytes = 0, excluded = 0, entries = 0
  const stateRoot = resolve(sandboxConfiguration().stateRoot)
  async function walk(directory: string): Promise<void> {
    for await (const entry of await fs.opendir(directory)) {
      assertActive(signal)
      if (++entries > policy.limits.files) throw new Error('工作区超过沙盒快照的总项数上限')
      const sourcePath = join(directory, entry.name), file = relative(source, sourcePath)
      if (within(stateRoot, sourcePath) || sandboxPathDenied(policy, sourcePath) || entry.isSymbolicLink() || entry.isDirectory() && excludedDirectories.has(entry.name.toLowerCase())) { excluded++; continue }
      await resolveSandboxPath(source, file)
      if (entry.isDirectory()) { directories.add(file); if (target) await fs.mkdir(join(target, file), { recursive: true }); await walk(sourcePath) }
      else if (entry.isFile()) {
        if (files.size >= policy.limits.files) throw new Error('工作区超过沙盒快照的文件数量上限')
        const bytes = await regularBytes(sourcePath, policy.limits.fileBytes)
        totalBytes += bytes.length
        if (totalBytes > policy.limits.bytes) throw new Error('工作区超过沙盒快照的容量上限')
        if (target) {
          await fs.mkdir(dirname(join(target, file)), { recursive: true })
          await fs.writeFile(join(target, file), bytes, { flag: 'wx' })
        }
        files.set(file, { hash: digest(bytes), size: bytes.length })
      } else throw new Error('工作区包含不支持的特殊文件')
    }
  }
  await walk(source)
  return { files, directories, excluded }
}
export async function stageChanges(stageRoot: string, before: Snapshot, beforeDirectories: Set<string>, policy: SandboxEffectivePolicy, nodeDependencies: boolean, signal?: AbortSignal): Promise<{ files: FileChange[]; directories: DirectoryChange[] }> {
  const after = new Map<string, Buffer>(), changes: FileChange[] = []
  const directories = new Set<string>()
  let totalBytes = 0, entries = 0
  async function walk(directory: string): Promise<void> {
    for await (const entry of await fs.opendir(directory)) {
      assertActive(signal)
      if (++entries > policy.limits.files) throw new Error('沙盒产物超过总项数上限')
      const absolute = join(directory, entry.name), file = relative(stageRoot, absolute)
      if (nodeDependencies && file === 'node_modules') continue
      if (sandboxPathDenied(policy, join(policy.workspaceRoot, file)) || entry.isSymbolicLink()) throw new Error('沙盒产生了受保护路径或链接；未写回工作区')
      await resolveSandboxPath(stageRoot, file)
      if (entry.isDirectory()) { directories.add(file); await walk(absolute) }
      else if (entry.isFile()) {
        if (after.size >= policy.limits.files) throw new Error('沙盒产物超过文件数量上限')
        const bytes = await regularBytes(absolute, policy.limits.fileBytes)
        totalBytes += bytes.length
        if (totalBytes > policy.limits.bytes) throw new Error('沙盒产物超过容量上限')
        after.set(file, bytes)
      } else throw new Error('沙盒产生了不支持的特殊文件')
    }
  }
  await walk(stageRoot)
  const directoryChanges: DirectoryChange[] = []
  for (const directory of beforeDirectories) if (!directories.has(directory)) directoryChanges.push({ path: directory, before: true, after: false })
  for (const directory of directories) if (!beforeDirectories.has(directory)) directoryChanges.push({ path: directory, before: false, after: true })
  for (const [file, bytes] of after) {
    const old = before.get(file)
    if (old?.hash !== digest(bytes)) changes.push({ path: file, beforeHash: old?.hash ?? null, after: bytes })
  }
  for (const [file, old] of before) if (!after.has(file)) changes.push({ path: file, beforeHash: old.hash, after: null })
  return { files: changes, directories: directoryChanges }
}