|
| 1 | +import { beforeEach, describe, expect, it, vi } from "vitest"; |
| 2 | + |
| 3 | +/** |
| 4 | + * A delegated (agent/PAT-minted) env JWT rotates its token value every turn. Keying the |
| 5 | + * rate limiter on the token would hand each turn a fresh bucket, so the limiter keys on |
| 6 | + * env+acting-user instead — stable across turns, namespaced away from PRIVATE-key buckets. |
| 7 | + */ |
| 8 | + |
| 9 | +const mocks = vi.hoisted(() => ({ |
| 10 | + authenticateAuthorizationHeader: vi.fn<(...args: any[]) => Promise<any>>(), |
| 11 | + resolvePrivateApiKeyRateLimitScope: vi.fn<(...args: any[]) => Promise<any>>(), |
| 12 | +})); |
| 13 | + |
| 14 | +// Importing the module constructs the real middleware at load; stub the constructor and env |
| 15 | +// so the test doesn't reach for redis or the env contract. |
| 16 | +vi.mock("~/services/authorizationRateLimitMiddleware.server", () => ({ |
| 17 | + authorizationRateLimitMiddleware: () => (_req: any, _res: any, next: any) => next(), |
| 18 | +})); |
| 19 | +vi.mock("~/env.server", () => ({ |
| 20 | + env: { API_RATE_LIMIT_JWT_WINDOW: "1m", API_RATE_LIMIT_JWT_TOKENS: 100 }, |
| 21 | +})); |
| 22 | +vi.mock("~/models/runtimeEnvironment.server", () => ({ |
| 23 | + resolvePrivateApiKeyRateLimitScope: mocks.resolvePrivateApiKeyRateLimitScope, |
| 24 | +})); |
| 25 | +vi.mock("~/runEngine/concerns/batchStreamGrantsInstance.server", () => ({ |
| 26 | + batchStreamGrants: { spend: vi.fn() }, |
| 27 | +})); |
| 28 | +vi.mock("~/services/apiAuth.server", () => ({ |
| 29 | + authenticateAuthorizationHeader: mocks.authenticateAuthorizationHeader, |
| 30 | +})); |
| 31 | + |
| 32 | +import { |
| 33 | + jwtActorRateLimitIdentifier, |
| 34 | + resolveApiRateLimitOverride, |
| 35 | +} from "~/services/apiRateLimit.server"; |
| 36 | + |
| 37 | +describe("jwtActorRateLimitIdentifier", () => { |
| 38 | + it("is stable across token value — depends only on env + acting user", () => { |
| 39 | + // Two turns of the same agent: different JWTs, same env + act.sub. |
| 40 | + const first = jwtActorRateLimitIdentifier("env_123", "usr_abc"); |
| 41 | + const second = jwtActorRateLimitIdentifier("env_123", "usr_abc"); |
| 42 | + |
| 43 | + expect(first).toBe(second); |
| 44 | + expect(first).toBe("jwt-actor:env_123:usr_abc"); |
| 45 | + }); |
| 46 | + |
| 47 | + it("is namespaced so it can't collide with a PRIVATE-key bucket (bare env id)", () => { |
| 48 | + const identifier = jwtActorRateLimitIdentifier("env_123", "usr_abc"); |
| 49 | + |
| 50 | + expect(identifier.startsWith("jwt-actor:")).toBe(true); |
| 51 | + expect(identifier).not.toBe("env_123"); |
| 52 | + }); |
| 53 | + |
| 54 | + it("is compound: a different user in the same env gets a different bucket", () => { |
| 55 | + expect(jwtActorRateLimitIdentifier("env_123", "usr_abc")).not.toBe( |
| 56 | + jwtActorRateLimitIdentifier("env_123", "usr_xyz") |
| 57 | + ); |
| 58 | + expect(jwtActorRateLimitIdentifier("env_123", "usr_abc")).not.toBe( |
| 59 | + jwtActorRateLimitIdentifier("env_999", "usr_abc") |
| 60 | + ); |
| 61 | + }); |
| 62 | +}); |
| 63 | + |
| 64 | +/** |
| 65 | + * The PUBLIC_JWT branch of the real override. These exercise the `actor?.sub` guard itself — |
| 66 | + * they fail if the guard is deleted or the (environmentId, actor.sub) args are swapped. |
| 67 | + */ |
| 68 | +describe("resolveApiRateLimitOverride — PUBLIC_JWT branch", () => { |
| 69 | + // A JWT bearer isn't `tr_`-prefixed, so it skips the private-key branch and hits auth. |
| 70 | + const JWT_BEARER = "Bearer eyJ.delegated.jwt"; |
| 71 | + |
| 72 | + beforeEach(() => { |
| 73 | + mocks.authenticateAuthorizationHeader.mockReset(); |
| 74 | + mocks.resolvePrivateApiKeyRateLimitScope.mockReset(); |
| 75 | + }); |
| 76 | + |
| 77 | + it("keys a delegated JWT (act.sub present) on jwt-actor:${env}:${sub}", async () => { |
| 78 | + mocks.authenticateAuthorizationHeader.mockResolvedValue({ |
| 79 | + ok: true, |
| 80 | + type: "PUBLIC_JWT", |
| 81 | + environment: { id: "env_777" }, |
| 82 | + actor: { sub: "usr_555" }, |
| 83 | + }); |
| 84 | + |
| 85 | + const override = await resolveApiRateLimitOverride(JWT_BEARER); |
| 86 | + |
| 87 | + expect(override?.identifier).toBe("jwt-actor:env_777:usr_555"); |
| 88 | + expect(override?.config).toBeDefined(); |
| 89 | + }); |
| 90 | + |
| 91 | + it("leaves a realtime JWT (no act) on the hashed-token fallback (no identifier)", async () => { |
| 92 | + mocks.authenticateAuthorizationHeader.mockResolvedValue({ |
| 93 | + ok: true, |
| 94 | + type: "PUBLIC_JWT", |
| 95 | + environment: { id: "env_777" }, |
| 96 | + // no `actor` |
| 97 | + }); |
| 98 | + |
| 99 | + const override = await resolveApiRateLimitOverride(JWT_BEARER); |
| 100 | + |
| 101 | + expect(override?.identifier).toBeUndefined(); |
| 102 | + expect(override?.config).toBeDefined(); |
| 103 | + }); |
| 104 | +}); |
0 commit comments