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

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

const crypto = new SubtleCryptoProvider();

describe('RateLimiter', () => {
  test('allows requests up to capacity', async () => {
    const store = new MemoryRateLimitStore();
    const rl = new RateLimiter(store, { capacity: 5, refillPerSecond: 1 });

    for (let i = 0; i < 5; i++) {
      await rl.consume('user:1');
    }
  });

  test('throws RateLimitError when exhausted', async () => {
    const store = new MemoryRateLimitStore();
    const rl = new RateLimiter(store, { capacity: 3, refillPerSecond: 0.1 });

    for (let i = 0; i < 3; i++) await rl.consume('user:1');
    expect(rl.consume('user:1')).rejects.toThrow(RateLimitError);
  });

  test('different keys have independent limits', async () => {
    const store = new MemoryRateLimitStore();
    const rl = new RateLimiter(store, { capacity: 2, refillPerSecond: 0.1 });

    await rl.consume('user:1');
    await rl.consume('user:1');
    // user:1 exhausted
    expect(rl.consume('user:1')).rejects.toThrow(RateLimitError);

    // user:2 still has full capacity
    await rl.consume('user:2');
    await rl.consume('user:2');
  });

  test('refills over time', async () => {
    const store = new MemoryRateLimitStore();
    const rl = new RateLimiter(store, { capacity: 2, refillPerSecond: 100 }); // fast refill for test

    await rl.consume('user:1');
    await rl.consume('user:1');

    // Wait a bit for refill (100 tokens/sec → 2 tokens in 20ms)
    await new Promise((r) => setTimeout(r, 30));

    // Should be refilled
    await rl.consume('user:1');
  });

  test('RateLimitError has retryAfterSeconds', async () => {
    const store = new MemoryRateLimitStore();
    const rl = new RateLimiter(store, { capacity: 1, refillPerSecond: 0.5 });

    await rl.consume('user:1');
    try {
      await rl.consume('user:1');
      expect.unreachable();
    } catch (err) {
      expect(err).toBeInstanceOf(RateLimitError);
      expect((err as RateLimitError).retryAfterSeconds).toBeGreaterThan(0);
    }
  });
});

describe('Rate limiting integration with routes', () => {
  test('register endpoint rate-limits per IP', async () => {
    const app = createPrekeyServer({ crypto, store: new MemoryPrekeyStore() });

    async function doRegister(addressSuffix: number) {
      const identity = await generateIdentityKeyPair(crypto);
      const body: any = {
        address: `user${addressSuffix}`,
        identitySigningKey: Buffer.from(identity.signingPublicKey).toString('base64'),
        identityDHKey: Buffer.from(identity.dhPublicKey).toString('base64'),
        signedPreKey: {
          keyId: 1,
          publicKey: Buffer.from(crypto.randomBytes(32)).toString('base64'),
          signature: Buffer.from(crypto.randomBytes(64)).toString('base64'),
        },
      };
      const signed = await signPayload(crypto, identity.signingPrivateKey, body);
      return app.request('/v1/keys/register', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json', 'x-forwarded-for': '203.0.113.1' },
        body: JSON.stringify(signed),
      });
    }

    // Register limit is 5/hour, so after 5 successful, the 6th should fail
    const results: number[] = [];
    for (let i = 0; i < 7; i++) {
      const res = await doRegister(i);
      results.push(res.status);
    }

    // First 5 should succeed (200), rest should be rate-limited (429)
    expect(results.filter((s) => s === 200).length).toBeGreaterThanOrEqual(5);
    expect(results.filter((s) => s === 429).length).toBeGreaterThanOrEqual(1);
  });

  test('rate limit returns Retry-After header', async () => {
    const app = createPrekeyServer({ crypto, store: new MemoryPrekeyStore() });

    // Burn through the delete limit (5/hour)
    const identity = await generateIdentityKeyPair(crypto);

    // First register
    const regBody: any = {
      address: 'ratelimit',
      identitySigningKey: Buffer.from(identity.signingPublicKey).toString('base64'),
      identityDHKey: Buffer.from(identity.dhPublicKey).toString('base64'),
      signedPreKey: {
        keyId: 1,
        publicKey: Buffer.from(crypto.randomBytes(32)).toString('base64'),
        signature: Buffer.from(crypto.randomBytes(64)).toString('base64'),
      },
    };
    const signedReg = await signPayload(crypto, identity.signingPrivateKey, regBody);
    await app.request('/v1/keys/register', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json', 'x-forwarded-for': '203.0.113.2' },
      body: JSON.stringify(signedReg),
    });

    // Burn through the delete limit
    for (let i = 0; i < 5; i++) {
      const delBody = await signPayload(crypto, identity.signingPrivateKey, { address: 'ratelimit' });
      await app.request('/v1/keys/ratelimit', {
        method: 'DELETE',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(delBody),
      });
    }

    // 6th should be rate-limited
    const delBody = await signPayload(crypto, identity.signingPrivateKey, { address: 'ratelimit' });
    const res = await app.request('/v1/keys/ratelimit', {
      method: 'DELETE',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(delBody),
    });
    expect(res.status).toBe(429);
    expect(res.headers.get('Retry-After')).toBeTruthy();
  });
});