src/lib/documents/server-storage.ts

Linked repository file from this build.

import { createHash } from "node:crypto";
import { mkdir, readdir, readFile, rm, stat, writeFile } from "node:fs/promises";
import path from "node:path";
import {
  DeleteObjectCommand,
  GetObjectCommand,
  HeadObjectCommand,
  ListObjectsV2Command,
  PutObjectCommand,
  S3Client,
} from "@aws-sdk/client-s3";
import { createPresignedPost } from "@aws-sdk/s3-presigned-post";
import { getSignedUrl } from "@aws-sdk/s3-request-presigner";
import { z } from "zod";
import { MAX_PDF_BYTES, validatePdf } from "./storage";

export class DocumentError extends Error {
  constructor(
    message: string,
    public status = 400,
  ) {
    super(message);
  }
}

export const documentKeySchema = z.uuid();
export const uploadRequestSchema = z.object({
  name: z
    .string()
    .max(255)
    .refine((name) => name.toLowerCase().endsWith(".pdf")),
  size: z.number().int().min(1).max(MAX_PDF_BYTES),
});

export function storageConfig(env: NodeJS.ProcessEnv = process.env) {
  const mode = env.DOCUMENT_STORAGE ?? (env.NODE_ENV === "production" ? "neon" : "local");
  if (mode === "local") return { mode: "local" as const };
  if (mode !== "neon") throw new DocumentError("Invalid document storage configuration.", 503);
  const endpoint = env.AWS_ENDPOINT_URL_S3;
  const region = env.AWS_REGION;
  const accessKeyId = env.AWS_ACCESS_KEY_ID;
  const secretAccessKey = env.AWS_SECRET_ACCESS_KEY;
  const bucket = env.DOCUMENT_STORAGE_BUCKET;
  if (!endpoint || !region || !accessKeyId || !secretAccessKey || !bucket)
    throw new DocumentError("Document storage is not configured. Contact your administrator.", 503);
  if (!endpoint.startsWith("https://"))
    throw new DocumentError("Document storage requires an HTTPS endpoint.", 503);
  return {
    mode: "neon" as const,
    endpoint,
    region,
    credentials: { accessKeyId, secretAccessKey },
    bucket,
  };
}

export function documentStorage() {
  const config = storageConfig();
  if (config.mode === "local") return { mode: "local" as const };
  const client = new S3Client({
    endpoint: config.endpoint,
    region: config.region,
    credentials: config.credentials,
    forcePathStyle: true,
    requestChecksumCalculation: "WHEN_REQUIRED",
    responseChecksumValidation: "WHEN_REQUIRED",
  });
  return { mode: "neon" as const, ...createNeonDocumentStorage(client, config.bucket) };
}

export function createNeonDocumentStorage(client: S3Client, bucket: string) {
  return {
    async prepare(input: z.infer<typeof uploadRequestSchema>, owner: string) {
      const { size } = uploadRequestSchema.parse(input);
      const uploadId = crypto.randomUUID();
      const upload = await createPresignedPost(client, {
        Bucket: bucket,
        Key: `pending/${uploadId}`,
        Expires: 300,
        Fields: { "Content-Type": "application/pdf", "x-amz-meta-owner": owner },
        Conditions: [
          ["content-length-range", size, size],
          ["eq", "$Content-Type", "application/pdf"],
          ["eq", "$x-amz-meta-owner", owner],
        ],
      });
      return { uploadId, ...upload };
    },
    async complete(uploadId: string, owner: string) {
      documentKeySchema.parse(uploadId);
      const pendingKey = `pending/${uploadId}`;
      // Read at most the allowed size plus one byte, even if the object was uploaded elsewhere.
      const object = await client.send(
        new GetObjectCommand({
          Bucket: bucket,
          Key: pendingKey,
          Range: `bytes=0-${MAX_PDF_BYTES}`,
        }),
      );
      if (object.Metadata?.owner !== owner)
        throw new DocumentError("This upload belongs to another user.", 403);
      if (!object.Body || !object.ContentLength || object.ContentLength > MAX_PDF_BYTES)
        throw new DocumentError("PDF must be 10 MB or smaller and cannot be empty.");
      const bytes = await object.Body.transformToByteArray();
      const file = new File([new Uint8Array(bytes)], "document.pdf", { type: object.ContentType });
      await validatePdf(file);
      const key = crypto.randomUUID();
      // The browser can only write pending objects. Final documents cannot be overwritten by an upload URL.
      await client.send(
        new PutObjectCommand({
          Bucket: bucket,
          Key: key,
          Body: bytes,
          ContentType: "application/pdf",
          Metadata: { owner },
        }),
      );
      await client.send(new DeleteObjectCommand({ Bucket: bucket, Key: pendingKey })).catch(() => {
        console.error("Unable to remove a completed pending document upload.");
      });
      return key;
    },
    /** Size and uploader of a final document, or undefined when it does not exist. */
    async describe(key: string): Promise<StoredDocument | undefined> {
      documentKeySchema.parse(key);
      const head = await client
        .send(new HeadObjectCommand({ Bucket: bucket, Key: key }))
        .catch((error: { name?: string }) => {
          if (error.name === "NotFound") return undefined;
          throw error;
        });
      return head && { size: head.ContentLength ?? 0, owner: head.Metadata?.owner };
    },
    /** Every object, including unfinished uploads under pending/. */
    async list(): Promise<StoredObject[]> {
      const objects: StoredObject[] = [];
      let token: string | undefined;
      do {
        const page = await client.send(
          new ListObjectsV2Command({ Bucket: bucket, ContinuationToken: token }),
        );
        for (const item of page.Contents ?? [])
          if (item.Key)
            objects.push({ key: item.Key, modifiedAt: item.LastModified ?? new Date() });
        token = page.NextContinuationToken;
      } while (token);
      return objects;
    },
    async remove(key: string) {
      await client.send(new DeleteObjectCommand({ Bucket: bucket, Key: key }));
    },
    async url(key: string) {
      documentKeySchema.parse(key);
      return getSignedUrl(
        client,
        new GetObjectCommand({
          Bucket: bucket,
          Key: key,
          ResponseContentType: "application/pdf",
          ResponseCacheControl: "private, no-store",
        }),
        { expiresIn: 60 },
      );
    },
  };
}

export function uploadOwner(program: string, staffId: string) {
  return createHash("sha256")
    .update(JSON.stringify([program, staffId]))
    .digest("hex");
}

export interface StoredObject {
  key: string;
  modifiedAt: Date;
}

export interface StoredDocument {
  size: number;
  owner?: string;
}

const localDirectory = () => path.join(process.cwd(), ".demo-documents");

export async function putLocalDocument(file: File, owner: string) {
  await validatePdf(file);
  const key = crypto.randomUUID();
  await mkdir(localDirectory(), { recursive: true });
  await writeFile(path.join(localDirectory(), key), Buffer.from(await file.arrayBuffer()), {
    flag: "wx",
  });
  await writeFile(path.join(localDirectory(), `${key}.owner`), owner, { flag: "wx" });
  return key;
}

export function readLocalDocument(key: string) {
  documentKeySchema.parse(key);
  return readFile(path.join(localDirectory(), key));
}

async function describeLocalDocument(key: string): Promise<StoredDocument | undefined> {
  documentKeySchema.parse(key);
  const file = await stat(path.join(localDirectory(), key)).catch(() => undefined);
  if (!file) return undefined;
  const owner = await readFile(path.join(localDirectory(), `${key}.owner`), "utf8").catch(
    () => undefined,
  );
  return { size: file.size, owner };
}

/**
 * Check a key before a review attaches it: the file must exist and come from this uploader.
 * Returns the stored size, so the browser cannot set it.
 */
export async function attachableDocumentSize(key: string, owner: string): Promise<number> {
  const storage = documentStorage();
  const found =
    storage.mode === "neon" ? await storage.describe(key) : await describeLocalDocument(key);
  if (!found || found.owner !== owner)
    throw new DocumentError("Upload the PDF again before attaching it.");
  return found.size;
}

/** Local files, listed like bucket objects. Owner sidecars go with their document. */
export async function listLocalDocuments(): Promise<StoredObject[]> {
  const names = await readdir(localDirectory()).catch(() => []);
  return Promise.all(
    names
      .filter((name) => !name.endsWith(".owner"))
      .map(async (key) => ({
        key,
        modifiedAt: (await stat(path.join(localDirectory(), key))).mtime,
      })),
  );
}

export async function removeLocalDocument(key: string) {
  await rm(path.join(localDirectory(), key), { force: true });
  await rm(path.join(localDirectory(), `${key}.owner`), { force: true });
}

/**
 * Objects no review snapshot references. Recent objects are kept, because an upload
 * finishes before the review save that attaches it.
 */
export function orphanedObjects(
  objects: StoredObject[],
  referenced: ReadonlySet<string>,
  now: Date,
  minimumAgeMs = 24 * 60 * 60 * 1000,
) {
  return objects.filter(
    (object) =>
      !referenced.has(object.key) && now.getTime() - object.modifiedAt.getTime() >= minimumAgeMs,
  );
}