142 lines
3.6 KiB
TypeScript
142 lines
3.6 KiB
TypeScript
import fs from "node:fs/promises";
|
|
import path from "node:path";
|
|
import { randomUUID } from "node:crypto";
|
|
import { appConfig } from "@/lib/config";
|
|
import type { GeneratedImage, StoredGeneratedImageInput } from "@/lib/types";
|
|
|
|
export interface PersistedImageFile {
|
|
absolutePath: string;
|
|
storageKey: string;
|
|
}
|
|
|
|
export async function persistGeneratedImages({
|
|
userId,
|
|
jobId,
|
|
images,
|
|
createdAt = new Date(),
|
|
}: {
|
|
userId: string;
|
|
jobId: string;
|
|
images: GeneratedImage[];
|
|
createdAt?: Date;
|
|
}): Promise<{
|
|
images: StoredGeneratedImageInput[];
|
|
files: PersistedImageFile[];
|
|
}> {
|
|
const year = String(createdAt.getUTCFullYear());
|
|
const month = String(createdAt.getUTCMonth() + 1).padStart(2, "0");
|
|
const writtenFiles: PersistedImageFile[] = [];
|
|
|
|
try {
|
|
const persistedImages: StoredGeneratedImageInput[] = [];
|
|
|
|
for (const image of images) {
|
|
const extension = resolveFileExtension(image);
|
|
const imageId = randomUUID();
|
|
const safeFileName = sanitizeFileName(image.fileName, extension);
|
|
const storageKey = path.posix.join(
|
|
userId,
|
|
year,
|
|
month,
|
|
jobId,
|
|
`${imageId}.${extension}`,
|
|
);
|
|
const absolutePath = resolveStoredImageAbsolutePath(storageKey);
|
|
const directory = path.dirname(absolutePath);
|
|
const buffer = Buffer.from(image.base64, "base64");
|
|
|
|
await fs.mkdir(directory, { recursive: true });
|
|
await fs.writeFile(absolutePath, buffer);
|
|
|
|
writtenFiles.push({
|
|
absolutePath,
|
|
storageKey,
|
|
});
|
|
persistedImages.push({
|
|
id: imageId,
|
|
fileName: safeFileName,
|
|
storageKey,
|
|
mimeType: image.mimeType,
|
|
fileSize: buffer.byteLength,
|
|
revisedPrompt: image.revisedPrompt,
|
|
});
|
|
}
|
|
|
|
return {
|
|
images: persistedImages,
|
|
files: writtenFiles,
|
|
};
|
|
} catch (error) {
|
|
await removePersistedImageFiles(writtenFiles);
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
export async function removePersistedImageFiles(
|
|
files: PersistedImageFile[],
|
|
): Promise<void> {
|
|
await Promise.all(
|
|
files.map(async (file) => {
|
|
try {
|
|
await fs.unlink(file.absolutePath);
|
|
} catch {
|
|
return;
|
|
}
|
|
}),
|
|
);
|
|
}
|
|
|
|
export async function readStoredImageFile(storageKey: string): Promise<Buffer> {
|
|
return fs.readFile(resolveStoredImageAbsolutePath(storageKey));
|
|
}
|
|
|
|
export function buildStoredImageUrl(imageId: string): string {
|
|
return `/api/images/${imageId}`;
|
|
}
|
|
|
|
export function buildStoredImageDownloadUrl(imageId: string): string {
|
|
return `/api/images/${imageId}?download=1`;
|
|
}
|
|
|
|
function resolveStoredImageAbsolutePath(storageKey: string): string {
|
|
const rootPath = path.resolve(appConfig.generatedImagesPath);
|
|
const absolutePath = path.resolve(rootPath, storageKey);
|
|
|
|
if (
|
|
absolutePath !== rootPath &&
|
|
!absolutePath.startsWith(`${rootPath}${path.sep}`)
|
|
) {
|
|
throw new Error("Invalid storage key.");
|
|
}
|
|
|
|
return absolutePath;
|
|
}
|
|
|
|
function resolveFileExtension(image: GeneratedImage): string {
|
|
const explicitExtension = path.extname(image.fileName).replace(/^\./, "");
|
|
|
|
if (explicitExtension) {
|
|
return explicitExtension.toLowerCase();
|
|
}
|
|
|
|
if (image.mimeType === "image/jpeg") {
|
|
return "jpg";
|
|
}
|
|
|
|
if (image.mimeType === "image/webp") {
|
|
return "webp";
|
|
}
|
|
|
|
return "png";
|
|
}
|
|
|
|
function sanitizeFileName(fileName: string, extension: string): string {
|
|
const baseName = path.basename(fileName, path.extname(fileName));
|
|
const safeBaseName = baseName
|
|
.replace(/[^a-zA-Z0-9-_]+/g, "-")
|
|
.replace(/^-+|-+$/g, "")
|
|
.slice(0, 64);
|
|
|
|
return `${safeBaseName || "generated-image"}.${extension}`;
|
|
}
|