malice

import { statusForKind } from './lib/wire-status';

import { decodeAttachmentIv, importAttachmentKey, RangeDecryptor } from './lib/media-decrypt';

import { attachmentContentType } from './lib/media-sniff';

import { createAvatarLoader } from './lib/avatar-cache';

export const AVATAR_CACHE = 'malice-avatars-v3';

interface MediaReply {

  bytes?: Uint8Array<ArrayBuffer>;

  content_type?: string | null;

  error?: { kind: string; message: string };

}

interface AttachmentEncryptionWire {

  key: string;

  iv: string;

  sha256: string;

}

interface MediaFileWire {

  url: string;

  mimetype: string | null;

  declared_size: number | null;

  encryption: AttachmentEncryptionWire | null;

}

type MediaLocationWire = { File: MediaFileWire } | { ServerThumbnail: { url: string; fallback: MediaFileWire } } | 'Local';

interface LocateReply {

  location?: MediaLocationWire;

  token?: string | null;

  error?: { kind: string; message: string };

}

const MEDIA_TIMEOUT_MS = 30000;

const pendingMedia = new Map<string, (reply: MediaReply & LocateReply) => void>();

let nextMediaRequestId = 0;

let mediaChannel: BroadcastChannel | undefined;

function channel(): BroadcastChannel {

  if (!mediaChannel) {

    mediaChannel = new BroadcastChannel('malice-media');

    mediaChannel.onmessage = (ev: MessageEvent) => {

      const msg = ev.data as MediaReply & LocateReply & { id: string };

      const resolve = pendingMedia.get(msg.id);

      if (!resolve) return;

      pendingMedia.delete(msg.id);

      resolve(msg);

    };

  }

  return mediaChannel;

}

function requestMedia(kind: 'avatar' | 'media', fields: Record<string, string>): Promise<MediaReply | null> {

  const id = `${Date.now()}-${nextMediaRequestId++}`;

  return new Promise((resolve) => {

    const timer = setTimeout(() => {

      pendingMedia.delete(id);

      resolve(null);

    }, MEDIA_TIMEOUT_MS);

    pendingMedia.set(id, (reply) => {

      clearTimeout(timer);

      resolve(reply);

    });

    channel().postMessage({ id, kind, ...fields });

  });

}

interface LocateEntry {

  location: MediaLocationWire;

  token: string | null;

  expiresAt: number;

}

const LOCATE_TTL_MS = 5 * 60 * 1000;

const locateCache = new Map<string, LocateEntry>();

const sniffCache = new Map<string, { contentType: string | null; expiresAt: number }>();

function requestLocate(mediaId: string, size: string): Promise<{ location: MediaLocationWire; token: string | null } | null> {

  const id = `${Date.now()}-${nextMediaRequestId++}`;

  return new Promise((resolve) => {

    const timer = setTimeout(() => {

      pendingMedia.delete(id);

      resolve(null);

    }, MEDIA_TIMEOUT_MS);

    pendingMedia.set(id, (reply) => {

      clearTimeout(timer);

      resolve(reply.error || !reply.location ? null : { location: reply.location, token: reply.token ?? null });

    });

    channel().postMessage({ id, kind: 'locate', media_id: mediaId, size });

  });

}

/** In-memory only: the access token never touches Cache Storage or IndexedDB, and expires with the entry. */

async function locate(mediaId: string, size: string, force: boolean): Promise<LocateEntry | null> {

  if (!force) {

    const cached = locateCache.get(mediaId);

    if (cached && cached.expiresAt > Date.now()) return cached;

  }

  const reply = await requestLocate(mediaId, size);

  if (!reply) return null;

  const entry: LocateEntry = { ...reply, expiresAt: Date.now() + LOCATE_TTL_MS };

  locateCache.set(mediaId, entry);

  return entry;

}

interface RangeSlice {

  body: Uint8Array<ArrayBuffer>;

  status: number;

  contentRange?: string;

}

/** The engine hands back the whole attachment at once, so seeking in <video>/<audio> needs this SW to slice it. */

function sliceForRange(bytes: Uint8Array<ArrayBuffer>, rangeHeader: string | null): RangeSlice {

  const match = rangeHeader ? /^bytes=(\d*)-(\d*)$/.exec(rangeHeader) : null;

  if (!match) return { body: bytes, status: 200 };

  const total = bytes.byteLength;

  let start = match[1] ? parseInt(match[1], 10) : undefined;

  let end = match[2] ? parseInt(match[2], 10) : undefined;

  if (start === undefined && end !== undefined) {

    start = Math.max(total - end, 0);

    end = total - 1;

  } else if (end === undefined) {

    end = total - 1;

  }

  if (start === undefined || start > end || start >= total) return { body: bytes, status: 200 };

  end = Math.min(end, total - 1);

  return { body: bytes.subarray(start, end + 1), status: 206, contentRange: `bytes ${start}-${end}/${total}` };

}

function mediaResponse(reply: MediaReply | null, rangeHeader: string | null, attachment: boolean): Response {

  if (!reply) return new Response(null, { status: 504 });

  if (reply.error) return new Response(reply.error.message, { status: statusForKind(reply.error.kind) });

  const bytes = reply.bytes ?? new Uint8Array();

  const { body, status, contentRange } = sliceForRange(bytes, rangeHeader);

  const headers = new Headers({

    'Content-Type': reply.content_type ?? 'application/octet-stream',

    'Cache-Control': attachment ? 'private, max-age=3600' : 'private, max-age=31536000, immutable',

    'X-Content-Type-Options': 'nosniff'

  });

  if (attachment) {

    headers.set('Content-Disposition', reply.content_type ? 'inline' : 'attachment');

    headers.set('Content-Security-Policy', 'sandbox');

  }

  if (contentRange) {

    headers.set('Content-Range', contentRange);

    headers.set('Accept-Ranges', 'bytes');

  }

  return new Response(body, { status, headers });

}

interface ByteRange {

  start: number;

  end: number;

}

/** Unlike sliceForRange, the total comes from the server's declared_size, not a buffer already in hand. */

function parseByteRange(rangeHeader: string | null, total: number): ByteRange | null {

  const match = rangeHeader ? /^bytes=(\d*)-(\d*)$/.exec(rangeHeader) : null;

  if (!match) return null;

  let start = match[1] ? parseInt(match[1], 10) : undefined;

  let end = match[2] ? parseInt(match[2], 10) : undefined;

  if (start === undefined && end !== undefined) {

    start = Math.max(total - end, 0);

    end = total - 1;

  } else if (end === undefined) {

    start = start ?? 0;

    end = total - 1;

  }

  if (start === undefined || start > end || start >= total) return null;

  return { start, end: Math.min(end, total - 1) };

}

function attachmentHeaders(contentType: string | null): Headers {

  return new Headers({

    'Content-Type': contentType ?? 'application/octet-stream',

    'Cache-Control': 'private, max-age=3600',

    'X-Content-Type-Options': 'nosniff',

    'Content-Disposition': contentType ? 'inline' : 'attachment',

    'Content-Security-Policy': 'sandbox'

  });

}

/**

 * The Authorization header carries the account's access token, so it's only worth the risk on the

 * exact authenticated media endpoints, and only when the URL being fetched is the one `locate` named

 * (never a redirect target, should a future caller pass one through).

 */

function shouldAuthorize(fetchUrl: string, locatedUrl: string): boolean {

  try {

    const target = new URL(fetchUrl);

    const located = new URL(locatedUrl);

    if (target.origin !== located.origin) return false;

    return target.pathname.startsWith('/_matrix/client/v1/media/') || target.pathname.startsWith('/_matrix/media/');

  } catch {

    return false;

  }

}

function concatChunks(parts: Uint8Array[]): Uint8Array<ArrayBuffer> {

  const total = parts.reduce((sum, part) => sum + part.byteLength, 0);

  const out = new Uint8Array(total);

  let offset = 0;

  for (const part of parts) {

    out.set(part, offset);

    offset += part.byteLength;

  }

  return out;

}

async function readUpTo(

  reader: ReadableStreamDefaultReader<Uint8Array<ArrayBuffer>>,

  limit: number

): Promise<Uint8Array<ArrayBuffer>> {

  const chunks: Uint8Array[] = [];

  let total = 0;

  while (total < limit) {

    const { value, done } = await reader.read();

    if (done) break;

    chunks.push(value);

    total += value.byteLength;

  }

  reader.cancel().catch(() => undefined);

  return concatChunks(chunks).subarray(0, Math.min(limit, total));

}

/**

 * Ciphertext in, plaintext out: skips the bytes the upstream sent but the caller didn't ask for

 * (only nonzero when the upstream ignored the Range request), decrypts, drops the sub-16-byte

 * alignment prefix from the very first output, and stops once `limit` bytes have been produced.

 */

async function* decryptedChunks(

  reader: ReadableStreamDefaultReader<Uint8Array<ArrayBuffer>>,

  decryptor: RangeDecryptor | null,

  skipBytes: number,

  dropPrefix: number,

  limit: number

): AsyncGenerator<Uint8Array> {

  let toSkip = skipBytes;

  let toDrop = dropPrefix;

  let remaining = limit;

  while (remaining > 0) {

    const { value, done } = await reader.read();

    if (done) break;

    let chunk = value;

    if (toSkip > 0) {

      if (chunk.byteLength <= toSkip) {

        toSkip -= chunk.byteLength;

        continue;

      }

      chunk = chunk.subarray(toSkip);

      toSkip = 0;

    }

    const plain = decryptor ? await decryptor.push(chunk) : chunk;

    if (plain.byteLength === 0) continue;

    const emitted = emit(plain);

    if (emitted) yield emitted;

  }

  if (remaining > 0 && decryptor) {

    const emitted = emit(await decryptor.finish());

    if (emitted) yield emitted;

  }

  reader.cancel().catch(() => undefined);

  function emit(chunk: Uint8Array): Uint8Array | null {

    let out = chunk;

    if (toDrop > 0) {

      if (out.byteLength <= toDrop) {

        toDrop -= out.byteLength;

        return null;

      }

      out = out.subarray(toDrop);

      toDrop = 0;

    }

    if (out.byteLength > remaining) out = out.subarray(0, remaining);

    remaining -= out.byteLength;

    return out.byteLength > 0 ? out : null;

  }

}

async function headSniff(mediaId: string, file: MediaFileWire, token: string | null, retried: boolean): Promise<string | null> {

  const cached = sniffCache.get(mediaId);

  if (cached && cached.expiresAt > Date.now()) return cached.contentType;

  const headers = new Headers({ Range: 'bytes=0-79' });

  if (shouldAuthorize(file.url, file.url) && token) headers.set('Authorization', `Bearer ${token}`);

  let upstream: Response;

  try {

    upstream = await fetch(file.url, { headers });

  } catch {

    return null;

  }

  if (upstream.status === 401 && !retried) {

    const fresh = await locate(mediaId, 'Full', true);

    if (fresh && typeof fresh.location === 'object' && 'File' in fresh.location) {

      return headSniff(mediaId, fresh.location.File, fresh.token, true);

    }

  }

  if (!upstream.body || (!upstream.ok && upstream.status !== 206)) return null;

  const head = await readUpTo(upstream.body.getReader(), 80);

  let contentType: string | null;

  if (file.encryption) {

    const key = await importAttachmentKey(file.encryption.key);

    const iv = decodeAttachmentIv(file.encryption.iv);

    const decryptor = new RangeDecryptor(key, iv, 0);

    contentType = attachmentContentType(concatChunks([await decryptor.push(head), await decryptor.finish()]).subarray(0, 64));

  } else {

    contentType = attachmentContentType(head.subarray(0, 64));

  }

  sniffCache.set(mediaId, { contentType, expiresAt: Date.now() + LOCATE_TTL_MS });

  return contentType;

}

function wholeFileFallback(mediaId: string, rangeHeader: string | null): Promise<Response> {

  return requestMedia('media', { media_id: mediaId, size: 'Full' }).then((reply) => mediaResponse(reply, rangeHeader, true));

}

async function streamFullMedia(mediaId: string, rangeHeader: string | null, entry?: LocateEntry, retried = false): Promise<Response> {

  const located = entry ?? (await locate(mediaId, 'Full', false));

  if (!located) return wholeFileFallback(mediaId, rangeHeader);

  if (typeof located.location !== 'object' || !('File' in located.location)) return wholeFileFallback(mediaId, rangeHeader);

  const file = located.location.File;

  const total = file.declared_size;

  if (total == null) return wholeFileFallback(mediaId, rangeHeader);

  const requested = parseByteRange(rangeHeader, total);

  const start = requested?.start ?? 0;

  const end = requested?.end ?? total - 1;

  const alignedStart = start - (start % 16);

  const headers = new Headers({ Range: `bytes=${alignedStart}-${end}` });

  if (shouldAuthorize(file.url, file.url) && located.token) headers.set('Authorization', `Bearer ${located.token}`);

  let upstream: Response;

  try {

    upstream = await fetch(file.url, { headers });

  } catch {

    return new Response('media upstream unreachable', { status: 502 });

  }

  if (upstream.status === 401 && !retried) {

    const fresh = await locate(mediaId, 'Full', true);

    if (fresh) return streamFullMedia(mediaId, rangeHeader, fresh, true);

  }

  if (!upstream.body || (!upstream.ok && upstream.status !== 206)) {

    return new Response('media upstream error', { status: 502 });

  }

  const decryptor = file.encryption

    ? new RangeDecryptor(await importAttachmentKey(file.encryption.key), decodeAttachmentIv(file.encryption.iv), alignedStart / 16)

    : null;

  const skipBytes = upstream.status === 206 ? 0 : alignedStart;

  const dropPrefix = start - alignedStart;

  const limit = end - start + 1;

  const reader = upstream.body.getReader();

  const chunks = decryptedChunks(reader, decryptor, skipBytes, dropPrefix, limit);

  const prefetched: Uint8Array[] = [];

  let prefetchedLen = 0;

  let chunksDone = false;

  let contentType: string | null;

  if (alignedStart === 0) {

    while (prefetchedLen < 64) {

      const step = await chunks.next();

      if (step.done) {

        chunksDone = true;

        break;

      }

      prefetched.push(step.value);

      prefetchedLen += step.value.byteLength;

    }

    contentType = attachmentContentType(concatChunks(prefetched).subarray(0, Math.min(64, prefetchedLen)));

    sniffCache.set(mediaId, { contentType, expiresAt: Date.now() + LOCATE_TTL_MS });

  } else {

    contentType = await headSniff(mediaId, file, located.token, retried);

  }

  const body = new ReadableStream<Uint8Array>({

    async pull(controller) {

      const next = prefetched.shift();

      if (next) {

        controller.enqueue(next);

        return;

      }

      if (chunksDone) {

        controller.close();

        return;

      }

      const step = await chunks.next();

      if (step.done) controller.close();

      else controller.enqueue(step.value);

    },

    cancel() {

      reader.cancel().catch(() => undefined);

    }

  });

  const responseHeaders = attachmentHeaders(contentType);

  responseHeaders.set('Accept-Ranges', 'bytes');

  if (requested) {

    responseHeaders.set('Content-Range', `bytes ${start}-${end}/${total}`);

    return new Response(body, { status: 206, headers: responseHeaders });

  }

  return new Response(body, { status: 200, headers: responseHeaders });

}

/** Keyed by the request URL, which names image and size, so cached bytes outlive the access token and survive restarts. */

let avatarLoader: ReturnType<typeof createAvatarLoader> | undefined;

function cachedAvatar(url: string, query: string): Promise<Response> {

  avatarLoader ??= createAvatarLoader(

    async () => caches.open(AVATAR_CACHE),

    async (q) => mediaResponse(await requestMedia('avatar', Object.fromEntries(new URLSearchParams(q))), null, false)

  );

  return avatarLoader(url, query);

}