orca/cloud/apps/push/src/client-ip-rate-limit.test.ts

import { PUSH_LIMITS } from '@orca-cloud/push-contract'
import { Hono } from 'hono'
import { describe, expect, it, vi } from 'vitest'
import { ClientIpRateLimiter, clientIpRateLimit } from './client-ip-rate-limit.js'

const CAPACITY = PUSH_LIMITS.unauthenticatedRequestsPerMinutePerIp

function limiterApp(limiter: ClientIpRateLimiter, trustedProxyHops = 0): Hono {
  const app = new Hono()
  app.post('/probe', clientIpRateLimit(limiter, { trustedProxyHops }), (context) =>
    context.json({ ok: true })
  )
  return app
}

describe('client ip rate limiter', () => {
  it('admits exactly the per-minute allowance and refuses the next request', () => {
    const limiter = new ClientIpRateLimiter({ now: () => 1_000 })
    for (let index = 0; index < CAPACITY; index++) {
      expect(limiter.allow('203.0.113.7')).toBe(true)
    }
    expect(limiter.allow('203.0.113.7')).toBe(false)
  })

  it('keeps one client ip from spending another one budget', () => {
    const limiter = new ClientIpRateLimiter({ now: () => 1_000 })
    for (let index = 0; index < CAPACITY; index++) limiter.allow('203.0.113.7')
    expect(limiter.allow('203.0.113.7')).toBe(false)
    expect(limiter.allow('198.51.100.9')).toBe(true)
  })

  it('refills over the window rather than resetting on a boundary', () => {
    let clock = 1_000
    const limiter = new ClientIpRateLimiter({ now: () => clock })
    for (let index = 0; index < CAPACITY; index++) limiter.allow('203.0.113.7')
    expect(limiter.allow('203.0.113.7')).toBe(false)

    // Half a window buys back half the allowance, no more.
    clock += 30_000
    for (let index = 0; index < CAPACITY / 2; index++) {
      expect(limiter.allow('203.0.113.7')).toBe(true)
    }
    expect(limiter.allow('203.0.113.7')).toBe(false)
  })

  it('bounds what it remembers when a flood of distinct ips arrives', () => {
    let clock = 1_000
    const limiter = new ClientIpRateLimiter({ now: () => clock, maxTrackedIps: 8 })
    for (let index = 0; index < 200; index++) {
      clock += 1
      limiter.allow(`198.51.100.${index}`)
    }
    expect(limiter.trackedIpCount()).toBeLessThanOrEqual(8)
  })

  it('evicts the least recently used bucket without scanning the map', () => {
    const limiter = new ClientIpRateLimiter({ capacity: 1, maxTrackedIps: 2, now: () => 1_000 })
    limiter.allow('old')
    limiter.allow('recent')
    expect(limiter.allow('old')).toBe(false)
    const entries = vi.spyOn(Map.prototype, 'entries')
    const iterator = vi.spyOn(Map.prototype, Symbol.iterator)
    try {
      limiter.allow('new')
      expect(entries).not.toHaveBeenCalled()
      expect(iterator).not.toHaveBeenCalled()
    } finally {
      entries.mockRestore()
      iterator.mockRestore()
    }
    expect(limiter.available('old')).toBe(false)
    expect(limiter.available('recent')).toBe(true)
    expect(limiter.trackedIpCount()).toBe(2)
  })

  it('answers 429 with a rate_limited body once the bucket is empty', async () => {
    const app = limiterApp(new ClientIpRateLimiter({ now: () => 1_000 }))
    const headers = { 'x-forwarded-for': '10.0.0.1, 10.0.0.2, 203.0.113.7' }
    for (let index = 0; index < CAPACITY; index++) {
      expect((await app.request('/probe', { method: 'POST', headers })).status).toBe(200)
    }
    const limited = await app.request('/probe', { method: 'POST', headers })
    expect(limited.status).toBe(429)
    expect(await limited.json()).toEqual({ error: 'rate_limited' })
  })

  it('buckets on the last forwarded hop, the only one the platform appended', async () => {
    const app = limiterApp(new ClientIpRateLimiter({ now: () => 1_000 }))
    for (let index = 0; index < CAPACITY; index++) {
      await app.request('/probe', {
        method: 'POST',
        headers: { 'x-forwarded-for': `10.0.0.${index}, 203.0.113.7` }
      })
    }
    const sameClient = await app.request('/probe', {
      method: 'POST',
      headers: { 'x-forwarded-for': '10.9.9.9, 203.0.113.7' }
    })
    expect(sameClient.status).toBe(429)
    const otherClient = await app.request('/probe', {
      method: 'POST',
      headers: { 'x-forwarded-for': '10.0.0.1, 198.51.100.9' }
    })
    expect(otherClient.status).toBe(200)
  })

  it('gives a spoofed left-most hop no escape from the caller own bucket', async () => {
    const app = limiterApp(new ClientIpRateLimiter({ now: () => 1_000 }))
    // A caller that rewrites its own x-forwarded-for on every request still ends
    // up behind the one value Cloud Run appended.
    for (let index = 0; index < CAPACITY; index++) {
      const allowed = await app.request('/probe', {
        method: 'POST',
        headers: { 'x-forwarded-for': `198.51.100.${index}, 203.0.113.7` }
      })
      expect(allowed.status).toBe(200)
    }
    const spoofed = await app.request('/probe', {
      method: 'POST',
      headers: { 'x-forwarded-for': '198.51.100.250, 10.1.1.1, 203.0.113.7' }
    })
    expect(spoofed.status).toBe(429)
  })

  it('skips the configured trusted proxies when counting from the right', async () => {
    const app = limiterApp(new ClientIpRateLimiter({ now: () => 1_000, capacity: 1 }), 1)
    // <client>, <cloud run>, <load balancer>: one trusted hop after the client.
    const headers = { 'x-forwarded-for': '203.0.113.7, 10.0.0.1' }
    expect((await app.request('/probe', { method: 'POST', headers })).status).toBe(200)
    expect((await app.request('/probe', { method: 'POST', headers })).status).toBe(429)
    expect(
      (
        await app.request('/probe', {
          method: 'POST',
          headers: { 'x-forwarded-for': '198.51.100.9, 10.0.0.1' }
        })
      ).status
    ).toBe(200)
  })

  it('trusts nothing when the header is shorter than the configured depth', async () => {
    const app = limiterApp(new ClientIpRateLimiter({ now: () => 1_000, capacity: 1 }), 1)
    // Only one hop, so the client value the depth points at does not exist.
    const headers = { 'x-forwarded-for': '203.0.113.7' }
    expect((await app.request('/probe', { method: 'POST', headers })).status).toBe(200)
    expect(
      (
        await app.request('/probe', {
          method: 'POST',
          headers: { 'x-forwarded-for': '198.51.100.9' }
        })
      ).status
    ).toBe(429)
  })

  it('ignores spoofable x-real-ip and uses a single shared bucket', async () => {
    const app = limiterApp(new ClientIpRateLimiter({ now: () => 1_000, capacity: 1 }))
    expect(
      (await app.request('/probe', { method: 'POST', headers: { 'x-real-ip': '198.51.100.9' } }))
        .status
    ).toBe(200)
    expect(
      (await app.request('/probe', { method: 'POST', headers: { 'x-real-ip': '203.0.113.7' } }))
        .status
    ).toBe(429)
    expect((await app.request('/probe', { method: 'POST' })).status).toBe(429)
    expect((await app.request('/probe', { method: 'POST' })).status).toBe(429)
  })
})