quickdrop/app/utils/crypto.ts

import {
  CHUNK_SIZE,
  GCM_TAG_BYTES,
  MAX_MANIFEST_PLAINTEXT,
  PROTOCOL_VERSION,
} from "~/../shared/types/drop";
import {
  decodeBase64Url,
  validateRoomId,
} from "~/../shared/utils/drop-contract";
import type { EncryptedManifest } from "~/../shared/types/drop";

const textEncoder = new TextEncoder();
const textDecoder = new TextDecoder();
const MANIFEST_COUNTER = 0xffffffff;

export function safeFilename(value: string): string {
  const cleaned = value
    .normalize("NFC")
    .replace(/[\\/\u0000-\u001F\u007F\u202A-\u202E\u2066-\u2069]/gu, "_")
    .trim();
  return [...cleaned].slice(0, 240).join("") || "download";
}
export function toBase64Url(bytes: Uint8Array): string {
  let binary = "";
  for (const byte of bytes) binary += String.fromCharCode(byte);
  return btoa(binary)
    .replaceAll("+", "-")
    .replaceAll("/", "_")
    .replace(/=+$/u, "");
}
export function fromBase64Url(value: string): Uint8Array {
  return decodeBase64Url(value);
}
export function keyFromHash(hash: string): string | null {
  return new URLSearchParams(hash.replace(/^#/u, "")).get("key");
}
export function accessFromHash(hash: string): string | null {
  return new URLSearchParams(hash.replace(/^#/u, "")).get("access");
}
export function generateKey(): string {
  return toBase64Url(crypto.getRandomValues(new Uint8Array(32)));
}
export function generateNoncePrefix(): string {
  return toBase64Url(crypto.getRandomValues(new Uint8Array(8)));
}

async function importKey(encoded: string, usage: KeyUsage): Promise<CryptoKey> {
  const raw = fromBase64Url(encoded);
  if (raw.byteLength !== 32)
    throw new Error("The private link contains an invalid key.");
  return crypto.subtle.importKey(
    "raw",
    asBuffer(raw),
    { name: "AES-GCM" },
    false,
    [usage],
  );
}
function iv(prefix: string, counter: number): Uint8Array {
  const bytes = fromBase64Url(prefix);
  if (bytes.byteLength !== 8) throw new Error("Invalid nonce prefix");
  const result = new Uint8Array(12);
  result.set(bytes);
  new DataView(result.buffer).setUint32(8, counter >>> 0, false);
  return result;
}
function aad(
  id: string,
  kind: 0 | 1,
  index: number,
  chunkSize: number,
  chunkCount: number,
  plainLength: number,
): Uint8Array {
  validateRoomId(id);
  const idBytes = textEncoder.encode(id);
  if (idBytes.byteLength !== 32) throw new Error("Invalid drop id");
  const result = new Uint8Array(4 + 32 + 1 + 16);
  result.set([0x51, 0x44, 0x32, 0], 0);
  result.set(idBytes, 4);
  result[36] = kind;
  const view = new DataView(result.buffer);
  view.setUint32(37, index >>> 0, false);
  view.setUint32(41, chunkSize >>> 0, false);
  view.setUint32(45, chunkCount >>> 0, false);
  view.setUint32(49, plainLength >>> 0, false);
  return result;
}
function asBuffer(data: Uint8Array): ArrayBuffer {
  return data.buffer.slice(
    data.byteOffset,
    data.byteOffset + data.byteLength,
  ) as ArrayBuffer;
}
async function seal(
  data: Uint8Array,
  encodedKey: string,
  id: string,
  prefix: string,
  kind: 0 | 1,
  index: number,
  chunkSize: number,
  count: number,
): Promise<Uint8Array> {
  const key = await importKey(encodedKey, "encrypt");
  const encrypted = await crypto.subtle.encrypt(
    {
      name: "AES-GCM",
      iv: asBuffer(iv(prefix, index)),
      additionalData: asBuffer(
        aad(id, kind, index, chunkSize, count, data.byteLength),
      ),
      tagLength: 128,
    },
    key,
    asBuffer(data),
  );
  return new Uint8Array(encrypted);
}
async function open(
  data: Uint8Array,
  encodedKey: string,
  id: string,
  prefix: string,
  kind: 0 | 1,
  index: number,
  chunkSize: number,
  count: number,
  plainLength: number,
): Promise<Uint8Array> {
  if (data.byteLength < GCM_TAG_BYTES)
    throw new Error("Encrypted data is truncated.");
  const key = await importKey(encodedKey, "decrypt");
  let decrypted: ArrayBuffer;
  try {
    decrypted = await crypto.subtle.decrypt(
      {
        name: "AES-GCM",
        iv: asBuffer(iv(prefix, index)),
        additionalData: asBuffer(
          aad(id, kind, index, chunkSize, count, plainLength),
        ),
        tagLength: 128,
      },
      key,
      asBuffer(data),
    );
  } catch {
    throw new Error("Encrypted data failed integrity verification.");
  }
  const result = new Uint8Array(decrypted);
  if (result.byteLength !== plainLength)
    throw new Error("Decrypted length does not match the manifest.");
  return result;
}
export async function encryptChunk(
  plain: Uint8Array | ArrayBuffer,
  key: string,
  id: string,
  noncePrefix: string,
  index: number,
  chunkSize = CHUNK_SIZE,
  chunkCount = 1,
): Promise<Uint8Array> {
  const data = plain instanceof Uint8Array ? plain : new Uint8Array(plain);
  return seal(data, key, id, noncePrefix, 1, index, chunkSize, chunkCount);
}
export async function decryptChunk(
  cipher: Uint8Array | ArrayBuffer,
  key: string,
  id: string,
  noncePrefix: string,
  index: number,
  plainLength: number,
  chunkSize = CHUNK_SIZE,
  chunkCount = 1,
): Promise<Uint8Array> {
  const data = cipher instanceof Uint8Array ? cipher : new Uint8Array(cipher);
  return open(
    data,
    key,
    id,
    noncePrefix,
    1,
    index,
    chunkSize,
    chunkCount,
    plainLength,
  );
}
export async function encryptManifest(
  manifest: EncryptedManifest,
  key: string,
  id: string,
  noncePrefix: string,
): Promise<Uint8Array> {
  const bytes = textEncoder.encode(JSON.stringify(manifest));
  if (bytes.byteLength > MAX_MANIFEST_PLAINTEXT)
    throw new Error("Manifest is too large.");
  return seal(
    bytes,
    key,
    id,
    noncePrefix,
    0,
    MANIFEST_COUNTER,
    manifest.chunkSize,
    manifest.chunkCount,
  );
}
export async function decryptManifest(
  cipher: Uint8Array | ArrayBuffer,
  key: string,
  id: string,
  noncePrefix: string,
  chunkSize: number,
  chunkCount: number,
): Promise<EncryptedManifest> {
  const data = cipher instanceof Uint8Array ? cipher : new Uint8Array(cipher);
  if (data.byteLength <= GCM_TAG_BYTES)
    throw new Error("The encrypted manifest is empty.");
  const bytes = await open(
    data,
    key,
    id,
    noncePrefix,
    0,
    MANIFEST_COUNTER,
    chunkSize,
    chunkCount,
    data.byteLength - GCM_TAG_BYTES,
  );
  let parsed: unknown;
  try {
    parsed = JSON.parse(textDecoder.decode(bytes));
  } catch {
    throw new Error("The encrypted manifest is invalid.");
  }
  if (!parsed || typeof parsed !== "object")
    throw new Error("The encrypted manifest is invalid.");
  const value = parsed as Record<string, unknown>;
  const fields = [
    "version",
    "name",
    "type",
    "size",
    "lastModified",
    "chunkSize",
    "chunkCount",
  ];
  if (
    Object.keys(value).some((field) => !fields.includes(field)) ||
    Object.keys(value).length !== fields.length ||
    value.version !== PROTOCOL_VERSION ||
    typeof value.name !== "string" ||
    typeof value.type !== "string" ||
    !Number.isSafeInteger(value.size) ||
    (value.size as number) < 0 ||
    !Number.isSafeInteger(value.lastModified) ||
    (value.lastModified as number) < 0 ||
    value.chunkSize !== chunkSize ||
    value.chunkCount !== chunkCount
  )
    throw new Error("The encrypted manifest is invalid.");
  if (value.name.length > 1024 || value.type.length > 255)
    throw new Error("The encrypted manifest is invalid.");
  return {
    version: 2,
    name: safeFilename(value.name),
    type: value.type || "application/octet-stream",
    size: value.size as number,
    lastModified: value.lastModified as number,
    chunkSize,
    chunkCount,
  };
}
export interface DropMetadata {
  name: string;
  type: string;
  size: number;
  lastModified: number;
}
export interface DecryptedDrop {
  blob: Blob;
  metadata: DropMetadata;
}
export function formatBytes(bytes: number): string {
  if (bytes === 0) return "0 B";
  const units = ["B", "KB", "MB", "GB"];
  const unit = Math.min(
    Math.floor(Math.log(bytes) / Math.log(1024)),
    units.length - 1,
  );
  const value = bytes / 1024 ** unit;
  return `${value >= 10 || unit === 0 ? value.toFixed(0) : value.toFixed(1)} ${units[unit]}`;
}