arkvault/src/lib/client/transfer/upload.svelte.ts

import { limitFunction } from "p-limit";
import { signMessageHmac } from "#lib/client/crypto/hmac.js";
import { extractExifDateTime } from "#lib/client/media/exif.js";
import { generateThumbnail } from "#lib/client/media/thumbnail.js";
import { uploadBlob } from "#lib/client/transfer/chunkedUpload.js";
import {
  encodeToBase64,
  generateDataKey,
  wrapDataKey,
  encryptString,
} from "#lib/shared/crypto/index.js";
import { getChunkCount } from "#lib/shared/format/index.js";
import type { MasterKey, HmacSecret } from "#lib/stores/index.js";
import { Scheduler } from "#lib/utils/index.js";
import { trpc } from "#trpc/client.js";

export interface FileUploadState {
  id: string;
  name: string;
  parentId: DirectoryId;
  status:
    | "queued"
    | "encryption-pending"
    | "encrypting"
    | "upload-pending"
    | "uploading"
    | "uploaded"
    | "canceled"
    | "error";
  progress?: number;
  rate?: number;
  estimated?: number;
}

export type LiveFileUploadState = FileUploadState & {
  status: "queued" | "encryption-pending" | "encrypting" | "upload-pending" | "uploading";
};

const scheduler = new Scheduler<
  { fileId: number; fileBuffer?: ArrayBuffer; thumbnailBuffer?: ArrayBuffer } | undefined
>();
let uploadingFiles: FileUploadState[] = $state([]);

const isFileUploading = (status: FileUploadState["status"]) =>
  ["queued", "encryption-pending", "encrypting", "upload-pending", "uploading"].includes(status);

export const getUploadingFiles = (parentId?: DirectoryId) => {
  return uploadingFiles.filter(
    (file) =>
      (parentId === undefined || file.parentId === parentId) && isFileUploading(file.status),
  );
};

export const clearUploadedFiles = () => {
  uploadingFiles = uploadingFiles.filter((file) => isFileUploading(file.status));
};

const requestDuplicateFileScan = limitFunction(
  async (
    state: FileUploadState,
    file: File,
    hmacSecret: HmacSecret,
    onDuplicate: () => Promise<boolean>,
  ) => {
    state.status = "encryption-pending";

    const fileSigned = encodeToBase64(await signMessageHmac(file, hmacSecret.secret));
    const files = await trpc.file.listByHash.query({
      hskVersion: hmacSecret.version,
      contentHmac: fileSigned,
    });
    if (files.length === 0 || (await onDuplicate())) {
      return { fileSigned };
    } else {
      return {};
    }
  },
  { concurrency: 1 },
);

const getFileType = (file: File) => {
  if (file.type) return file.type;
  if (file.name.endsWith(".heic")) return "image/heic";
  throw new Error("Unknown file type");
};

interface FileMetadata {
  parentId: "root" | number;
  name: string;
  createdAt?: Date;
  lastModifiedAt: Date;
}

const requestFileMetadataEncryption = limitFunction(
  async (
    state: FileUploadState,
    file: Blob,
    fileMetadata: FileMetadata,
    masterKey: MasterKey,
    hmacSecret: HmacSecret,
  ) => {
    state.status = "encrypting";

    const { dataKey, dataKeyVersion } = await generateDataKey();
    const dataKeyWrapped = await wrapDataKey(dataKey, masterKey.key);

    const [nameEncrypted, createdAtEncrypted, lastModifiedAtEncrypted, thumbnailBuffer] =
      await Promise.all([
        encryptString(fileMetadata.name, dataKey),
        fileMetadata.createdAt &&
          encryptString(fileMetadata.createdAt.getTime().toString(), dataKey),
        encryptString(fileMetadata.lastModifiedAt.getTime().toString(), dataKey),
        generateThumbnail(file).then((blob) => blob?.arrayBuffer()),
      ]);

    const { uploadId } = await trpc.upload.startFileUpload.mutate({
      chunks: getChunkCount(file.size),
      parent: fileMetadata.parentId,
      mekVersion: masterKey.version,
      dek: dataKeyWrapped,
      dekVersion: dataKeyVersion,
      hskVersion: hmacSecret.version,
      contentType: file.type,
      name: nameEncrypted.ciphertext,
      nameIv: nameEncrypted.iv,
      createdAt: createdAtEncrypted?.ciphertext,
      createdAtIv: createdAtEncrypted?.iv,
      lastModifiedAt: lastModifiedAtEncrypted.ciphertext,
      lastModifiedAtIv: lastModifiedAtEncrypted.iv,
    });

    state.status = "upload-pending";
    return { uploadId, thumbnailBuffer, dataKey, dataKeyVersion };
  },
  { concurrency: 4 },
);

const requestFileUpload = limitFunction(
  async (
    state: FileUploadState,
    uploadId: string,
    file: Blob,
    fileSigned: string,
    thumbnailBuffer: ArrayBuffer | undefined,
    dataKey: CryptoKey,
    dataKeyVersion: Date,
  ) => {
    state.status = "uploading";

    const { encContentHash } = await uploadBlob(uploadId, file, dataKey, {
      onProgress(s) {
        state.progress = s.progress;
        state.rate = s.rate;
      },
    });

    const { file: fileId } = await trpc.upload.completeFileUpload.mutate({
      uploadId,
      contentHmac: fileSigned,
      encContentHash,
    });

    if (thumbnailBuffer) {
      try {
        const { uploadId } = await trpc.upload.startFileThumbnailUpload.mutate({
          file: fileId,
          dekVersion: dataKeyVersion,
        });

        await uploadBlob(uploadId, new Blob([thumbnailBuffer]), dataKey);

        await trpc.upload.completeFileThumbnailUpload.mutate({ uploadId });
      } catch (e) {
        console.error(e);
      }
    }

    state.status = "uploaded";
    return { fileId };
  },
  { concurrency: 1 },
);

export const uploadFile = async (
  file: File,
  parentId: "root" | number,
  masterKey: MasterKey,
  hmacSecret: HmacSecret,
  onDuplicate: () => Promise<boolean>,
) => {
  uploadingFiles.push({
    id: crypto.randomUUID(),
    name: file.name,
    parentId,
    status: "queued",
  });
  const state = uploadingFiles.at(-1)!;

  return await scheduler.schedule(file.size, async () => {
    try {
      const { fileSigned } = await requestDuplicateFileScan(state, file, hmacSecret, onDuplicate);

      if (!fileSigned) {
        state.status = "canceled";
        uploadingFiles = uploadingFiles.filter((file) => file !== state);
        return;
      }

      let fileBuffer;
      const fileType = getFileType(file);
      const fileMetadata: FileMetadata = {
        parentId,
        name: file.name,
        lastModifiedAt: new Date(file.lastModified),
      };

      if (fileType.startsWith("image/")) {
        fileBuffer = await file.arrayBuffer();
        fileMetadata.createdAt = extractExifDateTime(fileBuffer);
      }

      const blob = new Blob([file], { type: fileType });

      const { uploadId, thumbnailBuffer, dataKey, dataKeyVersion } =
        await requestFileMetadataEncryption(state, blob, fileMetadata, masterKey, hmacSecret);

      const { fileId } = await requestFileUpload(
        state,
        uploadId,
        blob,
        fileSigned,
        thumbnailBuffer,
        dataKey,
        dataKeyVersion,
      );

      return { fileId, fileBuffer, thumbnailBuffer };
    } catch (e) {
      state.status = "error";
      throw e;
    }
  });
};