Shade/packages/shade-server/tests/server.test.ts

import { describe, test, expect, beforeEach } from 'bun:test';
import { createPrekeyServer, MemoryPrekeyStore } from '../src/index.js';
import type { PrekeyStore } from '../src/index.js';

function b64(bytes: Uint8Array): string {
  return Buffer.from(bytes).toString('base64');
}

function randBytes(n: number): Uint8Array {
  const buf = new Uint8Array(n);
  crypto.getRandomValues(buf);
  return buf;
}

describe('Shade Prekey Server', () => {
  let store: PrekeyStore;
  let app: ReturnType<typeof createPrekeyServer>;

  beforeEach(() => {
    store = new MemoryPrekeyStore();
    app = createPrekeyServer({ store });
  });

  function req(method: string, path: string, body?: any) {
    const init: RequestInit = { method, headers: { 'Content-Type': 'application/json' } };
    if (body) init.body = JSON.stringify(body);
    return app.request(path, init);
  }

  // ─── Registration ──────────────────────────────────────────

  describe('POST /v1/keys/register', () => {
    test('registers identity and signed prekey', async () => {
      const res = await req('POST', '/v1/keys/register', {
        address: 'alice',
        identitySigningKey: b64(randBytes(32)),
        identityDHKey: b64(randBytes(32)),
        signedPreKey: {
          keyId: 1,
          publicKey: b64(randBytes(32)),
          signature: b64(randBytes(64)),
        },
      });
      expect(res.status).toBe(200);
      expect(await res.json()).toEqual({ ok: true });
    });

    test('registers with one-time prekeys', async () => {
      const res = await req('POST', '/v1/keys/register', {
        address: 'alice',
        identitySigningKey: b64(randBytes(32)),
        identityDHKey: b64(randBytes(32)),
        signedPreKey: {
          keyId: 1,
          publicKey: b64(randBytes(32)),
          signature: b64(randBytes(64)),
        },
        oneTimePreKeys: [
          { keyId: 100, publicKey: b64(randBytes(32)) },
          { keyId: 101, publicKey: b64(randBytes(32)) },
          { keyId: 102, publicKey: b64(randBytes(32)) },
        ],
      });
      expect(res.status).toBe(200);

      // Verify count
      const countRes = await req('GET', '/v1/keys/count/alice');
      expect((await countRes.json()).count).toBe(3);
    });

    test('rejects missing fields', async () => {
      const res = await req('POST', '/v1/keys/register', { address: 'alice' });
      expect(res.status).toBe(400);
    });
  });

  // ─── Fetch Bundle ──────────────────────────────────────────

  describe('GET /v1/keys/bundle/:address', () => {
    test('returns bundle with one-time prekey', async () => {
      // Register first
      const sigKey = b64(randBytes(32));
      const dhKey = b64(randBytes(32));
      const spkPub = b64(randBytes(32));
      const spkSig = b64(randBytes(64));
      const otpkPub = b64(randBytes(32));

      await req('POST', '/v1/keys/register', {
        address: 'bob',
        identitySigningKey: sigKey,
        identityDHKey: dhKey,
        signedPreKey: { keyId: 1, publicKey: spkPub, signature: spkSig },
        oneTimePreKeys: [{ keyId: 100, publicKey: otpkPub }],
      });

      const res = await req('GET', '/v1/keys/bundle/bob');
      expect(res.status).toBe(200);

      const bundle = await res.json();
      expect(bundle.identitySigningKey).toBe(sigKey);
      expect(bundle.identityDHKey).toBe(dhKey);
      expect(bundle.signedPreKey.keyId).toBe(1);
      expect(bundle.signedPreKey.publicKey).toBe(spkPub);
      expect(bundle.signedPreKey.signature).toBe(spkSig);
      expect(bundle.oneTimePreKey.keyId).toBe(100);
      expect(bundle.oneTimePreKey.publicKey).toBe(otpkPub);
    });

    test('returns bundle without one-time prekey when depleted', async () => {
      await req('POST', '/v1/keys/register', {
        address: 'bob',
        identitySigningKey: b64(randBytes(32)),
        identityDHKey: b64(randBytes(32)),
        signedPreKey: { keyId: 1, publicKey: b64(randBytes(32)), signature: b64(randBytes(64)) },
      });

      const res = await req('GET', '/v1/keys/bundle/bob');
      expect(res.status).toBe(200);

      const bundle = await res.json();
      expect(bundle.oneTimePreKey).toBeUndefined();
    });

    test('consumes one-time prekeys on each fetch', async () => {
      await req('POST', '/v1/keys/register', {
        address: 'bob',
        identitySigningKey: b64(randBytes(32)),
        identityDHKey: b64(randBytes(32)),
        signedPreKey: { keyId: 1, publicKey: b64(randBytes(32)), signature: b64(randBytes(64)) },
        oneTimePreKeys: [
          { keyId: 100, publicKey: b64(randBytes(32)) },
          { keyId: 101, publicKey: b64(randBytes(32)) },
        ],
      });

      // First fetch consumes key 100
      const res1 = await req('GET', '/v1/keys/bundle/bob');
      expect((await res1.json()).oneTimePreKey.keyId).toBe(100);

      // Second fetch consumes key 101
      const res2 = await req('GET', '/v1/keys/bundle/bob');
      expect((await res2.json()).oneTimePreKey.keyId).toBe(101);

      // Third fetch has none left
      const res3 = await req('GET', '/v1/keys/bundle/bob');
      expect((await res3.json()).oneTimePreKey).toBeUndefined();

      // Count should be 0
      const countRes = await req('GET', '/v1/keys/count/bob');
      expect((await countRes.json()).count).toBe(0);
    });

    test('404 for unknown address', async () => {
      const res = await req('GET', '/v1/keys/bundle/nobody');
      expect(res.status).toBe(404);
    });
  });

  // ─── Replenish ─────────────────────────────────────────────

  describe('POST /v1/keys/replenish', () => {
    test('adds more one-time prekeys', async () => {
      await req('POST', '/v1/keys/register', {
        address: 'bob',
        identitySigningKey: b64(randBytes(32)),
        identityDHKey: b64(randBytes(32)),
        signedPreKey: { keyId: 1, publicKey: b64(randBytes(32)), signature: b64(randBytes(64)) },
        oneTimePreKeys: [{ keyId: 100, publicKey: b64(randBytes(32)) }],
      });

      const res = await req('POST', '/v1/keys/replenish', {
        address: 'bob',
        oneTimePreKeys: [
          { keyId: 200, publicKey: b64(randBytes(32)) },
          { keyId: 201, publicKey: b64(randBytes(32)) },
        ],
      });

      expect(res.status).toBe(200);
      const body = await res.json();
      expect(body.remaining).toBe(3); // 1 original + 2 new
    });
  });

  // ─── Delete ────────────────────────────────────────────────

  describe('DELETE /v1/keys/:address', () => {
    test('removes all keys for an address', async () => {
      await req('POST', '/v1/keys/register', {
        address: 'bob',
        identitySigningKey: b64(randBytes(32)),
        identityDHKey: b64(randBytes(32)),
        signedPreKey: { keyId: 1, publicKey: b64(randBytes(32)), signature: b64(randBytes(64)) },
        oneTimePreKeys: [{ keyId: 100, publicKey: b64(randBytes(32)) }],
      });

      const delRes = await req('DELETE', '/v1/keys/bob');
      expect(delRes.status).toBe(200);

      // Should be gone
      const bundleRes = await req('GET', '/v1/keys/bundle/bob');
      expect(bundleRes.status).toBe(404);

      const countRes = await req('GET', '/v1/keys/count/bob');
      expect((await countRes.json()).count).toBe(0);
    });
  });

  // ─── Multiple Addresses ────────────────────────────────────

  describe('multi-address isolation', () => {
    test('different addresses are independent', async () => {
      for (const addr of ['alice', 'bob', 'charlie']) {
        await req('POST', '/v1/keys/register', {
          address: addr,
          identitySigningKey: b64(randBytes(32)),
          identityDHKey: b64(randBytes(32)),
          signedPreKey: { keyId: 1, publicKey: b64(randBytes(32)), signature: b64(randBytes(64)) },
          oneTimePreKeys: [{ keyId: 1, publicKey: b64(randBytes(32)) }],
        });
      }

      // Delete bob, others remain
      await req('DELETE', '/v1/keys/bob');

      expect((await req('GET', '/v1/keys/bundle/alice')).status).toBe(200);
      expect((await req('GET', '/v1/keys/bundle/bob')).status).toBe(404);
      expect((await req('GET', '/v1/keys/bundle/charlie')).status).toBe(200);
    });
  });
});