| 1 | import { statusForKind } from './lib/wire-status'; |
| 2 | import { decodeAttachmentIv, importAttachmentKey, RangeDecryptor } from './lib/media-decrypt'; |
| 3 | import { attachmentContentType } from './lib/media-sniff'; |
| 4 | import { createAvatarLoader } from './lib/avatar-cache'; |
| 5 | |
| 6 | export const AVATAR_CACHE = 'malice-avatars-v3'; |
| 7 | |
| 8 | interface MediaReply { |
| 9 | bytes?: Uint8Array<ArrayBuffer>; |
| 10 | content_type?: string | null; |
| 11 | error?: { kind: string; message: string }; |
| 12 | } |
| 13 | |
| 14 | interface AttachmentEncryptionWire { |
| 15 | key: string; |
| 16 | iv: string; |
| 17 | sha256: string; |
| 18 | } |
| 19 | |
| 20 | interface MediaFileWire { |
| 21 | url: string; |
| 22 | mimetype: string | null; |
| 23 | declared_size: number | null; |
| 24 | encryption: AttachmentEncryptionWire | null; |
| 25 | } |
| 26 | |
| 27 | type MediaLocationWire = { File: MediaFileWire } | { ServerThumbnail: { url: string; fallback: MediaFileWire } } | 'Local'; |
| 28 | |
| 29 | interface LocateReply { |
| 30 | location?: MediaLocationWire; |
| 31 | token?: string | null; |
| 32 | error?: { kind: string; message: string }; |
| 33 | } |
| 34 | |
| 35 | const MEDIA_TIMEOUT_MS = 30000; |
| 36 | const pendingMedia = new Map<string, (reply: MediaReply & LocateReply) => void>(); |
| 37 | let nextMediaRequestId = 0; |
| 38 | let mediaChannel: BroadcastChannel | undefined; |
| 39 | |
| 40 | function channel(): BroadcastChannel { |
| 41 | if (!mediaChannel) { |
| 42 | mediaChannel = new BroadcastChannel('malice-media'); |
| 43 | mediaChannel.onmessage = (ev: MessageEvent) => { |
| 44 | const msg = ev.data as MediaReply & LocateReply & { id: string }; |
| 45 | const resolve = pendingMedia.get(msg.id); |
| 46 | if (!resolve) return; |
| 47 | pendingMedia.delete(msg.id); |
| 48 | resolve(msg); |
| 49 | }; |
| 50 | } |
| 51 | return mediaChannel; |
| 52 | } |
| 53 | |
| 54 | function requestMedia(kind: 'avatar' | 'media', fields: Record<string, string>): Promise<MediaReply | null> { |
| 55 | const id = `${Date.now()}-${nextMediaRequestId++}`; |
| 56 | return new Promise((resolve) => { |
| 57 | const timer = setTimeout(() => { |
| 58 | pendingMedia.delete(id); |
| 59 | resolve(null); |
| 60 | }, MEDIA_TIMEOUT_MS); |
| 61 | pendingMedia.set(id, (reply) => { |
| 62 | clearTimeout(timer); |
| 63 | resolve(reply); |
| 64 | }); |
| 65 | channel().postMessage({ id, kind, ...fields }); |
| 66 | }); |
| 67 | } |
| 68 | |
| 69 | interface LocateEntry { |
| 70 | location: MediaLocationWire; |
| 71 | token: string | null; |
| 72 | expiresAt: number; |
| 73 | } |
| 74 | |
| 75 | const LOCATE_TTL_MS = 5 * 60 * 1000; |
| 76 | const locateCache = new Map<string, LocateEntry>(); |
| 77 | const sniffCache = new Map<string, { contentType: string | null; expiresAt: number }>(); |
| 78 | |
| 79 | function requestLocate(mediaId: string, size: string): Promise<{ location: MediaLocationWire; token: string | null } | null> { |
| 80 | const id = `${Date.now()}-${nextMediaRequestId++}`; |
| 81 | return new Promise((resolve) => { |
| 82 | const timer = setTimeout(() => { |
| 83 | pendingMedia.delete(id); |
| 84 | resolve(null); |
| 85 | }, MEDIA_TIMEOUT_MS); |
| 86 | pendingMedia.set(id, (reply) => { |
| 87 | clearTimeout(timer); |
| 88 | resolve(reply.error || !reply.location ? null : { location: reply.location, token: reply.token ?? null }); |
| 89 | }); |
| 90 | channel().postMessage({ id, kind: 'locate', media_id: mediaId, size }); |
| 91 | }); |
| 92 | } |
| 93 | |
| 94 | /** In-memory only: the access token never touches Cache Storage or IndexedDB, and expires with the entry. */ |
| 95 | async function locate(mediaId: string, size: string, force: boolean): Promise<LocateEntry | null> { |
| 96 | if (!force) { |
| 97 | const cached = locateCache.get(mediaId); |
| 98 | if (cached && cached.expiresAt > Date.now()) return cached; |
| 99 | } |
| 100 | const reply = await requestLocate(mediaId, size); |
| 101 | if (!reply) return null; |
| 102 | const entry: LocateEntry = { ...reply, expiresAt: Date.now() + LOCATE_TTL_MS }; |
| 103 | locateCache.set(mediaId, entry); |
| 104 | return entry; |
| 105 | } |
| 106 | |
| 107 | interface RangeSlice { |
| 108 | body: Uint8Array<ArrayBuffer>; |
| 109 | status: number; |
| 110 | contentRange?: string; |
| 111 | } |
| 112 | |
| 113 | /** The engine hands back the whole attachment at once, so seeking in <video>/<audio> needs this SW to slice it. */ |
| 114 | function sliceForRange(bytes: Uint8Array<ArrayBuffer>, rangeHeader: string | null): RangeSlice { |
| 115 | const match = rangeHeader ? /^bytes=(\d*)-(\d*)$/.exec(rangeHeader) : null; |
| 116 | if (!match) return { body: bytes, status: 200 }; |
| 117 | const total = bytes.byteLength; |
| 118 | let start = match[1] ? parseInt(match[1], 10) : undefined; |
| 119 | let end = match[2] ? parseInt(match[2], 10) : undefined; |
| 120 | if (start === undefined && end !== undefined) { |
| 121 | start = Math.max(total - end, 0); |
| 122 | end = total - 1; |
| 123 | } else if (end === undefined) { |
| 124 | end = total - 1; |
| 125 | } |
| 126 | if (start === undefined || start > end || start >= total) return { body: bytes, status: 200 }; |
| 127 | end = Math.min(end, total - 1); |
| 128 | return { body: bytes.subarray(start, end + 1), status: 206, contentRange: `bytes ${start}-${end}/${total}` }; |
| 129 | } |
| 130 | |
| 131 | function mediaResponse(reply: MediaReply | null, rangeHeader: string | null, attachment: boolean): Response { |
| 132 | if (!reply) return new Response(null, { status: 504 }); |
| 133 | if (reply.error) return new Response(reply.error.message, { status: statusForKind(reply.error.kind) }); |
| 134 | const bytes = reply.bytes ?? new Uint8Array(); |
| 135 | const { body, status, contentRange } = sliceForRange(bytes, rangeHeader); |
| 136 | const headers = new Headers({ |
| 137 | 'Content-Type': reply.content_type ?? 'application/octet-stream', |
| 138 | 'Cache-Control': attachment ? 'private, max-age=3600' : 'private, max-age=31536000, immutable', |
| 139 | 'X-Content-Type-Options': 'nosniff' |
| 140 | }); |
| 141 | if (attachment) { |
| 142 | headers.set('Content-Disposition', reply.content_type ? 'inline' : 'attachment'); |
| 143 | headers.set('Content-Security-Policy', 'sandbox'); |
| 144 | } |
| 145 | if (contentRange) { |
| 146 | headers.set('Content-Range', contentRange); |
| 147 | headers.set('Accept-Ranges', 'bytes'); |
| 148 | } |
| 149 | return new Response(body, { status, headers }); |
| 150 | } |
| 151 | |
| 152 | interface ByteRange { |
| 153 | start: number; |
| 154 | end: number; |
| 155 | } |
| 156 | |
| 157 | /** Unlike sliceForRange, the total comes from the server's declared_size, not a buffer already in hand. */ |
| 158 | function parseByteRange(rangeHeader: string | null, total: number): ByteRange | null { |
| 159 | const match = rangeHeader ? /^bytes=(\d*)-(\d*)$/.exec(rangeHeader) : null; |
| 160 | if (!match) return null; |
| 161 | let start = match[1] ? parseInt(match[1], 10) : undefined; |
| 162 | let end = match[2] ? parseInt(match[2], 10) : undefined; |
| 163 | if (start === undefined && end !== undefined) { |
| 164 | start = Math.max(total - end, 0); |
| 165 | end = total - 1; |
| 166 | } else if (end === undefined) { |
| 167 | start = start ?? 0; |
| 168 | end = total - 1; |
| 169 | } |
| 170 | if (start === undefined || start > end || start >= total) return null; |
| 171 | return { start, end: Math.min(end, total - 1) }; |
| 172 | } |
| 173 | |
| 174 | function attachmentHeaders(contentType: string | null): Headers { |
| 175 | return new Headers({ |
| 176 | 'Content-Type': contentType ?? 'application/octet-stream', |
| 177 | 'Cache-Control': 'private, max-age=3600', |
| 178 | 'X-Content-Type-Options': 'nosniff', |
| 179 | 'Content-Disposition': contentType ? 'inline' : 'attachment', |
| 180 | 'Content-Security-Policy': 'sandbox' |
| 181 | }); |
| 182 | } |
| 183 | |
| 184 | /** |
| 185 | * The Authorization header carries the account's access token, so it's only worth the risk on the |
| 186 | * exact authenticated media endpoints, and only when the URL being fetched is the one `locate` named |
| 187 | * (never a redirect target, should a future caller pass one through). |
| 188 | */ |
| 189 | function shouldAuthorize(fetchUrl: string, locatedUrl: string): boolean { |
| 190 | try { |
| 191 | const target = new URL(fetchUrl); |
| 192 | const located = new URL(locatedUrl); |
| 193 | if (target.origin !== located.origin) return false; |
| 194 | return target.pathname.startsWith('/_matrix/client/v1/media/') || target.pathname.startsWith('/_matrix/media/'); |
| 195 | } catch { |
| 196 | return false; |
| 197 | } |
| 198 | } |
| 199 | |
| 200 | function concatChunks(parts: Uint8Array[]): Uint8Array<ArrayBuffer> { |
| 201 | const total = parts.reduce((sum, part) => sum + part.byteLength, 0); |
| 202 | const out = new Uint8Array(total); |
| 203 | let offset = 0; |
| 204 | for (const part of parts) { |
| 205 | out.set(part, offset); |
| 206 | offset += part.byteLength; |
| 207 | } |
| 208 | return out; |
| 209 | } |
| 210 | |
| 211 | async function readUpTo( |
| 212 | reader: ReadableStreamDefaultReader<Uint8Array<ArrayBuffer>>, |
| 213 | limit: number |
| 214 | ): Promise<Uint8Array<ArrayBuffer>> { |
| 215 | const chunks: Uint8Array[] = []; |
| 216 | let total = 0; |
| 217 | while (total < limit) { |
| 218 | const { value, done } = await reader.read(); |
| 219 | if (done) break; |
| 220 | chunks.push(value); |
| 221 | total += value.byteLength; |
| 222 | } |
| 223 | reader.cancel().catch(() => undefined); |
| 224 | return concatChunks(chunks).subarray(0, Math.min(limit, total)); |
| 225 | } |
| 226 | |
| 227 | /** |
| 228 | * Ciphertext in, plaintext out: skips the bytes the upstream sent but the caller didn't ask for |
| 229 | * (only nonzero when the upstream ignored the Range request), decrypts, drops the sub-16-byte |
| 230 | * alignment prefix from the very first output, and stops once `limit` bytes have been produced. |
| 231 | */ |
| 232 | async function* decryptedChunks( |
| 233 | reader: ReadableStreamDefaultReader<Uint8Array<ArrayBuffer>>, |
| 234 | decryptor: RangeDecryptor | null, |
| 235 | skipBytes: number, |
| 236 | dropPrefix: number, |
| 237 | limit: number |
| 238 | ): AsyncGenerator<Uint8Array> { |
| 239 | let toSkip = skipBytes; |
| 240 | let toDrop = dropPrefix; |
| 241 | let remaining = limit; |
| 242 | while (remaining > 0) { |
| 243 | const { value, done } = await reader.read(); |
| 244 | if (done) break; |
| 245 | let chunk = value; |
| 246 | if (toSkip > 0) { |
| 247 | if (chunk.byteLength <= toSkip) { |
| 248 | toSkip -= chunk.byteLength; |
| 249 | continue; |
| 250 | } |
| 251 | chunk = chunk.subarray(toSkip); |
| 252 | toSkip = 0; |
| 253 | } |
| 254 | const plain = decryptor ? await decryptor.push(chunk) : chunk; |
| 255 | if (plain.byteLength === 0) continue; |
| 256 | const emitted = emit(plain); |
| 257 | if (emitted) yield emitted; |
| 258 | } |
| 259 | if (remaining > 0 && decryptor) { |
| 260 | const emitted = emit(await decryptor.finish()); |
| 261 | if (emitted) yield emitted; |
| 262 | } |
| 263 | reader.cancel().catch(() => undefined); |
| 264 | |
| 265 | function emit(chunk: Uint8Array): Uint8Array | null { |
| 266 | let out = chunk; |
| 267 | if (toDrop > 0) { |
| 268 | if (out.byteLength <= toDrop) { |
| 269 | toDrop -= out.byteLength; |
| 270 | return null; |
| 271 | } |
| 272 | out = out.subarray(toDrop); |
| 273 | toDrop = 0; |
| 274 | } |
| 275 | if (out.byteLength > remaining) out = out.subarray(0, remaining); |
| 276 | remaining -= out.byteLength; |
| 277 | return out.byteLength > 0 ? out : null; |
| 278 | } |
| 279 | } |
| 280 | |
| 281 | async function headSniff(mediaId: string, file: MediaFileWire, token: string | null, retried: boolean): Promise<string | null> { |
| 282 | const cached = sniffCache.get(mediaId); |
| 283 | if (cached && cached.expiresAt > Date.now()) return cached.contentType; |
| 284 | const headers = new Headers({ Range: 'bytes=0-79' }); |
| 285 | if (shouldAuthorize(file.url, file.url) && token) headers.set('Authorization', `Bearer ${token}`); |
| 286 | let upstream: Response; |
| 287 | try { |
| 288 | upstream = await fetch(file.url, { headers }); |
| 289 | } catch { |
| 290 | return null; |
| 291 | } |
| 292 | if (upstream.status === 401 && !retried) { |
| 293 | const fresh = await locate(mediaId, 'Full', true); |
| 294 | if (fresh && typeof fresh.location === 'object' && 'File' in fresh.location) { |
| 295 | return headSniff(mediaId, fresh.location.File, fresh.token, true); |
| 296 | } |
| 297 | } |
| 298 | if (!upstream.body || (!upstream.ok && upstream.status !== 206)) return null; |
| 299 | const head = await readUpTo(upstream.body.getReader(), 80); |
| 300 | let contentType: string | null; |
| 301 | if (file.encryption) { |
| 302 | const key = await importAttachmentKey(file.encryption.key); |
| 303 | const iv = decodeAttachmentIv(file.encryption.iv); |
| 304 | const decryptor = new RangeDecryptor(key, iv, 0); |
| 305 | contentType = attachmentContentType(concatChunks([await decryptor.push(head), await decryptor.finish()]).subarray(0, 64)); |
| 306 | } else { |
| 307 | contentType = attachmentContentType(head.subarray(0, 64)); |
| 308 | } |
| 309 | sniffCache.set(mediaId, { contentType, expiresAt: Date.now() + LOCATE_TTL_MS }); |
| 310 | return contentType; |
| 311 | } |
| 312 | |
| 313 | function wholeFileFallback(mediaId: string, rangeHeader: string | null): Promise<Response> { |
| 314 | return requestMedia('media', { media_id: mediaId, size: 'Full' }).then((reply) => mediaResponse(reply, rangeHeader, true)); |
| 315 | } |
| 316 | |
| 317 | async function streamFullMedia(mediaId: string, rangeHeader: string | null, entry?: LocateEntry, retried = false): Promise<Response> { |
| 318 | const located = entry ?? (await locate(mediaId, 'Full', false)); |
| 319 | if (!located) return wholeFileFallback(mediaId, rangeHeader); |
| 320 | if (typeof located.location !== 'object' || !('File' in located.location)) return wholeFileFallback(mediaId, rangeHeader); |
| 321 | const file = located.location.File; |
| 322 | const total = file.declared_size; |
| 323 | if (total == null) return wholeFileFallback(mediaId, rangeHeader); |
| 324 | |
| 325 | const requested = parseByteRange(rangeHeader, total); |
| 326 | const start = requested?.start ?? 0; |
| 327 | const end = requested?.end ?? total - 1; |
| 328 | const alignedStart = start - (start % 16); |
| 329 | |
| 330 | const headers = new Headers({ Range: `bytes=${alignedStart}-${end}` }); |
| 331 | if (shouldAuthorize(file.url, file.url) && located.token) headers.set('Authorization', `Bearer ${located.token}`); |
| 332 | |
| 333 | let upstream: Response; |
| 334 | try { |
| 335 | upstream = await fetch(file.url, { headers }); |
| 336 | } catch { |
| 337 | return new Response('media upstream unreachable', { status: 502 }); |
| 338 | } |
| 339 | if (upstream.status === 401 && !retried) { |
| 340 | const fresh = await locate(mediaId, 'Full', true); |
| 341 | if (fresh) return streamFullMedia(mediaId, rangeHeader, fresh, true); |
| 342 | } |
| 343 | if (!upstream.body || (!upstream.ok && upstream.status !== 206)) { |
| 344 | return new Response('media upstream error', { status: 502 }); |
| 345 | } |
| 346 | |
| 347 | const decryptor = file.encryption |
| 348 | ? new RangeDecryptor(await importAttachmentKey(file.encryption.key), decodeAttachmentIv(file.encryption.iv), alignedStart / 16) |
| 349 | : null; |
| 350 | const skipBytes = upstream.status === 206 ? 0 : alignedStart; |
| 351 | const dropPrefix = start - alignedStart; |
| 352 | const limit = end - start + 1; |
| 353 | |
| 354 | const reader = upstream.body.getReader(); |
| 355 | const chunks = decryptedChunks(reader, decryptor, skipBytes, dropPrefix, limit); |
| 356 | |
| 357 | const prefetched: Uint8Array[] = []; |
| 358 | let prefetchedLen = 0; |
| 359 | let chunksDone = false; |
| 360 | let contentType: string | null; |
| 361 | |
| 362 | if (alignedStart === 0) { |
| 363 | while (prefetchedLen < 64) { |
| 364 | const step = await chunks.next(); |
| 365 | if (step.done) { |
| 366 | chunksDone = true; |
| 367 | break; |
| 368 | } |
| 369 | prefetched.push(step.value); |
| 370 | prefetchedLen += step.value.byteLength; |
| 371 | } |
| 372 | contentType = attachmentContentType(concatChunks(prefetched).subarray(0, Math.min(64, prefetchedLen))); |
| 373 | sniffCache.set(mediaId, { contentType, expiresAt: Date.now() + LOCATE_TTL_MS }); |
| 374 | } else { |
| 375 | contentType = await headSniff(mediaId, file, located.token, retried); |
| 376 | } |
| 377 | |
| 378 | const body = new ReadableStream<Uint8Array>({ |
| 379 | async pull(controller) { |
| 380 | const next = prefetched.shift(); |
| 381 | if (next) { |
| 382 | controller.enqueue(next); |
| 383 | return; |
| 384 | } |
| 385 | if (chunksDone) { |
| 386 | controller.close(); |
| 387 | return; |
| 388 | } |
| 389 | const step = await chunks.next(); |
| 390 | if (step.done) controller.close(); |
| 391 | else controller.enqueue(step.value); |
| 392 | }, |
| 393 | cancel() { |
| 394 | reader.cancel().catch(() => undefined); |
| 395 | } |
| 396 | }); |
| 397 | |
| 398 | const responseHeaders = attachmentHeaders(contentType); |
| 399 | responseHeaders.set('Accept-Ranges', 'bytes'); |
| 400 | if (requested) { |
| 401 | responseHeaders.set('Content-Range', `bytes ${start}-${end}/${total}`); |
| 402 | return new Response(body, { status: 206, headers: responseHeaders }); |
| 403 | } |
| 404 | return new Response(body, { status: 200, headers: responseHeaders }); |
| 405 | } |
| 406 | |
| 407 | /** Keyed by the request URL, which names image and size, so cached bytes outlive the access token and survive restarts. */ |
| 408 | let avatarLoader: ReturnType<typeof createAvatarLoader> | undefined; |
| 409 | |
| 410 | function cachedAvatar(url: string, query: string): Promise<Response> { |
| 411 | avatarLoader ??= createAvatarLoader( |
| 412 | async () => caches.open(AVATAR_CACHE), |
| 413 | async (q) => mediaResponse(await requestMedia('avatar', Object.fromEntries(new URLSearchParams(q))), null, false) |
| 414 | ); |
| 415 | return avatarLoader(url, query); |
| 416 | } |
| 417 | |
| 418 | export function handleMediaFetch(event: FetchEvent): boolean { |
| 419 | const url = new URL(event.request.url); |
| 420 | const range = event.request.headers.get('Range'); |
| 421 | |
| 422 | if (url.pathname === '/api/media') { |
| 423 | const mxc = url.searchParams.get('mxc'); |
| 424 | if (mxc) { |
| 425 | const size = url.searchParams.get('size') ?? '256'; |
| 426 | event.respondWith(cachedAvatar(event.request.url, new URLSearchParams({ mxc, size }).toString())); |
| 427 | return true; |
| 428 | } |
| 429 | } |
| 430 | if (url.pathname === '/api/message-media') { |
| 431 | const mediaId = url.searchParams.get('media_id'); |
| 432 | const size = url.searchParams.get('size'); |
| 433 | if (mediaId && size === 'Full') { |
| 434 | event.respondWith(streamFullMedia(mediaId, range)); |
| 435 | return true; |
| 436 | } |
| 437 | if (mediaId && size) { |
| 438 | event.respondWith(requestMedia('media', { media_id: mediaId, size }).then((reply) => mediaResponse(reply, range, true))); |
| 439 | return true; |
| 440 | } |
| 441 | } |
| 442 | return false; |
| 443 | } |