koc-loop/lib/object-store.ts

import { mkdir, readFile, rename, stat, writeFile } from "node:fs/promises";
import path from "node:path";
import { getRuntimeEnv } from "./runtime-env";

type PutOptions = {
  httpMetadata?: { contentType?: string };
  customMetadata?: Record<string, string>;
};

type StoredMetadata = {
  contentType: string;
  customMetadata: Record<string, string>;
  uploadedAt: string;
};

export class StoredObjectBody {
  readonly body: Uint8Array;
  private readonly bytes: Buffer;
  private readonly metadata: StoredMetadata;

  constructor(bytes: Buffer, metadata: StoredMetadata) {
    this.bytes = bytes;
    this.metadata = metadata;
    this.body = new Uint8Array(bytes);
  }

  async arrayBuffer() {
    return this.body.buffer.slice(
      this.body.byteOffset,
      this.body.byteOffset + this.body.byteLength,
    ) as ArrayBuffer;
  }

  writeHttpMetadata(headers: Headers) {
    headers.set("Content-Type", this.metadata.contentType);
  }

  get customMetadata() {
    return this.metadata.customMetadata;
  }
}

export function normalizeObjectKey(key: string) {
  const normalized = key.replaceAll("\\", "/").replace(/^\/+/, "");
  if (!normalized || normalized.includes("\0")) {
    throw new Error("文件存储键无效");
  }
  const parts = normalized.split("/");
  if (parts.some((part) => !part || part === "." || part === "..")) {
    throw new Error("文件存储键不安全");
  }
  return parts.join("/");
}

function defaultUploadDir() {
  return path.join(process.cwd(), ".data", "uploads");
}

export class LocalObjectStore {
  readonly root: string;

  constructor(root = getRuntimeEnv().UPLOAD_DIR || defaultUploadDir()) {
    this.root = root;
  }

  private resolve(key: string) {
    const normalized = normalizeObjectKey(key);
    const filePath = path.join(this.root, ...normalized.split("/"));
    return { normalized, filePath, metadataPath: `${filePath}.metadata.json` };
  }

  async put(key: string, input: ArrayBuffer | Uint8Array, options: PutOptions = {}) {
    const target = this.resolve(key);
    await mkdir(path.dirname(target.filePath), { recursive: true });
    const bytes = input instanceof Uint8Array ? input : new Uint8Array(input);
    const metadata: StoredMetadata = {
      contentType: options.httpMetadata?.contentType || "application/octet-stream",
      customMetadata: options.customMetadata ?? {},
      uploadedAt: new Date().toISOString(),
    };
    const suffix = `${process.pid}-${crypto.randomUUID()}`;
    const temporaryFile = `${target.filePath}.${suffix}.tmp`;
    const temporaryMetadata = `${target.metadataPath}.${suffix}.tmp`;
    await Promise.all([
      writeFile(temporaryFile, bytes),
      writeFile(temporaryMetadata, JSON.stringify(metadata), "utf8"),
    ]);
    await rename(temporaryFile, target.filePath);
    await rename(temporaryMetadata, target.metadataPath);
    return { key: target.normalized };
  }

  async get(key: string) {
    const target = this.resolve(key);
    try {
      const [bytes, metadataText] = await Promise.all([
        readFile(target.filePath),
        readFile(target.metadataPath, "utf8").catch(() => ""),
      ]);
      const fallback: StoredMetadata = {
        contentType: "application/octet-stream",
        customMetadata: {},
        uploadedAt: (await stat(target.filePath)).mtime.toISOString(),
      };
      const metadata = metadataText
        ? ({ ...fallback, ...JSON.parse(metadataText) } as StoredMetadata)
        : fallback;
      return new StoredObjectBody(bytes, metadata);
    } catch (error) {
      if ((error as NodeJS.ErrnoException).code === "ENOENT") return null;
      throw error;
    }
  }
}

declare global {
  var __kocLoopObjectStore: LocalObjectStore | undefined;
}

export function getObjectStore() {
  if (!globalThis.__kocLoopObjectStore) {
    globalThis.__kocLoopObjectStore = new LocalObjectStore();
  }
  return globalThis.__kocLoopObjectStore;
}