279 lines
13 KiB
TypeScript
279 lines
13 KiB
TypeScript
import fs, { type FileHandle } from "node:fs/promises";
|
|
import type { Stats } from "node:fs";
|
|
import path from "node:path";
|
|
|
|
import type { DreaminaGenerationInput } from "./dreamina-cli-contract.js";
|
|
import { DreaminaCliError } from "./dreamina-cli-process.js";
|
|
import type { StateLockLease } from "./dreamina-cli-state.js";
|
|
|
|
type ReferenceKind = "image" | "video" | "audio";
|
|
type ReferenceInput = DreaminaGenerationInput & {
|
|
referenceImages?: string[];
|
|
referenceVideos?: string[];
|
|
referenceAudios?: string[];
|
|
};
|
|
type PreparedReference = { handle: FileHandle; before: Stats; sourcePath: string; canonicalPath: string; extension: string; closed: boolean };
|
|
|
|
const IMAGE_MAX_BYTES = 32 * 1024 * 1024;
|
|
const MEDIA_MAX_BYTES = 512 * 1024 * 1024;
|
|
const STAGING_MAX_BYTES = 1024 * 1024 * 1024;
|
|
const STAGING_PREFIX = ".dreamina-references-";
|
|
const STAGING_DIRECTORY_PATTERN = /^\.dreamina-references-[A-Za-z0-9]{6}$/;
|
|
const STAGING_STALE_MS = 24 * 60 * 60 * 1000;
|
|
|
|
export async function stageReferences(
|
|
input: DreaminaGenerationInput,
|
|
roots: readonly string[],
|
|
ownerStateDirectory: string,
|
|
signal?: AbortSignal,
|
|
) {
|
|
throwIfCancelled(signal);
|
|
const source = input as ReferenceInput;
|
|
if (!source.referenceImages?.length && !source.referenceVideos?.length && !source.referenceAudios?.length) {
|
|
return { input, cleanup: async () => undefined };
|
|
}
|
|
const canonicalRoots: string[] = [];
|
|
try {
|
|
for (const root of roots) {
|
|
try {
|
|
const canonical = await fs.realpath(root);
|
|
if (!(await fs.stat(canonical)).isDirectory()) throw new Error("not a directory");
|
|
canonicalRoots.push(canonical);
|
|
} catch (error) {
|
|
if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
|
|
}
|
|
}
|
|
} catch {
|
|
throw referenceInvalid();
|
|
}
|
|
if (!canonicalRoots.length) throw referenceInvalid();
|
|
|
|
const prepared: PreparedReference[] = [];
|
|
let stagingDirectory = "";
|
|
try {
|
|
let totalBytes = 0;
|
|
const prepareGroup = async (values: string[] | undefined, kind: ReferenceKind) => {
|
|
if (!values) return undefined;
|
|
const result: PreparedReference[] = [];
|
|
for (const value of values) {
|
|
throwIfCancelled(signal);
|
|
const reference = await prepareReference(value, kind, canonicalRoots);
|
|
prepared.push(reference);
|
|
result.push(reference);
|
|
totalBytes += reference.before.size;
|
|
if (totalBytes > STAGING_MAX_BYTES) throw referenceBudgetExceeded();
|
|
}
|
|
return result;
|
|
};
|
|
const preparedImages = await prepareGroup(source.referenceImages, "image");
|
|
const preparedVideos = await prepareGroup(source.referenceVideos, "video");
|
|
const preparedAudios = await prepareGroup(source.referenceAudios, "audio");
|
|
throwIfCancelled(signal);
|
|
await fs.mkdir(ownerStateDirectory, { recursive: true, mode: 0o700 });
|
|
stagingDirectory = await fs.mkdtemp(path.join(ownerStateDirectory, STAGING_PREFIX));
|
|
await fs.chmod(stagingDirectory, 0o700);
|
|
const stageGroup = async (values: PreparedReference[] | undefined) => {
|
|
if (!values) return undefined;
|
|
const result: string[] = [];
|
|
for (const value of values) {
|
|
result.push(await stagePreparedReference(value, stagingDirectory, signal));
|
|
}
|
|
return result;
|
|
};
|
|
const referenceImages = await stageGroup(preparedImages);
|
|
const referenceVideos = await stageGroup(preparedVideos);
|
|
const referenceAudios = await stageGroup(preparedAudios);
|
|
return {
|
|
input: {
|
|
...input,
|
|
...(referenceImages ? { referenceImages } : {}),
|
|
...(referenceVideos ? { referenceVideos } : {}),
|
|
...(referenceAudios ? { referenceAudios } : {}),
|
|
} as DreaminaGenerationInput,
|
|
cleanup: () => removeStagingDirectory(stagingDirectory),
|
|
};
|
|
} catch (error) {
|
|
let cleanupFailed = false;
|
|
try { await closePreparedReferences(prepared); } catch { cleanupFailed = true; }
|
|
if (stagingDirectory) {
|
|
try { await fs.rm(stagingDirectory, { recursive: true, force: true }); }
|
|
catch { cleanupFailed = true; }
|
|
}
|
|
if (cleanupFailed) throw referenceCleanupFailed();
|
|
if (error instanceof DreaminaCliError) throw error;
|
|
throw referenceInvalid();
|
|
}
|
|
}
|
|
|
|
async function prepareReference(value: string, kind: ReferenceKind, roots: readonly string[]): Promise<PreparedReference> {
|
|
if (!path.isAbsolute(value) || value.includes("\0") || /^(?:\\\\|\/\/|\\\\[.?]\\)/.test(value)
|
|
|| value.split(/[\\/]+/).includes("..")) {
|
|
throw referenceInvalid();
|
|
}
|
|
let handle: FileHandle | undefined;
|
|
try {
|
|
const pathBefore = await fs.lstat(value);
|
|
if (pathBefore.isSymbolicLink()) throw new Error("link");
|
|
const canonical = await fs.realpath(value);
|
|
if (!roots.some((root) => isWithin(root, canonical))) throw new Error("outside root");
|
|
const canonicalBefore = await fs.stat(canonical);
|
|
if (!sameFileSnapshot(pathBefore, canonicalBefore)) throw new Error("reference path changed");
|
|
handle = await fs.open(canonical, "r");
|
|
const before = await handle.stat();
|
|
if (!sameFileSnapshot(pathBefore, before) || !sameFileSnapshot(canonicalBefore, before)) throw new Error("reference path changed");
|
|
const maxBytes = kind === "image" ? IMAGE_MAX_BYTES : MEDIA_MAX_BYTES;
|
|
if (!before.isFile() || before.nlink !== 1 || before.size < 4 || before.size > maxBytes) {
|
|
throw new Error("invalid file");
|
|
}
|
|
const header = Buffer.alloc(16);
|
|
const { bytesRead } = await handle.read(header, 0, header.length, 0);
|
|
const extension = path.extname(canonical).toLowerCase();
|
|
if (!validMedia(kind, extension, header.subarray(0, bytesRead))) throw new Error("invalid media");
|
|
return { handle, before, sourcePath: value, canonicalPath: canonical, extension, closed: false };
|
|
} catch (error) {
|
|
if (handle) {
|
|
try { await handle.close(); } catch { throw referenceCleanupFailed(); }
|
|
}
|
|
if (error instanceof DreaminaCliError) throw error;
|
|
throw referenceInvalid();
|
|
}
|
|
}
|
|
|
|
async function stagePreparedReference(reference: PreparedReference, stagingDirectory: string, signal?: AbortSignal) {
|
|
let stagedPath = "";
|
|
let staged: FileHandle | undefined;
|
|
let failure: unknown;
|
|
try {
|
|
throwIfCancelled(signal);
|
|
stagedPath = path.join(stagingDirectory, `${crypto.randomUUID()}${reference.extension}`);
|
|
staged = await fs.open(stagedPath, "wx", 0o600);
|
|
const buffer = Buffer.allocUnsafe(Math.min(1024 * 1024, reference.before.size));
|
|
let position = 0;
|
|
while (position < reference.before.size) {
|
|
throwIfCancelled(signal);
|
|
const length = Math.min(buffer.length, reference.before.size - position);
|
|
const copied = await reference.handle.read(buffer, 0, length, position);
|
|
if (copied.bytesRead !== length) throw new Error("reference changed");
|
|
let written = 0;
|
|
while (written < copied.bytesRead) {
|
|
throwIfCancelled(signal);
|
|
const result = await staged.write(buffer, written, copied.bytesRead - written, position + written);
|
|
if (result.bytesWritten <= 0 || result.bytesWritten > copied.bytesRead - written) {
|
|
throw new Error("staged write made no progress");
|
|
}
|
|
written += result.bytesWritten;
|
|
}
|
|
position += copied.bytesRead;
|
|
}
|
|
await staged.sync();
|
|
if ((await staged.stat()).size !== reference.before.size) throw new Error("staged size mismatch");
|
|
const after = await reference.handle.stat();
|
|
const pathAfter = await fs.lstat(reference.sourcePath);
|
|
const canonicalAfter = await fs.realpath(reference.sourcePath);
|
|
if (pathAfter.isSymbolicLink()
|
|
|| canonicalAfter !== reference.canonicalPath
|
|
|| !sameFileSnapshot(reference.before, after)
|
|
|| !sameFileSnapshot(pathAfter, after)) {
|
|
throw new Error("reference changed");
|
|
}
|
|
} catch (error) {
|
|
failure = error;
|
|
}
|
|
let cleanupFailed = false;
|
|
if (staged) {
|
|
try { await staged.close(); } catch { cleanupFailed = true; }
|
|
}
|
|
try { await closePreparedReference(reference); } catch { cleanupFailed = true; }
|
|
if (failure && stagedPath) {
|
|
try { await fs.rm(stagedPath, { force: true }); } catch { cleanupFailed = true; }
|
|
}
|
|
if (cleanupFailed) throw referenceCleanupFailed();
|
|
if (failure instanceof DreaminaCliError) throw failure;
|
|
if (failure) throw referenceInvalid();
|
|
return stagedPath;
|
|
}
|
|
|
|
async function closePreparedReference(reference: PreparedReference) {
|
|
if (reference.closed) return;
|
|
reference.closed = true;
|
|
await reference.handle.close();
|
|
}
|
|
|
|
async function closePreparedReferences(references: readonly PreparedReference[]) {
|
|
let failed = false;
|
|
for (const reference of references) {
|
|
try { await closePreparedReference(reference); } catch { failed = true; }
|
|
}
|
|
if (failed) throw referenceCleanupFailed();
|
|
}
|
|
|
|
async function removeStagingDirectory(directory: string) {
|
|
try { await fs.rm(directory, { recursive: true, force: true }); }
|
|
catch { throw referenceCleanupFailed(); }
|
|
}
|
|
|
|
export async function scavengeStaleStagingDirectories(directory: string, lease: StateLockLease) {
|
|
let entries;
|
|
try { entries = await fs.readdir(directory, { withFileTypes: true }); }
|
|
catch (error) {
|
|
if ((error as NodeJS.ErrnoException).code === "ENOENT") return;
|
|
throw referenceCleanupFailed();
|
|
}
|
|
for (const entry of entries) {
|
|
if (!entry.isDirectory() || !STAGING_DIRECTORY_PATTERN.test(entry.name)) continue;
|
|
const candidate = path.join(directory, entry.name);
|
|
let stats: Stats;
|
|
try { stats = await fs.lstat(candidate); }
|
|
catch (error) {
|
|
if ((error as NodeJS.ErrnoException).code === "ENOENT") continue;
|
|
throw referenceCleanupFailed();
|
|
}
|
|
if (!stats.isDirectory() || stats.isSymbolicLink() || Date.now() - stats.mtimeMs <= STAGING_STALE_MS) continue;
|
|
await lease.assertOwned();
|
|
try { await fs.rm(candidate, { recursive: true, force: true }); }
|
|
catch { throw referenceCleanupFailed(); }
|
|
}
|
|
}
|
|
|
|
function sameFileSnapshot(left: Stats, right: Stats) {
|
|
return left.isFile() && right.isFile()
|
|
&& left.dev === right.dev
|
|
&& left.ino === right.ino
|
|
&& left.nlink === right.nlink
|
|
&& left.size === right.size
|
|
&& left.mtimeMs === right.mtimeMs
|
|
&& left.ctimeMs === right.ctimeMs;
|
|
}
|
|
|
|
function isWithin(root: string, candidate: string) {
|
|
const relative = path.relative(root, candidate);
|
|
return relative !== "" && relative !== ".." && !relative.startsWith(`..${path.sep}`) && !path.isAbsolute(relative);
|
|
}
|
|
|
|
function validMedia(kind: ReferenceKind, extension: string, header: Buffer) {
|
|
const ascii = header.toString("ascii");
|
|
if (kind === "image") {
|
|
if (![".png", ".jpg", ".jpeg", ".webp"].includes(extension)) return false;
|
|
return (extension === ".png" && header.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a])))
|
|
|| ([".jpg", ".jpeg"].includes(extension) && header[0] === 0xff && header[1] === 0xd8 && header[2] === 0xff)
|
|
|| (extension === ".webp" && ascii.startsWith("RIFF") && ascii.slice(8, 12) === "WEBP");
|
|
}
|
|
if (kind === "video") {
|
|
if (![".mp4", ".mov", ".webm"].includes(extension)) return false;
|
|
return ([".mp4", ".mov"].includes(extension) && ascii.slice(4, 8) === "ftyp")
|
|
|| (extension === ".webm" && header.subarray(0, 4).equals(Buffer.from([0x1a, 0x45, 0xdf, 0xa3])));
|
|
}
|
|
if (![".mp3", ".wav", ".m4a", ".aac", ".flac"].includes(extension)) return false;
|
|
return (extension === ".mp3" && (ascii.startsWith("ID3") || (header[0] === 0xff && (header[1] & 0xe0) === 0xe0)))
|
|
|| (extension === ".wav" && ascii.startsWith("RIFF") && ascii.slice(8, 12) === "WAVE")
|
|
|| (extension === ".m4a" && ascii.slice(4, 8) === "ftyp")
|
|
|| (extension === ".aac" && header[0] === 0xff && (header[1] & 0xf0) === 0xf0)
|
|
|| (extension === ".flac" && ascii.startsWith("fLaC"));
|
|
}
|
|
|
|
function referenceInvalid() { return new DreaminaCliError("dreamina_reference_invalid", "Dreamina reference media is not an allowed owner resource", 400); }
|
|
function referenceBudgetExceeded() { return new DreaminaCliError("dreamina_reference_budget_exceeded", "Dreamina reference staging exceeds the request byte limit", 413); }
|
|
function referenceCleanupFailed() { return new DreaminaCliError("dreamina_reference_cleanup_failed", "Dreamina reference staging cleanup failed", 503); }
|
|
function cancelled() { return new DreaminaCliError("dreamina_cancelled", "Dreamina 操作已取消", 499); }
|
|
function throwIfCancelled(signal?: AbortSignal) { if (signal?.aborted) throw cancelled(); }
|