Shade/packages/shade-sdk/tests/thumbnail.test.ts

import { describe, test, expect, beforeAll, afterAll } from 'bun:test';
import {
  createShade,
  ShadeThumbnailCache,
  THUMBNAIL_MAX_BYTES,
  type IncomingTransfer,
  type Shade,
  type TransferHandle,
  type TransferResult,
} from '../src/index.js';
import {
  createPrekeyServer,
  MemoryPrekeyStore,
  PrekeyServerEvents,
} from '@shade/server';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import { sha256Once } from '@shade/streams';

const crypto = new SubtleCryptoProvider();

interface TestRig {
  alice: Shade;
  bob: Shade;
  prekeyStop: () => void;
  bobServerStop: () => void;
}

async function startPrekeyServer(): Promise<{ url: string; stop: () => void }> {
  const events = new PrekeyServerEvents();
  const server = createPrekeyServer({
    crypto,
    store: new MemoryPrekeyStore(),
    disableRateLimit: true,
    events,
  });
  const port = 22000 + Math.floor(Math.random() * 500);
  const handle = Bun.serve({ port, fetch: server.fetch });
  return { url: `http://localhost:${port}`, stop: () => handle.stop() };
}

async function setupRig(): Promise<TestRig> {
  const prekey = await startPrekeyServer();
  const alice = await createShade({ prekeyServer: prekey.url, address: 'alice' });
  const bob = await createShade({ prekeyServer: prekey.url, address: 'bob' });

  bob.configureTransfers({
    resolveBaseUrl: async () => {
      throw new Error('bob is receive-only');
    },
  });
  const bobApp = await bob.transferRoute();
  const port = 22500 + Math.floor(Math.random() * 500);
  const bobServer = Bun.serve({ port, fetch: bobApp.fetch });
  const bobBaseUrl = `http://localhost:${port}`;

  alice.configureTransfers({
    resolveBaseUrl: async (addr) => {
      if (addr === 'bob') return bobBaseUrl;
      throw new Error(`unknown peer ${addr}`);
    },
  });

  return {
    alice,
    bob,
    prekeyStop: prekey.stop,
    bobServerStop: () => bobServer.stop(),
  };
}

async function teardownRig(rig: TestRig): Promise<void> {
  await rig.alice.shutdown();
  await rig.bob.shutdown();
  rig.bobServerStop();
  rig.prekeyStop();
}

function fakeJpeg(size: number): Uint8Array {
  // Synthetic "JPEG-shaped" bytes: SOI marker + filler. The SDK only
  // hashes + ships these; the receiver-side validator we exercise is
  // MIME + size, not actual decode-ability. The widget renderer feeds
  // these to a real `<img>`; that's where format-correctness matters.
  const buf = new Uint8Array(size);
  buf[0] = 0xff;
  buf[1] = 0xd8;
  buf[2] = 0xff;
  for (let i = 3; i < size; i++) buf[i] = i & 0xff;
  return buf;
}

describe('V3.9 thumbnail roundtrip', () => {
  let rig: TestRig;
  beforeAll(async () => {
    rig = await setupRig();
  });
  afterAll(async () => {
    await teardownRig(rig);
  });

  test('upload with thumbnail attaches fileMetadata + ships separate stream', async () => {
    const main = crypto.randomBytes(64 * 1024);
    const thumb = fakeJpeg(4096);
    const expectedHashB64 = bytesToBase64(sha256Once(thumb));

    const incomings: IncomingTransfer[] = [];
    const recvHandles: TransferHandle[] = [];
    const unsub = await rig.bob.onIncomingTransfer(async (incoming) => {
      incomings.push(incoming);
      const handle = await incoming.accept({ output: { kind: 'buffer' } });
      recvHandles.push(handle);
    });

    const mainHandle = await rig.alice.upload({
      to: 'bob',
      input: main,
      thumbnail: { bytes: thumb, mime: 'image/jpeg' },
      lanes: 1,
      chunkSize: 16 * 1024,
      metadata: {
        fileMetadata: {
          filename: 'doc.pdf',
          mimeType: 'application/pdf',
        },
      },
    });
    await mainHandle.done();
    // Wait for any background thumbnail finish before tearing down.
    for (const h of recvHandles) await h.done();
    unsub();

    // Two transfers: thumb then main (thumb is shipped first).
    expect(incomings.length).toBe(2);
    const thumbIncoming = incomings.find(
      (i) => i.metadata.userMetadata?.shadeThumbnail === '1',
    );
    const mainIncoming = incomings.find(
      (i) => i.metadata.userMetadata?.shadeThumbnail !== '1',
    );
    expect(thumbIncoming).toBeDefined();
    expect(mainIncoming).toBeDefined();
    expect(thumbIncoming!.metadata.contentType).toBe('image/jpeg');
    expect(mainIncoming!.metadata.fileMetadata?.thumbnailHash).toBe(expectedHashB64);
    expect(mainIncoming!.metadata.fileMetadata?.thumbnailMime).toBe('image/jpeg');
    expect(mainIncoming!.metadata.fileMetadata?.thumbnailBytes).toBe(thumb.byteLength);
    expect(mainIncoming!.metadata.fileMetadata?.thumbnailStreamId).toBe(
      thumbIncoming!.streamId,
    );
    expect(mainIncoming!.metadata.fileMetadata?.filename).toBe('doc.pdf');
    expect(mainIncoming!.metadata.fileMetadata?.mimeType).toBe('application/pdf');
  });

  test('legacy receiver (no thumbnail handling) still receives main stream', async () => {
    const main = crypto.randomBytes(8 * 1024);
    const thumb = fakeJpeg(2048);

    let resolveMain!: (h: TransferHandle) => void;
    const mainHandlePromise = new Promise<TransferHandle>((r) => {
      resolveMain = r;
    });
    const allHandles: TransferHandle[] = [];

    const unsub = await rig.bob.onIncomingTransfer(async (incoming) => {
      // "Legacy" path: ignore fileMetadata entirely. Just accept everything.
      const handle = await incoming.accept({ output: { kind: 'buffer' } });
      allHandles.push(handle);
      // Resolve as soon as we see the main (non-thumb) stream.
      if (incoming.metadata.userMetadata?.shadeThumbnail !== '1') {
        resolveMain(handle);
      }
    });

    const senderHandle = await rig.alice.upload({
      to: 'bob',
      input: main,
      thumbnail: { bytes: thumb, mime: 'image/jpeg' },
      lanes: 1,
      chunkSize: 4 * 1024,
    });
    const mainHandle = await mainHandlePromise;
    const [senderResult, mainResult] = await Promise.all([
      senderHandle.done(),
      mainHandle.done(),
    ]);
    for (const h of allHandles) await h.done();
    unsub();

    expect(
      (mainResult as TransferResult & { bytes?: Uint8Array }).bytes,
    ).toEqual(main);
    expect(senderResult.bytesSent).toBe(main.byteLength);
  });

  test('throws on oversize thumbnail bytes', async () => {
    const main = crypto.randomBytes(1024);
    const tooLarge = fakeJpeg(THUMBNAIL_MAX_BYTES + 1);
    await expect(
      rig.alice.upload({
        to: 'bob',
        input: main,
        thumbnail: { bytes: tooLarge, mime: 'image/jpeg' },
      }),
    ).rejects.toThrow();
  });

  test('throws on disallowed thumbnail mime', async () => {
    const main = crypto.randomBytes(1024);
    await expect(
      rig.alice.upload({
        to: 'bob',
        input: main,
        // @ts-expect-error — testing runtime guard
        thumbnail: { bytes: fakeJpeg(1024), mime: 'image/svg+xml' },
      }),
    ).rejects.toThrow();
  });
});

describe('V3.9 ShadeThumbnailCache', () => {
  test('rejects oversize bytes', () => {
    const cache = new ShadeThumbnailCache();
    const oversize = new Uint8Array(THUMBNAIL_MAX_BYTES + 1);
    expect(cache.put('s1', oversize, 'image/jpeg')).toBe(false);
    expect(cache.size).toBe(0);
  });

  test('rejects disallowed mime', () => {
    const cache = new ShadeThumbnailCache();
    const tiny = new Uint8Array(8);
    expect(cache.put('s1', tiny, 'image/svg+xml')).toBe(false);
    expect(cache.size).toBe(0);
  });

  test('drops bytes whose hash does not match expectedHash', () => {
    const cache = new ShadeThumbnailCache();
    const bytes = new Uint8Array([1, 2, 3, 4]);
    const wrongHash = bytesToBase64(new Uint8Array(32)); // all zero
    cache.setExpectedHash('s1', wrongHash);
    expect(cache.put('s1', bytes, 'image/jpeg')).toBe(false);
    expect(cache.get('s1')).toBeNull();
  });

  test('round-trips when hash matches', () => {
    const cache = new ShadeThumbnailCache();
    const bytes = new Uint8Array([0xff, 0xd8, 0xff, 0x42]);
    const hash = bytesToBase64(sha256Once(bytes));
    cache.setExpectedHash('s1', hash);
    expect(cache.put('s1', bytes, 'image/jpeg')).toBe(true);
    const hit = cache.get('s1', hash);
    expect(hit).not.toBeNull();
    expect(hit!.bytes).toEqual(bytes);
    expect(hit!.mime).toBe('image/jpeg');
  });

  test('get drops entry on hash mismatch', () => {
    const cache = new ShadeThumbnailCache();
    const bytes = new Uint8Array([1, 2, 3]);
    cache.put('s1', bytes, 'image/png');
    const wrongHash = bytesToBase64(new Uint8Array(32));
    expect(cache.get('s1', wrongHash)).toBeNull();
    expect(cache.size).toBe(0);
  });

  test('emits onChange when an entry is added', () => {
    const cache = new ShadeThumbnailCache();
    const seen: string[] = [];
    cache.onChange((s) => seen.push(s));
    cache.put('s1', new Uint8Array([1]), 'image/jpeg');
    cache.put('s2', new Uint8Array([2]), 'image/jpeg');
    expect(seen).toEqual(['s1', 's2']);
  });
});

function bytesToBase64(bytes: Uint8Array): string {
  let bin = '';
  for (let i = 0; i < bytes.length; i++) bin += String.fromCharCode(bytes[i]!);
  return btoa(bin);
}