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

import { describe, test, expect, beforeEach } from 'bun:test';
import { createPrekeyServer, MemoryPrekeyStore, signPayload } from '../src/index.js';
import type { PrekeyStore } from '../src/index.js';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import { generateIdentityKeyPair } from '@shade/core';

const crypto = new SubtleCryptoProvider();

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

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

async function makeIdentity() {
  return generateIdentityKeyPair(crypto);
}

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

  beforeEach(() => {
    store = new MemoryPrekeyStore();
    app = createPrekeyServer({ crypto, store, disableRateLimit: true });
  });

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

  /** Helper: build a signed registration body for a given identity */
  async function signedRegisterBody(identity: Awaited<ReturnType<typeof makeIdentity>>, address: string, withOTPKs = true) {
    const body: any = {
      address,
      identitySigningKey: b64(identity.signingPublicKey),
      identityDHKey: b64(identity.dhPublicKey),
      signedPreKey: {
        keyId: 1,
        publicKey: b64(randBytes(32)),
        signature: b64(randBytes(64)),
      },
    };
    if (withOTPKs) {
      body.oneTimePreKeys = [
        { keyId: 100, publicKey: b64(randBytes(32)) },
        { keyId: 101, publicKey: b64(randBytes(32)) },
      ];
    }
    return signPayload(crypto, identity.signingPrivateKey, body);
  }

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

  describe('POST /v1/keys/register', () => {
    test('accepts valid signed registration', async () => {
      const alice = await makeIdentity();
      const body = await signedRegisterBody(alice, 'alice');
      const res = await req('POST', '/v1/keys/register', body);
      expect(res.status).toBe(200);
    });

    test('rejects unsigned registration', async () => {
      const alice = await makeIdentity();
      const res = await req('POST', '/v1/keys/register', {
        address: 'alice',
        identitySigningKey: b64(alice.signingPublicKey),
        identityDHKey: b64(alice.dhPublicKey),
        signedPreKey: { keyId: 1, publicKey: b64(randBytes(32)), signature: b64(randBytes(64)) },
      });
      // Missing signature/signedAt → validation error
      expect(res.status).toBe(400);
    });

    test('rejects registration with wrong signing key', async () => {
      const alice = await makeIdentity();
      const bob = await makeIdentity();
      // Sign with bob's key but claim alice's public key
      const body: any = {
        address: 'alice',
        identitySigningKey: b64(alice.signingPublicKey), // mismatch
        identityDHKey: b64(alice.dhPublicKey),
        signedPreKey: { keyId: 1, publicKey: b64(randBytes(32)), signature: b64(randBytes(64)) },
      };
      const signed = await signPayload(crypto, bob.signingPrivateKey, body);
      const res = await req('POST', '/v1/keys/register', signed);
      expect(res.status).toBe(401);
    });

    test('rejects registration with stale signedAt', async () => {
      const alice = await makeIdentity();
      const body = await signedRegisterBody(alice, 'alice');
      // Tamper with signedAt to be old
      body.signedAt = Date.now() - 10 * 60 * 1000; // 10 minutes ago
      const res = await req('POST', '/v1/keys/register', body);
      expect(res.status).toBe(409); // ReplayError
    });

    test('rejects invalid address format', async () => {
      const alice = await makeIdentity();
      const body = await signedRegisterBody(alice, '../evil');
      const res = await req('POST', '/v1/keys/register', body);
      expect(res.status).toBe(400);
    });

    test('accepts registration with one-time prekeys', async () => {
      const alice = await makeIdentity();
      const body = await signedRegisterBody(alice, 'alice');
      await req('POST', '/v1/keys/register', body);

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

  // ─── Fetch Bundle (anonymous) ──────────────────────────────

  describe('GET /v1/keys/bundle/:address', () => {
    test('returns bundle for registered address', async () => {
      const bob = await makeIdentity();
      const body = await signedRegisterBody(bob, 'bob');
      await req('POST', '/v1/keys/register', body);

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

      const bundle = await res.json();
      expect(bundle.identitySigningKey).toBe(b64(bob.signingPublicKey));
      expect(bundle.identityDHKey).toBe(b64(bob.dhPublicKey));
      expect(bundle.signedPreKey.keyId).toBe(1);
      expect(bundle.oneTimePreKey.keyId).toBe(100);
    });

    test('consumes one-time prekey on each fetch (FIFO)', async () => {
      const bob = await makeIdentity();
      await req('POST', '/v1/keys/register', await signedRegisterBody(bob, 'bob'));

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

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

      const res3 = await req('GET', '/v1/keys/bundle/bob');
      expect((await res3.json()).oneTimePreKey).toBeUndefined();
    });

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

    test('rejects invalid address in URL', async () => {
      const res = await req('GET', '/v1/keys/bundle/..evil');
      expect(res.status).toBe(400);
    });
  });

  // ─── Replenish (signed) ────────────────────────────────────

  describe('POST /v1/keys/replenish', () => {
    test('accepts signed replenishment from registered identity', async () => {
      const bob = await makeIdentity();
      await req('POST', '/v1/keys/register', await signedRegisterBody(bob, 'bob'));

      const replenishBody = await signPayload(crypto, bob.signingPrivateKey, {
        address: 'bob',
        oneTimePreKeys: [
          { keyId: 200, publicKey: b64(randBytes(32)) },
          { keyId: 201, publicKey: b64(randBytes(32)) },
        ],
      });

      const res = await req('POST', '/v1/keys/replenish', replenishBody);
      expect(res.status).toBe(200);
      const body = await res.json();
      expect(body.remaining).toBe(4); // 2 original + 2 new
    });

    test('rejects replenishment signed by wrong identity', async () => {
      const bob = await makeIdentity();
      const eve = await makeIdentity();
      await req('POST', '/v1/keys/register', await signedRegisterBody(bob, 'bob'));

      const evilBody = await signPayload(crypto, eve.signingPrivateKey, {
        address: 'bob',
        oneTimePreKeys: [{ keyId: 200, publicKey: b64(randBytes(32)) }],
      });

      const res = await req('POST', '/v1/keys/replenish', evilBody);
      expect(res.status).toBe(401);
    });

    test('rejects replenishment for unknown address', async () => {
      const bob = await makeIdentity();
      const body = await signPayload(crypto, bob.signingPrivateKey, {
        address: 'nobody',
        oneTimePreKeys: [{ keyId: 200, publicKey: b64(randBytes(32)) }],
      });
      const res = await req('POST', '/v1/keys/replenish', body);
      expect(res.status).toBe(404);
    });
  });

  // ─── Delete (signed) ───────────────────────────────────────

  describe('DELETE /v1/keys/:address', () => {
    test('accepts signed delete from registered identity', async () => {
      const bob = await makeIdentity();
      await req('POST', '/v1/keys/register', await signedRegisterBody(bob, 'bob'));

      const delBody = await signPayload(crypto, bob.signingPrivateKey, { address: 'bob' });
      const res = await req('DELETE', '/v1/keys/bob', delBody);
      expect(res.status).toBe(200);

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

    test('rejects delete signed by wrong identity', async () => {
      const bob = await makeIdentity();
      const eve = await makeIdentity();
      await req('POST', '/v1/keys/register', await signedRegisterBody(bob, 'bob'));

      const evilBody = await signPayload(crypto, eve.signingPrivateKey, { address: 'bob' });
      const res = await req('DELETE', '/v1/keys/bob', evilBody);
      expect(res.status).toBe(401);

      // Should still exist
      const bundleRes = await req('GET', '/v1/keys/bundle/bob');
      expect(bundleRes.status).toBe(200);
    });
  });

  // ─── Multi-address isolation ───────────────────────────────

  describe('multi-address isolation', () => {
    test('different addresses are independent', async () => {
      const alice = await makeIdentity();
      const bob = await makeIdentity();
      const charlie = await makeIdentity();

      await req('POST', '/v1/keys/register', await signedRegisterBody(alice, 'alice'));
      await req('POST', '/v1/keys/register', await signedRegisterBody(bob, 'bob'));
      await req('POST', '/v1/keys/register', await signedRegisterBody(charlie, 'charlie'));

      // Delete bob with his own signature
      const delBody = await signPayload(crypto, bob.signingPrivateKey, { address: 'bob' });
      await req('DELETE', '/v1/keys/bob', delBody);

      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);
    });
  });
});