From 8270c8a0b4f523b4554d0fbd49667a7e05c8bea6 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 1 Sep 2026 13:50:42 -0500 Subject: [PATCH 1/3] feat(auth): enable audience-aware worker readers in batch one --- packages/worker-utils/src/index.ts | 1 + .../worker-utils/src/kilo-token-auth.test.ts | 298 ++++++++++++++++++ packages/worker-utils/src/kilo-token-auth.ts | 46 ++- .../src/kilo-token-policy.test.ts | 149 +++++++++ .../worker-utils/src/kilo-token-policy.ts | 15 + packages/worker-utils/src/kilo-token.test.ts | 18 ++ pnpm-lock.yaml | 3 + services/ai-attribution/jest.config.mjs | 10 +- services/ai-attribution/package.json | 1 + services/ai-attribution/src/util/auth.test.ts | 225 +++++++++++++ services/ai-attribution/src/util/auth.ts | 9 +- .../middleware/kilo-auth.middleware.test.ts | 228 +++++++++++--- .../src/middleware/kilo-auth.middleware.ts | 9 +- .../middleware/kilo-auth.middleware.test.ts | 197 ++++++++++++ .../src/middleware/kilo-auth.middleware.ts | 9 +- 15 files changed, 1145 insertions(+), 73 deletions(-) create mode 100644 services/ai-attribution/src/util/auth.test.ts create mode 100644 services/wasteland/src/middleware/kilo-auth.middleware.test.ts diff --git a/packages/worker-utils/src/index.ts b/packages/worker-utils/src/index.ts index 816a67d77f..515809b87d 100644 --- a/packages/worker-utils/src/index.ts +++ b/packages/worker-utils/src/index.ts @@ -144,6 +144,7 @@ export { isKiloCredentialExchangeEligible, isKiloResourceAudienceAllowed, verifyKiloSessionForPolicy, + verifyKiloTokenForResource, verifyKiloTokenForPolicy, } from './kilo-token-policy.js'; export type { diff --git a/packages/worker-utils/src/kilo-token-auth.test.ts b/packages/worker-utils/src/kilo-token-auth.test.ts index 539450e6d6..a2fcd72c6d 100644 --- a/packages/worker-utils/src/kilo-token-auth.test.ts +++ b/packages/worker-utils/src/kilo-token-auth.test.ts @@ -183,6 +183,304 @@ describe('verifyKiloBearerAgainstCurrentPepper', () => { }) ).resolves.toBeNull(); }); + + it('uses the existing verifier strictly by default before account lookup', async () => { + const { token } = await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'resource-audience', + }); + let lookupCount = 0; + + await expect( + verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: TEST_JWT_SECRET, + workerEnv: 'production', + connectionString: 'postgres://test', + getUserPepper: async (...args) => { + lookupCount++; + return getUserPepper(...args); + }, + }) + ).resolves.toBeNull(); + expect(lookupCount).toBe(0); + }); + + it('dispatches the existing verifier audience option for string and array claims', async () => { + const now = Math.floor(Date.now() / 1000); + const arrayAudienceToken = await new SignJWT({ + version: 3, + kiloUserId: 'user-xyz-789', + apiTokenPepper: 'pepper-current', + env: 'production', + }) + .setProtectedHeader({ alg: 'HS256' }) + .setAudience(['other-resource', 'resource-audience']) + .setIssuedAt(now) + .setExpirationTime(now + 3600) + .sign(new TextEncoder().encode(TEST_JWT_SECRET)); + const matchingString = await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'resource-audience', + }); + const absentAudience = await signToken({ pepper: 'pepper-current', tokenSource: 'kilo-chat' }); + const mismatchedAudience = await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'other-resource', + }); + const params = { + nextAuthSecret: TEST_JWT_SECRET, + workerEnv: 'production', + connectionString: 'postgres://test', + audience: 'resource-audience', + } as const; + + for (const testCase of [ + { token: matchingString.token, result: { userId: 'user-xyz-789' } }, + { token: arrayAudienceToken, result: { userId: 'user-xyz-789' } }, + { token: absentAudience.token, result: null }, + { token: mismatchedAudience.token, result: null }, + ]) { + let lookupCount = 0; + await expect( + verifyKiloBearerAgainstCurrentPepper({ + ...params, + token: testCase.token, + getUserPepper: async (...args) => { + lookupCount++; + return getUserPepper(...args); + }, + }) + ).resolves.toEqual(testCase.result); + expect(lookupCount).toBe(testCase.result === null ? 0 : 1); + } + }); + + it('applies required and legacy resource audience policies before account lookup', async () => { + const { token } = await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + }); + const params = { + token, + nextAuthSecret: TEST_JWT_SECRET, + workerEnv: 'production', + connectionString: 'postgres://test', + }; + + await expect( + verifyKiloBearerAgainstCurrentPepper({ + ...params, + getUserPepper, + resourceAudience: { audience: 'resource-audience', mode: 'allow-legacy' }, + }) + ).resolves.toEqual({ userId: 'user-xyz-789' }); + let lookupCount = 0; + await expect( + verifyKiloBearerAgainstCurrentPepper({ + ...params, + getUserPepper: async (...args) => { + lookupCount++; + return getUserPepper(...args); + }, + resourceAudience: { audience: 'resource-audience', mode: 'required' }, + }) + ).resolves.toBeNull(); + expect(lookupCount).toBe(0); + }); + + it('compares present resource-token pepper claims against a non-null stored pepper', async () => { + const tokens = await Promise.all([ + signKiloToken({ + userId: 'user-xyz-789', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'resource-audience', + }), + signKiloToken({ + userId: 'user-xyz-789', + pepper: null, + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'resource-audience', + }), + signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-stale', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'resource-audience', + }), + signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'resource-audience', + }), + ]); + const params = { + nextAuthSecret: TEST_JWT_SECRET, + workerEnv: 'production', + connectionString: 'postgres://test', + getUserPepper, + resourceAudience: { audience: 'resource-audience', mode: 'required' } as const, + }; + + for (const [token, result] of [ + [tokens[0]!.token, { userId: 'user-xyz-789' }], + [tokens[1]!.token, null], + [tokens[2]!.token, null], + [tokens[3]!.token, { userId: 'user-xyz-789' }], + ] as const) { + await expect(verifyKiloBearerAgainstCurrentPepper({ ...params, token })).resolves.toEqual( + result + ); + } + }); + + it('denies missing and blocked resource-token accounts and never bypasses pepper checks', async () => { + const { token } = await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-stale', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'resource-audience', + }); + const params = { + token, + nextAuthSecret: TEST_JWT_SECRET, + workerEnv: 'production', + connectionString: 'postgres://test', + resourceAudience: { audience: 'resource-audience', mode: 'required' } as const, + }; + + userResultByUserId.clear(); + await expect( + verifyKiloBearerAgainstCurrentPepper({ ...params, getUserPepper }) + ).resolves.toBeNull(); + userResultByUserId.set('user-xyz-789', { + pepper: 'pepper-current', + blockedReason: 'manual block', + }); + await expect( + verifyKiloBearerAgainstCurrentPepper({ ...params, getUserPepper, allowBlocked: true }) + ).resolves.toBeNull(); + await expect( + verifyKiloBearerAgainstCurrentPepper({ ...params, getUserPepper }) + ).resolves.toBeNull(); + }); + + it('denies resource tokens with missing or mismatched env values when workerEnv is set', async () => { + const tokens = await Promise.all([ + signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + audience: 'resource-audience', + }), + signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'staging', + audience: 'resource-audience', + }), + ]); + for (const { token } of tokens) { + await expect( + verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: TEST_JWT_SECRET, + workerEnv: 'production', + connectionString: 'postgres://test', + getUserPepper, + resourceAudience: { audience: 'resource-audience', mode: 'required' }, + }) + ).resolves.toBeNull(); + } + }); + + it('propagates resource secret-provider failures without an account lookup', async () => { + const { token } = await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'resource-audience', + }); + let lookupCount = 0; + + await expect( + verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: { get: async () => Promise.reject(new Error('secrets store unavailable')) }, + workerEnv: 'production', + connectionString: 'postgres://test', + getUserPepper: async (...args) => { + lookupCount++; + return getUserPepper(...args); + }, + resourceAudience: { audience: 'resource-audience', mode: 'required' }, + }) + ).rejects.toThrow('secrets store unavailable'); + expect(lookupCount).toBe(0); + }); + + it('propagates account lookup failures after resource verification', async () => { + const { token } = await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + audience: 'resource-audience', + }); + + await expect( + verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: TEST_JWT_SECRET, + connectionString: 'postgres://test', + getUserPepper: async () => Promise.reject(new Error('connection refused')), + resourceAudience: { audience: 'resource-audience', mode: 'required' }, + }) + ).rejects.toThrow('connection refused'); + }); + + it('rejects mutually configured audience policies explicitly', async () => { + await expect( + verifyKiloBearerAgainstCurrentPepper({ + token: null, + nextAuthSecret: TEST_JWT_SECRET, + connectionString: 'postgres://test', + audience: 'legacy-audience', + resourceAudience: { audience: 'resource-audience', mode: 'required' }, + } as never) + ).rejects.toThrow('mutually exclusive'); + }); }); describe('internal service tokens (no apiTokenPepper, no env)', () => { diff --git a/packages/worker-utils/src/kilo-token-auth.ts b/packages/worker-utils/src/kilo-token-auth.ts index ca781ac541..9dae91252f 100644 --- a/packages/worker-utils/src/kilo-token-auth.ts +++ b/packages/worker-utils/src/kilo-token-auth.ts @@ -4,6 +4,7 @@ import { eq } from 'drizzle-orm'; import { getCachedSecret } from './cached-secret'; import { verifyKiloToken } from './kilo-token'; +import { verifyKiloTokenForResource, type KiloResourceAudiencePolicy } from './kilo-token-policy'; export type KiloBearerAuthResult = { userId: string; @@ -21,6 +22,16 @@ export type GetKiloUserPepper = ( userId: string ) => Promise; +type KiloBearerAudienceOptions = + | { + audience?: string; + resourceAudience?: never; + } + | { + audience?: never; + resourceAudience?: KiloResourceAudiencePolicy; + }; + export async function findKiloUserPepper( connectionString: string, userId: string @@ -47,15 +58,19 @@ export async function findKiloUserPepper( * A dependency failure (secret store, database) throws, so a caller can map an * outage to a retryable 503 instead of reporting it as an invalid token. */ -export async function verifyKiloBearerAgainstCurrentPepper(params: { - token: string | null; - nextAuthSecret: KiloSecretBinding | string; - workerEnv?: string; - connectionString: string; - getUserPepper?: GetKiloUserPepper; - audience?: string; - allowBlocked?: boolean; -}): Promise { +export async function verifyKiloBearerAgainstCurrentPepper( + params: { + token: string | null; + nextAuthSecret: KiloSecretBinding | string; + workerEnv?: string; + connectionString: string; + getUserPepper?: GetKiloUserPepper; + allowBlocked?: boolean; + } & KiloBearerAudienceOptions +): Promise { + if (params.audience !== undefined && params.resourceAudience !== undefined) { + throw new Error('Bearer audience and resource audience policies are mutually exclusive'); + } if (!params.token) return null; const getUserPepper = params.getUserPepper ?? findKiloUserPepper; @@ -67,11 +82,14 @@ export async function verifyKiloBearerAgainstCurrentPepper(params: { let payload: Awaited>; try { - payload = await verifyKiloToken( - params.token, - secret, - params.audience ? { audience: params.audience } : undefined - ); + payload = + params.resourceAudience === undefined + ? await verifyKiloToken( + params.token, + secret, + params.audience ? { audience: params.audience } : undefined + ) + : await verifyKiloTokenForResource(params.token, secret, params.resourceAudience); } catch { return null; } diff --git a/packages/worker-utils/src/kilo-token-policy.test.ts b/packages/worker-utils/src/kilo-token-policy.test.ts index 086309c0ab..4f96441319 100644 --- a/packages/worker-utils/src/kilo-token-policy.test.ts +++ b/packages/worker-utils/src/kilo-token-policy.test.ts @@ -6,6 +6,7 @@ import { isKiloCredentialExchangeEligible, isKiloResourceAudienceAllowed, verifyKiloSessionForPolicy, + verifyKiloTokenForResource, verifyKiloTokenForPolicy, type ModernKiloTokenClaims, type VerifiedKiloAuthContext, @@ -323,6 +324,154 @@ describe('verifyKiloTokenForPolicy', () => { }); }); +describe('verifyKiloTokenForResource', () => { + it.each([ + ['matching audience', legacyClaims({ aud: 'kilo-api' }), API_POLICY, true], + [ + 'matching audience array', + legacyClaims({ aud: ['kilo-gateway', 'kilo-api'] }), + API_POLICY, + true, + ], + ['missing required audience', legacyClaims(), API_POLICY, false], + ['missing legacy audience', legacyClaims(), LEGACY_POLICY, true], + ['mismatched audience', legacyClaims({ aud: 'kilo-gateway' }), API_POLICY, false], + ['null audience', legacyClaims({ aud: null }), API_POLICY, false], + ['empty audience', legacyClaims({ aud: '' }), API_POLICY, false], + ['numeric audience', legacyClaims({ aud: 1 }), API_POLICY, false], + ['empty audience array', legacyClaims({ aud: [] }), API_POLICY, false], + [ + 'duplicate audience members', + legacyClaims({ aud: ['kilo-api', 'kilo-api'] }), + API_POLICY, + false, + ], + [ + 'trimmed audience member', + legacyClaims({ aud: ['kilo-api', ' kilo-gateway'] }), + API_POLICY, + false, + ], + ['numeric audience member', legacyClaims({ aud: ['kilo-api', 1] }), API_POLICY, false], + ])('handles %s', async (_name, claims, policy, allowed) => { + vi.useFakeTimers(); + vi.setSystemTime(NOW); + const verification = verifyKiloTokenForResource(await sign(claims), SECRET, policy); + if (allowed) { + await expect(verification).resolves.toMatchObject({ kiloUserId: 'synthetic-user' }); + } else { + await expect(verification).rejects.toThrow(); + } + }); + + it.each(['', ' kilo-api', 'kilo-api '])( + 'fails closed for invalid configured audience %j', + async audience => { + vi.useFakeTimers(); + vi.setSystemTime(NOW); + await expect( + verifyKiloTokenForResource(await sign(legacyClaims({ aud: 'kilo-api' })), SECRET, { + audience, + mode: 'required', + }) + ).rejects.toThrow(); + } + ); + + it('preserves known legacy claims and strips ignored unknown claims', async () => { + vi.useFakeTimers(); + vi.setSystemTime(NOW); + const payload = await verifyKiloTokenForResource( + await sign( + legacyClaims({ + aud: 'kilo-api', + apiTokenPepper: null, + env: 'production', + botId: 'synthetic-bot', + tokenPurpose: 'human-api', + credentialExchange: true, + ignoredFutureResourceClaim: true, + }) + ), + SECRET, + API_POLICY + ); + + expect(payload).toMatchObject({ + kiloUserId: 'synthetic-user', + apiTokenPepper: null, + env: 'production', + botId: 'synthetic-bot', + }); + expect(payload).not.toHaveProperty('ignoredFutureResourceClaim'); + expect(payload).not.toHaveProperty('tokenPurpose'); + expect(payload).not.toHaveProperty('credentialExchange'); + }); + + it('preserves optional dates and jose temporal behavior', async () => { + vi.useFakeTimers(); + vi.setSystemTime(NOW); + const undated = await verifyKiloTokenForResource( + await sign({ version: 3, kiloUserId: 'synthetic-user', aud: 'kilo-api' }, { dates: false }), + SECRET, + API_POLICY + ); + expect(undated).not.toHaveProperty('iat'); + expect(undated).not.toHaveProperty('exp'); + await expect( + verifyKiloTokenForResource( + await sign(legacyClaims({ aud: 'kilo-api', iat: NOW_SECONDS + 1, exp: NOW_SECONDS + 10 })), + SECRET, + API_POLICY + ) + ).resolves.toMatchObject({ iat: NOW_SECONDS + 1 }); + await expect( + verifyKiloTokenForResource( + await sign(legacyClaims({ aud: 'kilo-api', exp: NOW_SECONDS + 1.5 })), + SECRET, + API_POLICY + ) + ).resolves.toMatchObject({ exp: NOW_SECONDS + 1.5 }); + }); + + it('rejects malformed, expired, not-before, signature, version, and user failures', async () => { + vi.useFakeTimers(); + vi.setSystemTime(NOW); + await expect(verifyKiloTokenForResource('not.a.token', SECRET, API_POLICY)).rejects.toThrow(); + await expect( + verifyKiloTokenForResource( + await sign(legacyClaims({ aud: 'kilo-api', exp: NOW_SECONDS - 1 })), + SECRET, + API_POLICY + ) + ).rejects.toThrow(); + await expect( + verifyKiloTokenForResource( + await sign(legacyClaims({ aud: 'kilo-api' }), { nbf: NOW_SECONDS + 1 }), + SECRET, + API_POLICY + ) + ).rejects.toThrow(); + await expect( + verifyKiloTokenForResource( + await sign(legacyClaims({ aud: 'kilo-api' }), { secret: 'another-secret' }), + SECRET, + API_POLICY + ) + ).rejects.toThrow(); + await expect( + verifyKiloTokenForResource( + await sign({ version: 2, kiloUserId: 'synthetic-user', aud: 'kilo-api' }), + SECRET, + API_POLICY + ) + ).rejects.toThrow(); + await expect( + verifyKiloTokenForResource(await sign({ version: 3, aud: 'kilo-api' }), SECRET, API_POLICY) + ).rejects.toThrow(); + }); +}); + describe('isKiloCredentialExchangeEligible', () => { it('only accepts authentic verified session contexts', async () => { const session = await verifyKiloSessionForPolicy(async () => ({ diff --git a/packages/worker-utils/src/kilo-token-policy.ts b/packages/worker-utils/src/kilo-token-policy.ts index 142245c25d..627930570f 100644 --- a/packages/worker-utils/src/kilo-token-policy.ts +++ b/packages/worker-utils/src/kilo-token-policy.ts @@ -2,6 +2,7 @@ import { jwtVerify } from 'jose'; import { z } from 'zod'; import { KILO_API_AUDIENCE } from './internal-service-token-audiences'; import { KILO_TOKEN_VERSION, kiloTokenPayload } from './kilo-token'; +import type { KiloTokenPayload } from './kilo-token'; const audienceName = z .string() @@ -76,6 +77,20 @@ export function isKiloResourceAudienceAllowed( : parsed.data.includes(policy.audience); } +export async function verifyKiloTokenForResource( + token: string, + secret: string, + audiencePolicy: KiloResourceAudiencePolicy +): Promise { + const { payload } = await jwtVerify(token, new TextEncoder().encode(secret), { + algorithms: ['HS256'], + }); + if (!isKiloResourceAudienceAllowed(payload.aud, audiencePolicy)) { + throw new Error('Unexpected token audience'); + } + return kiloTokenPayload.parse(payload); +} + const verifiedAuth = Symbol('verified-kilo-policy-auth'); type DeepReadonly = { readonly [K in keyof T]: DeepReadonly }; diff --git a/packages/worker-utils/src/kilo-token.test.ts b/packages/worker-utils/src/kilo-token.test.ts index 519aebb8d6..daf6525638 100644 --- a/packages/worker-utils/src/kilo-token.test.ts +++ b/packages/worker-utils/src/kilo-token.test.ts @@ -200,6 +200,24 @@ describe('verifyKiloToken', () => { await expect(verifyKiloToken(token, SECRET)).rejects.toThrow(); }); + it('retains jose audience-array handling for the strict reader', async () => { + const token = await new SignJWT({ + version: 3, + kiloUserId: 'user-123', + aud: ['git-token-service:bitbucket-repositories', 'another-service'], + }) + .setProtectedHeader({ alg: 'HS256' }) + .setIssuedAt() + .setExpirationTime('1h') + .sign(encode(SECRET)); + + await expect( + verifyKiloToken(token, SECRET, { audience: 'git-token-service:bitbucket-repositories' }) + ).resolves.toMatchObject({ kiloUserId: 'user-123' }); + await expect(verifyKiloToken(token, SECRET, { audience: 'other-service' })).rejects.toThrow(); + await expect(verifyKiloToken(token, SECRET)).rejects.toThrow('Unexpected token audience'); + }); + it('rejects wrong secret', async () => { const token = await sign({ version: 3, kiloUserId: 'user-123' }); await expect( diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index c1e89a7818..460a67b098 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1762,6 +1762,9 @@ importers: jest: specifier: 29.7.0 version: 29.7.0(@types/node@24.12.4)(node-notifier@10.0.1) + ts-jest: + specifier: 29.4.9 + version: 29.4.9(@babel/core@7.29.7)(@jest/transform@30.3.0)(@jest/types@30.3.0)(babel-jest@30.3.0(@babel/core@7.29.7))(jest-util@30.3.0)(jest@29.7.0(@types/node@24.12.4)(node-notifier@10.0.1))(typescript@5.9.3) typescript: specifier: 'catalog:' version: 5.9.3 diff --git a/services/ai-attribution/jest.config.mjs b/services/ai-attribution/jest.config.mjs index 6c7712724b..49737a6796 100644 --- a/services/ai-attribution/jest.config.mjs +++ b/services/ai-attribution/jest.config.mjs @@ -3,15 +3,21 @@ export default { preset: 'ts-jest', testEnvironment: 'node', testMatch: ['**/*.test.ts'], + moduleNameMapper: { + '^(\\.{1,2}/.*)\\.js$': '$1', + }, + transformIgnorePatterns: ['/node_modules/(?!.*(?:jose|workers-tagged-logger))'], transform: { - '^.+\\.ts$': [ + '^.+\\.[tj]s$': [ 'ts-jest', { tsconfig: { esModuleInterop: true, allowSyntheticDefaultImports: true, + allowJs: true, + module: 'CommonJS', + moduleResolution: 'Node', }, - useESM: true, }, ], }, diff --git a/services/ai-attribution/package.json b/services/ai-attribution/package.json index 156331bb23..dd65d7c038 100644 --- a/services/ai-attribution/package.json +++ b/services/ai-attribution/package.json @@ -24,6 +24,7 @@ "@typescript/native-preview": "catalog:", "drizzle-kit": "catalog:", "jest": "29.7.0", + "ts-jest": "29.4.9", "typescript": "catalog:", "wrangler": "catalog:" }, diff --git a/services/ai-attribution/src/util/auth.test.ts b/services/ai-attribution/src/util/auth.test.ts new file mode 100644 index 0000000000..462159257a --- /dev/null +++ b/services/ai-attribution/src/util/auth.test.ts @@ -0,0 +1,225 @@ +import { createHmac } from 'node:crypto'; +import { Hono } from 'hono'; +import { signKiloToken } from '@kilocode/worker-utils/kilo-token'; +import { AI_ATTRIBUTION_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import type { HonoContext } from '../ai-attribution.worker'; +import { authMiddleware, validateKiloToken } from './auth'; + +const TEST_SECRET = 'ai-attribution-test-secret'; +const USER_ID = 'user-123'; +const ORGANIZATION_ID = 'org-123'; + +function base64Url(value: string): string { + return Buffer.from(value).toString('base64url'); +} + +function signPayload(payload: Record, secret = TEST_SECRET): string { + const header = base64Url(JSON.stringify({ alg: 'HS256', typ: 'JWT' })); + const body = base64Url(JSON.stringify(payload)); + const signature = createHmac('sha256', secret).update(`${header}.${body}`).digest('base64url'); + return `${header}.${body}.${signature}`; +} + +function policyPayload(overrides: Record = {}): Record { + const now = Math.floor(Date.now() / 1000); + return { + version: 3, + kiloUserId: USER_ID, + organizationId: ORGANIZATION_ID, + organizationRole: 'member', + iat: now - 60, + exp: now + 60, + ...overrides, + }; +} + +async function organizationToken( + options: { + audience?: string; + organizationId?: string; + organizationRole?: 'owner' | 'member' | 'billing_manager'; + expiresInSeconds?: number; + } = {} +): Promise { + const { token } = await signKiloToken({ + userId: USER_ID, + pepper: null, + secret: TEST_SECRET, + expiresInSeconds: options.expiresInSeconds ?? 60, + audience: options.audience, + extra: { + organizationId: options.organizationId ?? ORGANIZATION_ID, + organizationRole: options.organizationRole ?? 'member', + }, + }); + return token; +} + +function buildApp() { + let downstreamCalls = 0; + const app = new Hono(); + app.use('*', authMiddleware); + app.get('*', c => { + downstreamCalls += 1; + return c.json({ + reached: true, + userId: c.get('user_id'), + organizationId: c.get('organization_id'), + organizationRole: c.get('organization_role'), + token: c.get('token'), + }); + }); + return { app, downstreamCalls: () => downstreamCalls }; +} + +describe('validateKiloToken', () => { + it('accepts an audience-bound organization token and preserves the bearer token', async () => { + const token = await organizationToken({ audience: AI_ATTRIBUTION_AUDIENCE }); + + await expect(validateKiloToken(`Bearer ${token}`, TEST_SECRET)).resolves.toEqual({ + success: true, + token, + kiloUserId: USER_ID, + organizationId: ORGANIZATION_ID, + organizationRole: 'member', + }); + }); + + it('accepts a legacy organization token with no audience', async () => { + const token = await organizationToken(); + + await expect(validateKiloToken(`Bearer ${token}`, TEST_SECRET)).resolves.toMatchObject({ + success: true, + organizationId: ORGANIZATION_ID, + }); + }); + + it('accepts a matching audience in an audience array', async () => { + const token = signPayload(policyPayload({ aud: ['another-service', AI_ATTRIBUTION_AUDIENCE] })); + + await expect(validateKiloToken(`Bearer ${token}`, TEST_SECRET)).resolves.toMatchObject({ + success: true, + token, + }); + }); + + it.each([ + ['mismatched audience', policyPayload({ aud: 'another-service' })], + ['malformed empty audience', policyPayload({ aud: '' })], + [ + 'malformed duplicate audience array', + policyPayload({ aud: [AI_ATTRIBUTION_AUDIENCE, AI_ATTRIBUTION_AUDIENCE] }), + ], + ['missing organization ID', policyPayload({ organizationId: undefined })], + ['billing-manager organization role', policyPayload({ organizationRole: 'billing_manager' })], + ['invalid organization role', policyPayload({ organizationRole: 'admin' })], + ['unsupported token version', policyPayload({ version: 2 })], + ['expired token', policyPayload({ iat: 1, exp: 2 })], + ])('rejects a token with %s', async (_name, payload) => { + const token = signPayload(payload); + + await expect(validateKiloToken(`Bearer ${token}`, TEST_SECRET)).resolves.toEqual({ + success: false, + error: 'Invalid or expired token', + }); + }); + + it('accepts a legacy token without an audience, issue date, or expiry date', async () => { + const token = signPayload(policyPayload({ iat: undefined, exp: undefined })); + + await expect(validateKiloToken(`Bearer ${token}`, TEST_SECRET)).resolves.toMatchObject({ + success: true, + token, + }); + }); + + it('rejects invalid signatures and malformed authorization headers', async () => { + const token = await organizationToken({ audience: AI_ATTRIBUTION_AUDIENCE }); + const invalidSignature = signPayload( + policyPayload({ aud: AI_ATTRIBUTION_AUDIENCE }), + 'other-secret' + ); + + await expect(validateKiloToken(`Bearer ${invalidSignature}`, TEST_SECRET)).resolves.toEqual({ + success: false, + error: 'Invalid or expired token', + }); + await expect(validateKiloToken(token, TEST_SECRET)).resolves.toEqual({ + success: false, + error: 'Missing or malformed Authorization header', + }); + }); +}); + +describe('authMiddleware', () => { + it('preserves server errors from the secret provider without calling downstream', async () => { + const { app, downstreamCalls } = buildApp(); + const token = await organizationToken({ audience: AI_ATTRIBUTION_AUDIENCE }); + const env = { + NEXTAUTH_SECRET: { + get: async (): Promise => { + throw new Error('test secret provider unavailable'); + }, + }, + } as Env; + + const response = await app.request( + '/attributions/whoami', + { + headers: { Authorization: `Bearer ${token}` }, + }, + env + ); + + expect(response.status).toBe(500); + expect(downstreamCalls()).toBe(0); + }); + + it('does not call the downstream handler for an invalid audience', async () => { + const { app, downstreamCalls } = buildApp(); + const token = await organizationToken({ audience: 'another-service' }); + const env = { NEXTAUTH_SECRET: { get: async () => TEST_SECRET } } as Env; + + const response = await app.request( + '/attributions/whoami', + { + headers: { Authorization: `Bearer ${token}` }, + }, + env + ); + + expect(response.status).toBe(401); + expect(downstreamCalls()).toBe(0); + await expect(response.json()).resolves.toEqual({ + success: false, + error: 'Invalid or expired token', + }); + }); + + it('sets the existing user and organization contexts for a valid token', async () => { + const { app, downstreamCalls } = buildApp(); + const token = await organizationToken({ + audience: AI_ATTRIBUTION_AUDIENCE, + organizationRole: 'owner', + }); + const env = { NEXTAUTH_SECRET: { get: async () => TEST_SECRET } } as Env; + + const response = await app.request( + '/attributions/whoami', + { + headers: { Authorization: `Bearer ${token}` }, + }, + env + ); + + expect(response.status).toBe(200); + expect(downstreamCalls()).toBe(1); + await expect(response.json()).resolves.toEqual({ + reached: true, + userId: USER_ID, + organizationId: ORGANIZATION_ID, + organizationRole: 'owner', + token, + }); + }); +}); diff --git a/services/ai-attribution/src/util/auth.ts b/services/ai-attribution/src/util/auth.ts index e00b3ad1ab..b3642fc4d9 100644 --- a/services/ai-attribution/src/util/auth.ts +++ b/services/ai-attribution/src/util/auth.ts @@ -1,5 +1,7 @@ import { createMiddleware } from 'hono/factory'; -import { verifyKiloToken, extractBearerToken } from '@kilocode/worker-utils'; +import { extractBearerToken } from '@kilocode/worker-utils/extract-bearer-token'; +import { AI_ATTRIBUTION_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { verifyKiloTokenForResource } from '@kilocode/worker-utils/kilo-token-policy'; import { logger } from './logger'; import type { HonoContext } from '../ai-attribution.worker'; import { OrganizationJWTPayload } from '../schemas'; @@ -23,7 +25,10 @@ export async function validateKiloToken( } try { - const raw = await verifyKiloToken(token, secret); + const raw = await verifyKiloTokenForResource(token, secret, { + audience: AI_ATTRIBUTION_AUDIENCE, + mode: 'allow-legacy', + }); const payload = OrganizationJWTPayload.parse(raw); return { diff --git a/services/gastown/src/middleware/kilo-auth.middleware.test.ts b/services/gastown/src/middleware/kilo-auth.middleware.test.ts index 1e4991c41f..1b770015d8 100644 --- a/services/gastown/src/middleware/kilo-auth.middleware.test.ts +++ b/services/gastown/src/middleware/kilo-auth.middleware.test.ts @@ -1,75 +1,201 @@ -import { describe, it, expect } from 'vitest'; -import { Hono } from 'hono'; +import { describe, expect, it } from 'vitest'; +import { Hono, type Context } from 'hono'; import { SignJWT } from 'jose'; +import { GASTOWN_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; import { kiloAuthMiddleware } from './kilo-auth.middleware'; import type { GastownEnv } from '../gastown.worker'; const TEST_SECRET = 'test-secret-that-is-long-enough-for-hs256'; function createApp() { + let downstreamCalls = 0; const app = new Hono(); app.use('/api/*', kiloAuthMiddleware); - app.get('/api/whoami', c => { - return c.json({ kiloUserId: c.get('kiloUserId') }); - }); - return app; + app.use('/trpc/*', kiloAuthMiddleware); + const handler = (c: Context) => { + downstreamCalls += 1; + return c.json({ + kiloUserId: c.get('kiloUserId'), + isAdmin: c.get('kiloIsAdmin'), + pepper: c.get('kiloApiTokenPepper'), + gastownAccess: c.get('kiloGastownAccess'), + memberships: c.get('kiloOrgMemberships'), + }); + }; + app.get('/api/whoami', handler); + app.get('/trpc/whoami', handler); + return { app, downstreamCalls: () => downstreamCalls }; } -async function signToken(payload: Record) { +async function signToken( + claims: Record = {}, + options: { secret?: string; dates?: boolean; expiration?: number } = {} +) { const now = Math.floor(Date.now() / 1000); - return new SignJWT(payload) - .setProtectedHeader({ alg: 'HS256' }) - .setIssuedAt(now) - .setExpirationTime(now + 3600) - .sign(new TextEncoder().encode(TEST_SECRET)); + let jwt = new SignJWT({ version: 3, kiloUserId: 'user-abc', ...claims }).setProtectedHeader({ + alg: 'HS256', + }); + if (options.dates !== false) { + jwt = jwt.setIssuedAt(now).setExpirationTime(options.expiration ?? now + 3600); + } + return jwt.sign(new TextEncoder().encode(options.secret ?? TEST_SECRET)); +} + +async function request( + app: Hono, + token: string | undefined, + secret: string | { get(): Promise } | null = TEST_SECRET +) { + return Promise.all( + ['/api/whoami', '/trpc/whoami'].map(path => + app.request( + path, + token === undefined ? {} : { headers: { Authorization: `Bearer ${token}` } }, + { NEXTAUTH_SECRET: secret } as never + ) + ) + ); } describe('kiloAuthMiddleware', () => { - it('rejects when no token is provided', async () => { - const app = createApp(); - const res = await app.request('/api/whoami', {}, { - NEXTAUTH_SECRET: TEST_SECRET, - } as never); - expect(res.status).toBe(401); + it.each([ + ['missing authentication', undefined], + ['malformed authentication', 'Bearer'], + ])('rejects %s before calling downstream', async (_name, authorization) => { + const { app, downstreamCalls } = createApp(); + const responses = + authorization === undefined + ? await request(app, undefined) + : await Promise.all( + ['/api/whoami', '/trpc/whoami'].map(path => + app.request(path, { headers: { Authorization: authorization } }, { + NEXTAUTH_SECRET: TEST_SECRET, + } as never) + ) + ); + + expect(responses.map(response => response.status)).toEqual([401, 401]); + expect(downstreamCalls()).toBe(0); }); - it('accepts a well-formed Kilo token', async () => { - const app = createApp(); - const token = await signToken({ - version: 3, - kiloUserId: 'user-abc', - env: 'development', - }); + it('accepts legacy tokens without an audience or dates', async () => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await signToken({}, { dates: false })); - const res = await app.request( - '/api/whoami', - { headers: { Authorization: `Bearer ${token}` } }, - { NEXTAUTH_SECRET: TEST_SECRET } as never - ); - expect(res.status).toBe(200); - const body = (await res.json()) as { kiloUserId: string }; - expect(body.kiloUserId).toBe('user-abc'); + expect(responses.map(response => response.status)).toEqual([200, 200]); + expect(downstreamCalls()).toBe(2); }); -}); -describe('C15 deviceSessionId compatibility', () => { - it('accepts a token carrying deviceSessionId claim', async () => { - const app = createApp(); - const token = await signToken({ - version: 3, - kiloUserId: 'user-abc', - apiTokenPepper: null, - env: 'development', - deviceSessionId: 'session-gastown-test', - }); + it('accepts its audience as a string or an array', async () => { + for (const aud of [GASTOWN_AUDIENCE, ['kilo-api', GASTOWN_AUDIENCE]]) { + const { app } = createApp(); + const responses = await request(app, await signToken({ aud })); + expect(responses.map(response => response.status)).toEqual([200, 200]); + } + }); + + it.each(['wasteland', 'kilo-api', 'kilo-gateway', 'git-token-service:github-user-access-token'])( + 'rejects a token for %s before calling downstream', + async aud => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await signToken({ aud })); + + expect(responses.map(response => response.status)).toEqual([401, 401]); + expect(downstreamCalls()).toBe(0); + } + ); + + it.each([ + false, + null, + '', + ` ${GASTOWN_AUDIENCE}`, + [GASTOWN_AUDIENCE, ''], + [GASTOWN_AUDIENCE, GASTOWN_AUDIENCE], + ])('rejects malformed explicit audiences before calling downstream', async aud => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await signToken({ aud })); - const res = await app.request( - '/api/whoami', - { headers: { Authorization: `Bearer ${token}` } }, - { NEXTAUTH_SECRET: TEST_SECRET } as never + expect(responses.map(response => response.status)).toEqual([401, 401]); + expect(downstreamCalls()).toBe(0); + }); + + it.each([ + [ + 'invalid signature', + () => signToken({}, { secret: 'different-secret-that-is-long-enough-for-hs256' }), + ], + ['expired token', () => signToken({}, { expiration: Math.floor(Date.now() / 1000) - 60 })], + ['wrong version', () => signToken({ version: 2 })], + ['missing user', () => signToken({ kiloUserId: '' })], + ])('rejects %s before calling downstream', async (_name, token) => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await token()); + + expect(responses.map(response => response.status)).toEqual([401, 401]); + expect(downstreamCalls()).toBe(0); + }); + + it('preserves claim values and defaults', async () => { + const { app } = createApp(); + const responses = await request( + app, + await signToken({ + aud: GASTOWN_AUDIENCE, + isAdmin: true, + apiTokenPepper: 'pepper', + gastownAccess: true, + orgMemberships: [{ orgId: 'org-a', role: 'owner' }], + deviceSessionId: 'session-gastown-test', + }) ); - expect(res.status).toBe(200); - const body = (await res.json()) as { kiloUserId: string }; - expect(body.kiloUserId).toBe('user-abc'); + + for (const response of responses) { + await expect(response.json()).resolves.toEqual({ + kiloUserId: 'user-abc', + isAdmin: true, + pepper: 'pepper', + gastownAccess: true, + memberships: [{ orgId: 'org-a', role: 'owner' }], + }); + } + + const defaults = await request(app, await signToken()); + for (const response of defaults) { + await expect(response.json()).resolves.toEqual({ + kiloUserId: 'user-abc', + isAdmin: false, + pepper: null, + gastownAccess: false, + memberships: [], + }); + } + }); + + it('preserves an explicit false admin claim', async () => { + const { app } = createApp(); + const responses = await request(app, await signToken({ isAdmin: false, orgMemberships: [] })); + + for (const response of responses) { + await expect(response.json()).resolves.toMatchObject({ isAdmin: false, memberships: [] }); + } + }); + + it.each([ + ['a missing secret', null], + [ + 'a failing secret provider', + { + get: async () => { + throw new Error('unavailable'); + }, + }, + ], + ])('returns 500 for %s without calling downstream', async (_name, secret) => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await signToken(), secret); + + expect(responses.map(response => response.status)).toEqual([500, 500]); + expect(downstreamCalls()).toBe(0); }); }); diff --git a/services/gastown/src/middleware/kilo-auth.middleware.ts b/services/gastown/src/middleware/kilo-auth.middleware.ts index 93f490db3f..12b990d5cc 100644 --- a/services/gastown/src/middleware/kilo-auth.middleware.ts +++ b/services/gastown/src/middleware/kilo-auth.middleware.ts @@ -1,5 +1,7 @@ import { createMiddleware } from 'hono/factory'; -import { verifyKiloToken, extractBearerToken } from '@kilocode/worker-utils'; +import { extractBearerToken } from '@kilocode/worker-utils/extract-bearer-token'; +import { GASTOWN_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { verifyKiloTokenForResource } from '@kilocode/worker-utils/kilo-token-policy'; import { resError } from '../util/res.util'; import type { GastownEnv } from '../gastown.worker'; import { resolveSecret } from '../util/secret.util'; @@ -30,7 +32,10 @@ export const kiloAuthMiddleware = createMiddleware(async (c, next) = } try { - const payload = await verifyKiloToken(token, secret); + const payload = await verifyKiloTokenForResource(token, secret, { + audience: GASTOWN_AUDIENCE, + mode: 'allow-legacy', + }); c.set('kiloUserId', payload.kiloUserId); c.set('kiloIsAdmin', payload.isAdmin === true); c.set('kiloApiTokenPepper', payload.apiTokenPepper ?? null); diff --git a/services/wasteland/src/middleware/kilo-auth.middleware.test.ts b/services/wasteland/src/middleware/kilo-auth.middleware.test.ts new file mode 100644 index 0000000000..e3e265cb7f --- /dev/null +++ b/services/wasteland/src/middleware/kilo-auth.middleware.test.ts @@ -0,0 +1,197 @@ +import { describe, expect, it } from 'vitest'; +import { Hono, type Context } from 'hono'; +import { SignJWT } from 'jose'; +import { WASTELAND_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { kiloAuthMiddleware } from './kilo-auth.middleware'; +import type { WastelandEnv } from '../wasteland.worker'; + +const TEST_SECRET = 'test-secret-that-is-long-enough-for-hs256'; + +function createApp() { + let downstreamCalls = 0; + const app = new Hono(); + app.use('/api/*', kiloAuthMiddleware); + app.use('/trpc/*', kiloAuthMiddleware); + const handler = (c: Context) => { + downstreamCalls += 1; + return c.json({ + kiloUserId: c.get('kiloUserId'), + isAdmin: c.get('kiloIsAdmin'), + pepper: c.get('kiloApiTokenPepper'), + memberships: c.get('kiloOrgMemberships'), + }); + }; + app.get('/api/whoami', handler); + app.get('/trpc/whoami', handler); + return { app, downstreamCalls: () => downstreamCalls }; +} + +async function signToken( + claims: Record = {}, + options: { secret?: string; dates?: boolean; expiration?: number } = {} +) { + const now = Math.floor(Date.now() / 1000); + let jwt = new SignJWT({ version: 3, kiloUserId: 'user-abc', ...claims }).setProtectedHeader({ + alg: 'HS256', + }); + if (options.dates !== false) { + jwt = jwt.setIssuedAt(now).setExpirationTime(options.expiration ?? now + 3600); + } + return jwt.sign(new TextEncoder().encode(options.secret ?? TEST_SECRET)); +} + +async function request( + app: Hono, + token: string | undefined, + secret: string | { get(): Promise } | null = TEST_SECRET +) { + return Promise.all( + ['/api/whoami', '/trpc/whoami'].map(path => + app.request( + path, + token === undefined ? {} : { headers: { Authorization: `Bearer ${token}` } }, + { NEXTAUTH_SECRET: secret } as never + ) + ) + ); +} + +describe('kiloAuthMiddleware', () => { + it.each([ + ['missing authentication', undefined], + ['malformed authentication', 'Bearer'], + ])('rejects %s before calling downstream', async (_name, authorization) => { + const { app, downstreamCalls } = createApp(); + const responses = + authorization === undefined + ? await request(app, undefined) + : await Promise.all( + ['/api/whoami', '/trpc/whoami'].map(path => + app.request(path, { headers: { Authorization: authorization } }, { + NEXTAUTH_SECRET: TEST_SECRET, + } as never) + ) + ); + + expect(responses.map(response => response.status)).toEqual([401, 401]); + expect(downstreamCalls()).toBe(0); + }); + + it('accepts legacy tokens without an audience or dates', async () => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await signToken({}, { dates: false })); + + expect(responses.map(response => response.status)).toEqual([200, 200]); + expect(downstreamCalls()).toBe(2); + }); + + it('accepts its audience as a string or an array', async () => { + for (const aud of [WASTELAND_AUDIENCE, ['kilo-api', WASTELAND_AUDIENCE]]) { + const { app } = createApp(); + const responses = await request(app, await signToken({ aud })); + expect(responses.map(response => response.status)).toEqual([200, 200]); + } + }); + + it.each(['gastown', 'kilo-api', 'kilo-gateway', 'git-token-service:github-user-access-token'])( + 'rejects a token for %s before calling downstream', + async aud => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await signToken({ aud })); + + expect(responses.map(response => response.status)).toEqual([401, 401]); + expect(downstreamCalls()).toBe(0); + } + ); + + it.each([ + false, + null, + '', + ` ${WASTELAND_AUDIENCE}`, + [WASTELAND_AUDIENCE, ''], + [WASTELAND_AUDIENCE, WASTELAND_AUDIENCE], + ])('rejects malformed explicit audiences before calling downstream', async aud => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await signToken({ aud })); + + expect(responses.map(response => response.status)).toEqual([401, 401]); + expect(downstreamCalls()).toBe(0); + }); + + it.each([ + [ + 'invalid signature', + () => signToken({}, { secret: 'different-secret-that-is-long-enough-for-hs256' }), + ], + ['expired token', () => signToken({}, { expiration: Math.floor(Date.now() / 1000) - 60 })], + ['wrong version', () => signToken({ version: 2 })], + ['missing user', () => signToken({ kiloUserId: '' })], + ])('rejects %s before calling downstream', async (_name, token) => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await token()); + + expect(responses.map(response => response.status)).toEqual([401, 401]); + expect(downstreamCalls()).toBe(0); + }); + + it('preserves claim values and defaults', async () => { + const { app } = createApp(); + const responses = await request( + app, + await signToken({ + aud: WASTELAND_AUDIENCE, + isAdmin: true, + apiTokenPepper: 'pepper', + orgMemberships: [{ orgId: 'org-a', role: 'owner' }], + deviceSessionId: 'session-wasteland-test', + }) + ); + + for (const response of responses) { + await expect(response.json()).resolves.toEqual({ + kiloUserId: 'user-abc', + isAdmin: true, + pepper: 'pepper', + memberships: [{ orgId: 'org-a', role: 'owner' }], + }); + } + + const defaults = await request(app, await signToken()); + for (const response of defaults) { + await expect(response.json()).resolves.toEqual({ + kiloUserId: 'user-abc', + isAdmin: false, + pepper: null, + memberships: [], + }); + } + }); + + it('preserves an explicit false admin claim', async () => { + const { app } = createApp(); + const responses = await request(app, await signToken({ isAdmin: false, orgMemberships: [] })); + + for (const response of responses) { + await expect(response.json()).resolves.toMatchObject({ isAdmin: false, memberships: [] }); + } + }); + + it.each([ + ['a missing secret', null], + [ + 'a failing secret provider', + { + get: async () => { + throw new Error('unavailable'); + }, + }, + ], + ])('returns 500 for %s without calling downstream', async (_name, secret) => { + const { app, downstreamCalls } = createApp(); + const responses = await request(app, await signToken(), secret); + + expect(responses.map(response => response.status)).toEqual([500, 500]); + expect(downstreamCalls()).toBe(0); + }); +}); diff --git a/services/wasteland/src/middleware/kilo-auth.middleware.ts b/services/wasteland/src/middleware/kilo-auth.middleware.ts index ed3e15796f..8517c45620 100644 --- a/services/wasteland/src/middleware/kilo-auth.middleware.ts +++ b/services/wasteland/src/middleware/kilo-auth.middleware.ts @@ -1,5 +1,7 @@ import { createMiddleware } from 'hono/factory'; -import { verifyKiloToken, extractBearerToken } from '@kilocode/worker-utils'; +import { extractBearerToken } from '@kilocode/worker-utils/extract-bearer-token'; +import { WASTELAND_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { verifyKiloTokenForResource } from '@kilocode/worker-utils/kilo-token-policy'; import { resError } from '../util/res.util'; import type { WastelandEnv } from '../wasteland.worker'; import { resolveSecret } from '../util/secret.util'; @@ -30,7 +32,10 @@ export const kiloAuthMiddleware = createMiddleware(async (c, next) } try { - const payload = await verifyKiloToken(token, secret); + const payload = await verifyKiloTokenForResource(token, secret, { + audience: WASTELAND_AUDIENCE, + mode: 'allow-legacy', + }); c.set('kiloUserId', payload.kiloUserId); c.set('kiloIsAdmin', payload.isAdmin === true); c.set('kiloApiTokenPepper', payload.apiTokenPepper ?? null); From aff7f5aa0344a8e35f0146ed98be57e0938c19e0 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 1 Sep 2026 14:39:35 -0500 Subject: [PATCH 2/3] feat(auth): enable audience-aware readers for chat and KiloClaw --- .../event-service/src/__tests__/auth.test.ts | 153 +++++++++- services/event-service/src/auth.ts | 2 + services/kilo-chat/src/__tests__/auth.test.ts | 267 ++++++++++------- services/kilo-chat/src/auth.ts | 2 + services/kiloclaw/src/auth/jwt.test.ts | 51 ++++ services/kiloclaw/src/auth/jwt.ts | 10 +- services/kiloclaw/src/auth/middleware.test.ts | 155 +++++++++- .../src/routes/access-gateway.test.ts | 76 ++++- .../notifications/src/__tests__/auth.test.ts | 282 ++++++++++-------- services/notifications/src/auth.ts | 2 + 10 files changed, 764 insertions(+), 236 deletions(-) diff --git a/services/event-service/src/__tests__/auth.test.ts b/services/event-service/src/__tests__/auth.test.ts index f512a4cc66..9f8aeeca67 100644 --- a/services/event-service/src/__tests__/auth.test.ts +++ b/services/event-service/src/__tests__/auth.test.ts @@ -1,5 +1,9 @@ -import { beforeEach, describe, expect, it } from 'vitest'; -import { clearSecretCacheForTest, signKiloToken } from '@kilocode/worker-utils'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { + clearSecretCacheForTest, + EVENT_SERVICE_AUDIENCE, + signKiloToken, +} from '@kilocode/worker-utils'; import { type KiloUserPepperResult } from '@kilocode/worker-utils/kilo-token-auth'; import { type AuthEnv, authenticateToken } from '../auth'; @@ -25,6 +29,45 @@ function authenticateTestToken(token: string | null) { return authenticateToken(token, makeEnv(), { getUserPepper }); } +function signEventServiceToken(params: { pepper?: string | null; env?: string } = {}) { + return signKiloToken({ + userId: 'user-xyz-789', + pepper: params.pepper, + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: params.env, + audience: EVENT_SERVICE_AUDIENCE, + }); +} + +async function signWithAudience(aud: unknown): Promise { + const now = Math.floor(Date.now() / 1000); + const encode = (bytes: Uint8Array) => + btoa(String.fromCharCode(...bytes)) + .replace(/=/g, '') + .replace(/\+/g, '-') + .replace(/\//g, '_'); + const json = (value: unknown) => encode(new TextEncoder().encode(JSON.stringify(value))); + const input = `${json({ alg: 'HS256', typ: 'JWT' })}.${json({ + version: 3, + kiloUserId: 'user-xyz-789', + apiTokenPepper: 'pepper-current', + env: 'production', + aud, + iat: now, + exp: now + 3600, + })}`; + const key = await crypto.subtle.importKey( + 'raw', + new TextEncoder().encode(TEST_JWT_SECRET), + { name: 'HMAC', hash: 'SHA-256' }, + false, + ['sign'] + ); + const signature = await crypto.subtle.sign('HMAC', key, new TextEncoder().encode(input)); + return `${input}.${encode(new Uint8Array(signature))}`; +} + describe('authenticateToken', () => { beforeEach(() => { clearSecretCacheForTest(); @@ -32,7 +75,7 @@ describe('authenticateToken', () => { currentPepperByUserId.set('user-xyz-789', { pepper: 'pepper-current', blockedReason: null }); }); - it('authenticates a kilo-chat token with the current pepper', async () => { + it('authenticates a legacy one-hour kilo-chat token without an audience', async () => { const { token } = await signKiloToken({ userId: 'user-xyz-789', pepper: 'pepper-current', @@ -60,6 +103,49 @@ describe('authenticateToken', () => { }); }); + it('authenticates event-service audience claims as strings and arrays', async () => { + const arrayAudienceToken = await signWithAudience(['another-service', EVENT_SERVICE_AUDIENCE]); + const stringAudienceToken = await signEventServiceToken({ + pepper: 'pepper-current', + env: 'production', + }); + + await expect(authenticateTestToken(stringAudienceToken.token)).resolves.toEqual({ + userId: 'user-xyz-789', + }); + await expect(authenticateTestToken(arrayAudienceToken)).resolves.toEqual({ + userId: 'user-xyz-789', + }); + }); + + it('rejects wrong or malformed event-service audiences before pepper lookup', async () => { + const malformedAudienceToken = await signWithAudience([ + EVENT_SERVICE_AUDIENCE, + EVENT_SERVICE_AUDIENCE, + ]); + const wrongAudienceToken = await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + audience: 'another-service', + }); + let lookupCount = 0; + const lookup = async (...args: Parameters) => { + lookupCount++; + return getUserPepper(...args); + }; + + await expect( + authenticateToken(wrongAudienceToken.token, makeEnv(), { getUserPepper: lookup }) + ).resolves.toBeNull(); + await expect( + authenticateToken(malformedAudienceToken, makeEnv(), { getUserPepper: lookup }) + ).resolves.toBeNull(); + expect(lookupCount).toBe(0); + }); + it('rejects a valid kilo-chat JWT with a stale pepper', async () => { const { token } = await signKiloToken({ userId: 'user-xyz-789', @@ -86,6 +172,16 @@ describe('authenticateToken', () => { await expect(authenticateTestToken(token)).resolves.toBeNull(); }); + it('rejects event-service tokens with missing or mismatched environments', async () => { + const [missingEnvironment, mismatchedEnvironment] = await Promise.all([ + signEventServiceToken({ pepper: 'pepper-current' }), + signEventServiceToken({ pepper: 'pepper-current', env: 'development' }), + ]); + + await expect(authenticateTestToken(missingEnvironment.token)).resolves.toBeNull(); + await expect(authenticateTestToken(mismatchedEnvironment.token)).resolves.toBeNull(); + }); + it('rejects a token for a blocked user even when pepper matches', async () => { currentPepperByUserId.set('user-xyz-789', { pepper: 'pepper-current', @@ -102,4 +198,55 @@ describe('authenticateToken', () => { await expect(authenticateTestToken(token)).resolves.toBeNull(); }); + + it('rejects missing users, stale peppers, and a null pepper against a stored pepper', async () => { + const [stalePepper, nullPepper] = await Promise.all([ + signEventServiceToken({ pepper: 'pepper-stale', env: 'production' }), + signEventServiceToken({ pepper: null, env: 'production' }), + ]); + + await expect(authenticateTestToken(stalePepper.token)).resolves.toBeNull(); + await expect(authenticateTestToken(nullPepper.token)).resolves.toBeNull(); + const validToken = await signEventServiceToken({ pepper: 'pepper-current', env: 'production' }); + await expect(authenticateTestToken(validToken.token)).resolves.toEqual({ + userId: 'user-xyz-789', + }); + currentPepperByUserId.clear(); + const lookup = vi.fn(getUserPepper); + await expect( + authenticateToken(validToken.token, makeEnv(), { getUserPepper: lookup }) + ).resolves.toBeNull(); + expect(lookup).toHaveBeenCalledWith('postgres://test', 'user-xyz-789'); + }); + + it('allows absent peppers but rejects explicit null peppers for a user with a current pepper', async () => { + const absentPepper = await signEventServiceToken({ env: 'production' }); + const nullPepper = await signEventServiceToken({ pepper: null, env: 'production' }); + + await expect(authenticateTestToken(absentPepper.token)).resolves.toEqual({ + userId: 'user-xyz-789', + }); + await expect(authenticateTestToken(nullPepper.token)).resolves.toBeNull(); + }); + + it('propagates secret and pepper lookup failures', async () => { + const { token } = await signEventServiceToken({ pepper: 'pepper-current', env: 'production' }); + + await expect( + authenticateToken( + token, + { + ...makeEnv(), + NEXTAUTH_SECRET: { get: async () => Promise.reject(new Error('secret unavailable')) }, + }, + { getUserPepper } + ) + ).rejects.toThrow('secret unavailable'); + clearSecretCacheForTest(); + await expect( + authenticateToken(token, makeEnv(), { + getUserPepper: async () => Promise.reject(new Error('lookup unavailable')), + }) + ).rejects.toThrow('lookup unavailable'); + }); }); diff --git a/services/event-service/src/auth.ts b/services/event-service/src/auth.ts index c9623bc1a7..1991eeb1e3 100644 --- a/services/event-service/src/auth.ts +++ b/services/event-service/src/auth.ts @@ -1,3 +1,4 @@ +import { EVENT_SERVICE_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; import { type GetKiloUserPepper, verifyKiloBearerAgainstCurrentPepper, @@ -21,6 +22,7 @@ export async function authenticateToken( nextAuthSecret: env.NEXTAUTH_SECRET, workerEnv: env.WORKER_ENV, connectionString: env.HYPERDRIVE.connectionString, + resourceAudience: { audience: EVENT_SERVICE_AUDIENCE, mode: 'allow-legacy' }, ...(options.getUserPepper ? { getUserPepper: options.getUserPepper } : {}), }); } diff --git a/services/kilo-chat/src/__tests__/auth.test.ts b/services/kilo-chat/src/__tests__/auth.test.ts index aae4b2fc5f..ab1d326453 100644 --- a/services/kilo-chat/src/__tests__/auth.test.ts +++ b/services/kilo-chat/src/__tests__/auth.test.ts @@ -1,8 +1,13 @@ -import { beforeEach, describe, it, expect, vi } from 'vitest'; import { Hono } from 'hono'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import { signKiloToken } from '@kilocode/worker-utils'; -import { authMiddleware } from '../auth'; -import type { AuthContext } from '../auth'; +import { + EVENT_SERVICE_AUDIENCE, + KILO_CHAT_AUDIENCE, + KILO_GATEWAY_AUDIENCE, + NOTIFICATIONS_AUDIENCE, +} from '@kilocode/worker-utils/internal-service-token-audiences'; +import { authMiddleware, type AuthContext } from '../auth'; type MockEnv = { NEXTAUTH_SECRET: { get: () => Promise }; @@ -11,151 +16,197 @@ type MockEnv = { }; const TEST_JWT_SECRET = 'test-secret-that-is-long-enough-for-hs256'; -const currentPepperByUserId = vi.hoisted(() => new Map()); +const dbState = vi.hoisted(() => ({ lookupCount: 0, fails: false })); +const userRow = vi.hoisted(() => ({ + pepper: 'pepper-current' as string | null, + blockedReason: null as string | null, +})); vi.mock('@kilocode/db/client', () => ({ getWorkerDb: () => ({ select: () => ({ from: () => ({ where: () => ({ - limit: async () => [ - { - api_token_pepper: currentPepperByUserId.get('user-xyz-789'), - blocked_reason: null, - }, - ], + limit: async () => { + dbState.lookupCount++; + if (dbState.fails) throw new Error('connection refused'); + return [{ api_token_pepper: userRow.pepper, blocked_reason: userRow.blockedReason }]; + }, }), }), }), }), })); -function makeApp(_env: MockEnv) { - const app = new Hono<{ Bindings: MockEnv; Variables: AuthContext }>(); - app.use('*', authMiddleware); - app.get('/test', c => c.json({ callerId: c.get('callerId'), callerKind: c.get('callerKind') })); - return app; -} - const defaultEnv: MockEnv = { NEXTAUTH_SECRET: { get: async () => TEST_JWT_SECRET }, HYPERDRIVE: { connectionString: 'postgres://test' }, WORKER_ENV: 'production', }; +function makeApp() { + let downstreamCalls = 0; + const app = new Hono<{ Bindings: MockEnv; Variables: AuthContext }>(); + app.use('*', authMiddleware); + app.get('/test', c => { + downstreamCalls++; + return c.json({ callerId: c.get('callerId'), callerKind: c.get('callerKind') }); + }); + return { app, downstreamCalls: () => downstreamCalls }; +} + +async function signToken( + params: { + audience?: string; + pepper?: string | null; + env?: string; + extra?: { tokenSource?: string; botId?: string; deviceSessionId?: string }; + } = {} +) { + return ( + await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper' in params ? params.pepper : 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'env' in params ? params.env : 'production', + audience: params.audience, + extra: params.extra, + }) + ).token; +} + +async function signWithAudience(aud: unknown): Promise { + const now = Math.floor(Date.now() / 1000); + const encode = (bytes: Uint8Array) => + btoa(String.fromCharCode(...bytes)) + .replace(/=/g, '') + .replace(/\+/g, '-') + .replace(/\//g, '_'); + const json = (value: unknown) => encode(new TextEncoder().encode(JSON.stringify(value))); + const input = `${json({ alg: 'HS256', typ: 'JWT' })}.${json({ + version: 3, + kiloUserId: 'user-xyz-789', + apiTokenPepper: 'pepper-current', + env: 'production', + aud, + iat: now, + exp: now + 3600, + })}`; + const key = await crypto.subtle.importKey( + 'raw', + new TextEncoder().encode(TEST_JWT_SECRET), + { name: 'HMAC', hash: 'SHA-256' }, + false, + ['sign'] + ); + const signature = await crypto.subtle.sign('HMAC', key, new TextEncoder().encode(input)); + return `${input}.${encode(new Uint8Array(signature))}`; +} + +async function request(token: string, env = defaultEnv) { + const testApp = makeApp(); + const response = await testApp.app.request( + '/test', + { headers: { authorization: `Bearer ${token}` } }, + env + ); + return { response, ...testApp }; +} + describe('authMiddleware', () => { beforeEach(() => { - currentPepperByUserId.set('user-xyz-789', 'pepper-current'); + dbState.lookupCount = 0; + dbState.fails = false; + userRow.pepper = 'pepper-current'; + userRow.blockedReason = null; }); it('returns 401 with no authorization header', async () => { - const res = await makeApp(defaultEnv).request('/test', {}, defaultEnv); - expect(res.status).toBe(401); - expect(await res.json()).toEqual({ error: 'Unauthorized' }); + const testApp = makeApp(); + const response = await testApp.app.request('/test', {}, defaultEnv); + expect(response.status).toBe(401); + expect(await response.json()).toEqual({ error: 'Unauthorized' }); }); - it('authenticates with a valid JWT and sets user identity', async () => { - const { token } = await signKiloToken({ - userId: 'user-xyz-789', - pepper: 'pepper-current', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'production', - extra: { tokenSource: 'kilo-chat' }, - }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv - ); - expect(res.status).toBe(200); - expect(await res.json()).toEqual({ - callerId: 'user-xyz-789', - callerKind: 'user', - }); + it.each([ + ['a matching string audience', () => signToken({ audience: KILO_CHAT_AUDIENCE })], + ['a matching array audience', () => signWithAudience([KILO_CHAT_AUDIENCE, 'other-service'])], + [ + 'a legacy kilo-chat token with bot and device-session claims', + () => + signToken({ + extra: { + tokenSource: 'kilo-chat', + botId: 'bot-123', + deviceSessionId: 'device-123', + }, + }), + ], + ])('authenticates %s', async (_name, createToken) => { + const { response } = await request(await createToken()); + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ callerId: 'user-xyz-789', callerKind: 'user' }); }); - it('authenticates a valid JWT from another token source', async () => { - const { token } = await signKiloToken({ - userId: 'user-xyz-789', - pepper: 'pepper-current', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'production', - extra: { tokenSource: 'cloud-agent' }, - }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv - ); - expect(res.status).toBe(200); - expect(await res.json()).toEqual({ - callerId: 'user-xyz-789', - callerKind: 'user', - }); + it.each([ + ['another batch', NOTIFICATIONS_AUDIENCE], + ['the gateway', KILO_GATEWAY_AUDIENCE], + ['another service', EVENT_SERVICE_AUDIENCE], + ])('rejects an audience for %s before database access', async (_name, audience) => { + const { response, downstreamCalls } = await request(await signToken({ audience })); + expect(response.status).toBe(401); + expect(dbState.lookupCount).toBe(0); + expect(downstreamCalls()).toBe(0); }); - it('returns 401 when the chat JWT has a stale pepper', async () => { - const { token } = await signKiloToken({ - userId: 'user-xyz-789', - pepper: 'pepper-stale', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'production', - extra: { tokenSource: 'kilo-chat' }, - }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv - ); - expect(res.status).toBe(401); + it.each([null, '', ['kilo-chat', 'kilo-chat'], ['kilo-chat', 1], []])( + 'rejects malformed audience claims before database access', + async audience => { + const { response, downstreamCalls } = await request(await signWithAudience(audience)); + expect(response.status).toBe(401); + expect(dbState.lookupCount).toBe(0); + expect(downstreamCalls()).toBe(0); + } + ); + + it.each([ + ['a stale string pepper', 'pepper-stale', 'pepper-current', 401], + ['a null claim against a non-null pepper', null, 'pepper-current', 401], + ['a null claim against a null pepper', null, null, 200], + ])('preserves pepper semantics for %s', async (_name, pepper, currentPepper, status) => { + userRow.pepper = currentPepper; + const { response } = await request(await signToken({ audience: KILO_CHAT_AUDIENCE, pepper })); + expect(response.status).toBe(status); }); - it('returns 401 when the chat JWT was minted for a different environment', async () => { + it('preserves absent pepper claim semantics', async () => { const { token } = await signKiloToken({ userId: 'user-xyz-789', - pepper: 'pepper-current', secret: TEST_JWT_SECRET, expiresInSeconds: 3600, - env: 'development', - extra: { tokenSource: 'kilo-chat' }, + env: 'production', + audience: KILO_CHAT_AUDIENCE, }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv - ); - expect(res.status).toBe(401); + const { response } = await request(token); + expect(response.status).toBe(200); }); - it('returns 401 with an expired JWT', async () => { - const { token } = await signKiloToken({ - userId: 'user-xyz-789', - pepper: null, - secret: TEST_JWT_SECRET, - expiresInSeconds: -1, - env: 'production', - extra: { tokenSource: 'kilo-chat' }, - }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv + it.each([ + ['a blocked account', () => (userRow.blockedReason = 'manual block'), 'production'], + ['a missing env claim', () => undefined, undefined], + ['a mismatched environment', () => undefined, 'development'], + ])('returns 401 for %s', async (_name, arrange, tokenEnv) => { + arrange(); + const { response } = await request( + await signToken({ audience: KILO_CHAT_AUDIENCE, env: tokenEnv }) ); - expect(res.status).toBe(401); - expect(await res.json()).toEqual({ error: 'Unauthorized' }); + expect(response.status).toBe(401); }); - it('does not accept arbitrary bearers as bots — there is no HTTP bot surface', async () => { - // Bots reach kilo-chat via service-binding RPC only; no HTTP path grants - // bot identity. Any non-JWT bearer must fail closed. - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: 'Bearer not-a-jwt' } }, - defaultEnv - ); - expect(res.status).toBe(401); + it('maps dependency failures to 401', async () => { + dbState.fails = true; + const { response } = await request(await signToken({ audience: KILO_CHAT_AUDIENCE })); + expect(response.status).toBe(401); }); }); diff --git a/services/kilo-chat/src/auth.ts b/services/kilo-chat/src/auth.ts index 1be48bc694..c6fcf183ef 100644 --- a/services/kilo-chat/src/auth.ts +++ b/services/kilo-chat/src/auth.ts @@ -1,6 +1,7 @@ import { createMiddleware } from 'hono/factory'; import { verifyKiloBearerAgainstCurrentPepper } from '@kilocode/worker-utils/kilo-token-auth'; import { extractBearerToken } from '@kilocode/worker-utils'; +import { KILO_CHAT_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; import { logger } from './util/logger'; export type AuthContext = { @@ -31,6 +32,7 @@ export const authMiddleware = createMiddleware<{ nextAuthSecret: c.env.NEXTAUTH_SECRET, workerEnv: c.env.WORKER_ENV, connectionString: c.env.HYPERDRIVE.connectionString, + resourceAudience: { audience: KILO_CHAT_AUDIENCE, mode: 'allow-legacy' }, }); if (!auth) { return c.json({ error: 'Unauthorized' }, 401); diff --git a/services/kiloclaw/src/auth/jwt.test.ts b/services/kiloclaw/src/auth/jwt.test.ts index 17f200e18a..0af8f0203c 100644 --- a/services/kiloclaw/src/auth/jwt.test.ts +++ b/services/kiloclaw/src/auth/jwt.test.ts @@ -2,6 +2,7 @@ import { describe, it, expect } from 'vitest'; import { SignJWT } from 'jose'; import { validateKiloToken } from './jwt'; import { KILO_TOKEN_VERSION } from '../config'; +import { KILOCLAW_AUDIENCE } from '@kilocode/worker-utils'; const TEST_SECRET = 'test-secret-for-jwt-verification'; @@ -37,6 +38,56 @@ describe('validateKiloToken', () => { }); }); + it('accepts the KiloClaw audience as a string or array member', async () => { + for (const aud of [KILOCLAW_AUDIENCE, ['another-resource', KILOCLAW_AUDIENCE]]) { + const token = await signToken({ + kiloUserId: 'user_123', + apiTokenPepper: 'pepper_abc', + version: KILO_TOKEN_VERSION, + aud, + }); + + await expect(validateKiloToken(token, TEST_SECRET, undefined)).resolves.toMatchObject({ + success: true, + userId: 'user_123', + token, + pepper: 'pepper_abc', + }); + } + }); + + it('rejects wrong and malformed audiences', async () => { + for (const aud of ['another-resource', [], [' kiloclaw']]) { + const token = await signToken({ + kiloUserId: 'user_123', + apiTokenPepper: 'pepper_abc', + version: KILO_TOKEN_VERSION, + aud, + }); + + await expect(validateKiloToken(token, TEST_SECRET, undefined)).resolves.toMatchObject({ + success: false, + }); + } + }); + + it('preserves legacy tokens without an audience or date claims', async () => { + const token = await new SignJWT({ + kiloUserId: 'user_123', + apiTokenPepper: 'pepper_abc', + version: KILO_TOKEN_VERSION, + }) + .setProtectedHeader({ alg: 'HS256' }) + .sign(new TextEncoder().encode(TEST_SECRET)); + + await expect(validateKiloToken(token, TEST_SECRET, undefined)).resolves.toEqual({ + success: true, + userId: 'user_123', + token, + pepper: 'pepper_abc', + }); + }); + it('rejects wrong token version', async () => { const token = await signToken({ kiloUserId: 'user_123', diff --git a/services/kiloclaw/src/auth/jwt.ts b/services/kiloclaw/src/auth/jwt.ts index 392a5f782f..0643749778 100644 --- a/services/kiloclaw/src/auth/jwt.ts +++ b/services/kiloclaw/src/auth/jwt.ts @@ -1,5 +1,6 @@ import { SignJWT } from 'jose'; -import { verifyKiloToken } from '@kilocode/worker-utils'; +import { KILOCLAW_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { verifyKiloTokenForResource } from '@kilocode/worker-utils/kilo-token-policy'; import { KILO_TOKEN_VERSION, KILOCLAW_AUTH_COOKIE_MAX_AGE } from '../config'; export type ValidateResult = @@ -17,9 +18,12 @@ export async function validateKiloToken( secret: string, expectedEnv: string | undefined ): Promise { - let payload: Awaited>; + let payload: Awaited>; try { - payload = await verifyKiloToken(token, secret); + payload = await verifyKiloTokenForResource(token, secret, { + audience: KILOCLAW_AUDIENCE, + mode: 'allow-legacy', + }); } catch (err) { const message = err instanceof Error ? err.message : 'JWT verification failed'; return { success: false, error: message }; diff --git a/services/kiloclaw/src/auth/middleware.test.ts b/services/kiloclaw/src/auth/middleware.test.ts index ccc23b5e31..ad715fee12 100644 --- a/services/kiloclaw/src/auth/middleware.test.ts +++ b/services/kiloclaw/src/auth/middleware.test.ts @@ -4,12 +4,16 @@ import { SignJWT } from 'jose'; import type { AppEnv } from '../types'; import { authMiddleware, internalApiMiddleware } from './middleware'; import { KILO_TOKEN_VERSION, KILOCLAW_AUTH_COOKIE } from '../config'; +import { KILOCLAW_AUDIENCE } from '@kilocode/worker-utils'; +import { findPepperByUserId, getWorkerDb } from '../db'; + +let downstreamExecutions = 0; vi.mock('../db', () => ({ getWorkerDb: vi.fn(() => ({})), findPepperByUserId: vi.fn(async (_db: unknown, userId: string) => ({ id: userId, - api_token_pepper: `pepper_for_${userId}`, + api_token_pepper: userId === 'pepperless_user' ? null : `pepper_for_${userId}`, blocked_reason: userId === 'blocked_user' ? 'abuse' : null, })), })); @@ -36,6 +40,7 @@ function createTestApp() { // Auth-protected route app.use('/protected/*', authMiddleware); app.get('/protected/whoami', c => { + downstreamExecutions += 1; return c.json({ userId: c.get('userId'), authToken: c.get('authToken') }); }); @@ -62,6 +67,8 @@ describe('authMiddleware', () => { let app: ReturnType; beforeEach(() => { + vi.clearAllMocks(); + downstreamExecutions = 0; app = createTestApp(); }); @@ -131,6 +138,98 @@ describe('authMiddleware', () => { expect(body.userId).toBe('user_cookie'); }); + it.each([ + { tokenPepper: 'absent', storedPepper: null, expectedStatus: 200 }, + { tokenPepper: null, storedPepper: null, expectedStatus: 200 }, + { tokenPepper: 'absent', storedPepper: 'rotated_pepper', expectedStatus: 401 }, + { tokenPepper: null, storedPepper: 'rotated_pepper', expectedStatus: 401 }, + ])( + 'validates $tokenPepper token pepper against $storedPepper stored pepper', + async ({ tokenPepper, storedPepper, expectedStatus }) => { + const token = await signToken({ + kiloUserId: 'pepperless_user', + ...(tokenPepper === null ? { apiTokenPepper: null } : {}), + version: KILO_TOKEN_VERSION, + aud: KILOCLAW_AUDIENCE, + }); + const lookup = vi.mocked(findPepperByUserId); + lookup.mockResolvedValueOnce({ + id: 'pepperless_user', + api_token_pepper: storedPepper, + blocked_reason: null, + }); + + const res = await app.request( + '/protected/whoami', + { headers: { Authorization: `Bearer ${token}` } }, + ENV_WITH_HYPERDRIVE + ); + + expect(res.status).toBe(expectedStatus); + expect(lookup).toHaveBeenCalledOnce(); + if (storedPepper === null) { + expect(await jsonBody(res)).toEqual({ userId: 'pepperless_user', authToken: token }); + } else { + expect(await jsonBody(res)).toEqual({ error: 'Token revoked' }); + } + } + ); + + it('authenticates correct-audience Bearer and cookie tokens', async () => { + for (const headers of [ + { Authorization: 'Bearer TOKEN' }, + { Cookie: `${KILOCLAW_AUTH_COOKIE}=TOKEN` }, + ]) { + const token = await signToken({ + kiloUserId: 'user_123', + apiTokenPepper: pepperFor('user_123'), + version: KILO_TOKEN_VERSION, + aud: KILOCLAW_AUDIENCE, + }); + const resolvedHeaders = Object.fromEntries( + Object.entries(headers).map(([name, value]) => [name, value.replace('TOKEN', token)]) + ); + + const res = await app.request( + '/protected/whoami', + { headers: resolvedHeaders }, + ENV_WITH_HYPERDRIVE + ); + expect(res.status).toBe(200); + expect(await jsonBody(res)).toEqual({ userId: 'user_123', authToken: token }); + } + }); + + it('rejects wrong-audience Bearer and cookie tokens before database lookup', async () => { + const lookup = vi.mocked(findPepperByUserId); + const workerDb = vi.mocked(getWorkerDb); + for (const headers of [ + { Authorization: 'Bearer TOKEN' }, + { Cookie: `${KILOCLAW_AUTH_COOKIE}=TOKEN` }, + ]) { + const token = await signToken({ + kiloUserId: 'user_123', + apiTokenPepper: pepperFor('user_123'), + version: KILO_TOKEN_VERSION, + aud: 'another-resource', + }); + const resolvedHeaders = Object.fromEntries( + Object.entries(headers).map(([name, value]) => [name, value.replace('TOKEN', token)]) + ); + lookup.mockClear(); + + const res = await app.request( + '/protected/whoami', + { headers: resolvedHeaders }, + ENV_WITH_HYPERDRIVE + ); + expect(res.status).toBe(401); + expect(lookup).not.toHaveBeenCalled(); + expect(workerDb).not.toHaveBeenCalled(); + expect(downstreamExecutions).toBe(0); + } + }); + it('prefers Bearer header over cookie', async () => { const bearerToken = await signToken({ kiloUserId: 'user_bearer', @@ -158,6 +257,53 @@ describe('authMiddleware', () => { expect(body.userId).toBe('user_bearer'); }); + it('does not fall back to a valid cookie when the Bearer token has the wrong audience', async () => { + const bearerToken = await signToken({ + kiloUserId: 'user_bearer', + apiTokenPepper: pepperFor('user_bearer'), + version: KILO_TOKEN_VERSION, + aud: 'another-resource', + }); + const cookieToken = await signToken({ + kiloUserId: 'user_cookie', + apiTokenPepper: pepperFor('user_cookie'), + version: KILO_TOKEN_VERSION, + aud: KILOCLAW_AUDIENCE, + }); + + const res = await app.request( + '/protected/whoami', + { + headers: { + Authorization: `Bearer ${bearerToken}`, + Cookie: `${KILOCLAW_AUTH_COOKIE}=${cookieToken}`, + }, + }, + ENV_WITH_HYPERDRIVE + ); + expect(res.status).toBe(401); + expect(vi.mocked(findPepperByUserId)).not.toHaveBeenCalled(); + expect(vi.mocked(getWorkerDb)).not.toHaveBeenCalled(); + expect(downstreamExecutions).toBe(0); + }); + + it('rejects a correct-audience token with a stale pepper', async () => { + const token = await signToken({ + kiloUserId: 'user_123', + apiTokenPepper: 'stale_pepper', + version: KILO_TOKEN_VERSION, + aud: KILOCLAW_AUDIENCE, + }); + + const res = await app.request( + '/protected/whoami', + { headers: { Authorization: `Bearer ${token}` } }, + ENV_WITH_HYPERDRIVE + ); + expect(res.status).toBe(401); + expect((await jsonBody(res)).error).toContain('revoked'); + }); + it('rejects when pepper does not match', async () => { const token = await signToken({ kiloUserId: 'user_123', @@ -226,6 +372,8 @@ describe('blocked users', () => { let app: ReturnType; beforeEach(() => { + vi.clearAllMocks(); + downstreamExecutions = 0; app = createTestApp(); }); @@ -234,6 +382,7 @@ describe('blocked users', () => { kiloUserId: 'blocked_user', apiTokenPepper: pepperFor('blocked_user'), version: KILO_TOKEN_VERSION, + aud: KILOCLAW_AUDIENCE, }); const res = await app.request( @@ -249,6 +398,8 @@ describe('C15 deviceSessionId compatibility', () => { let app: ReturnType; beforeEach(() => { + vi.clearAllMocks(); + downstreamExecutions = 0; app = createTestApp(); }); @@ -276,6 +427,8 @@ describe('internalApiMiddleware', () => { let app: ReturnType; beforeEach(() => { + vi.clearAllMocks(); + downstreamExecutions = 0; app = createTestApp(); }); diff --git a/services/kiloclaw/src/routes/access-gateway.test.ts b/services/kiloclaw/src/routes/access-gateway.test.ts index 22fc9a1eea..4ae5fc1ed5 100644 --- a/services/kiloclaw/src/routes/access-gateway.test.ts +++ b/services/kiloclaw/src/routes/access-gateway.test.ts @@ -1,12 +1,28 @@ -import { describe, expect, it, vi } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import { Hono } from 'hono'; import type { AppEnv } from '../types'; +import type * as dbModule from '../db'; + +vi.mock('../db', async importOriginal => ({ + ...(await importOriginal()), + getWorkerDb: vi.fn(() => ({})), + findPepperByUserId: vi.fn(), + validateAndRedeemAccessCode: vi.fn(), +})); + import { accessGatewayRoutes } from './access-gateway'; import { signKiloToken } from '../auth/jwt'; import { deriveGatewayToken } from '../auth/gateway-token'; import { sandboxIdFromUserId } from '../auth/sandbox-id'; import { sandboxIdFromInstanceId } from '@kilocode/worker-utils/instance-id'; -import { KILOCLAW_AUTH_COOKIE, KILOCLAW_ACTIVE_INSTANCE_COOKIE } from '../config'; +import { + KILO_TOKEN_VERSION, + KILOCLAW_AUTH_COOKIE, + KILOCLAW_ACTIVE_INSTANCE_COOKIE, +} from '../config'; +import { KILOCLAW_AUDIENCE } from '@kilocode/worker-utils'; +import { SignJWT } from 'jose'; +import { findPepperByUserId, getWorkerDb, validateAndRedeemAccessCode } from '../db'; const NEXTAUTH_SECRET = 'test-nextauth-secret'; const GATEWAY_TOKEN_SECRET = 'test-gateway-secret'; @@ -50,6 +66,19 @@ async function signedAuthCookie(): Promise { }); } +async function signedAudienceAuthCookie(aud: string | string[]): Promise { + return new SignJWT({ + kiloUserId: USER_ID, + apiTokenPepper: null, + version: KILO_TOKEN_VERSION, + aud, + }) + .setProtectedHeader({ alg: 'HS256' }) + .setIssuedAt() + .setExpirationTime('1h') + .sign(new TextEncoder().encode(NEXTAUTH_SECRET)); +} + function parseSetCookies(response: Response): Record { const cookies: Record = {}; for (const line of response.headers.getSetCookie?.() ?? []) { @@ -66,6 +95,7 @@ function envBindings(overrides: Record = {}) { NEXTAUTH_SECRET, GATEWAY_TOKEN_SECRET, WORKER_ENV: 'test', + HYPERDRIVE: { connectionString: 'postgresql://fake' }, KILOCLAW_INSTANCE: buildInstanceBinding(USER_ID), KILOCLAW_INSTANCE_HOST_SUFFIX: '.kiloclaw.ai', KILOCLAW_INSTANCE_URL_SCHEME: 'https', @@ -74,6 +104,48 @@ function envBindings(overrides: Record = {}) { } describe('access-gateway cookie scoping', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('rejects a wrong-audience cookie without an access-code lookup', async () => { + const app = buildApp(); + const token = await signedAudienceAuthCookie('another-resource'); + + const response = await app.fetch( + new Request(`https://claw.kilosessions.ai/kilo-access-gateway?userId=${USER_ID}`, { + headers: { Cookie: `${KILOCLAW_AUTH_COOKIE}=${token}` }, + }), + envBindings() + ); + + expect(response.status).toBe(200); + expect(await response.text()).toContain('Enter the access code'); + expect(vi.mocked(getWorkerDb)).not.toHaveBeenCalled(); + expect(vi.mocked(findPepperByUserId)).not.toHaveBeenCalled(); + expect(vi.mocked(validateAndRedeemAccessCode)).not.toHaveBeenCalled(); + }); + + it('accepts a correct-audience cookie without a pepper database lookup', async () => { + const app = buildApp(); + const token = await signedAudienceAuthCookie([KILOCLAW_AUDIENCE]); + + const response = await app.fetch( + new Request(`https://claw.kilosessions.ai/kilo-access-gateway?userId=${USER_ID}`, { + headers: { Cookie: `${KILOCLAW_AUTH_COOKIE}=${token}` }, + }), + envBindings() + ); + + expect(response.status).toBe(302); + expect(response.headers.get('Location')).toBe( + `/#token=${await deriveGatewayToken(sandboxIdFromUserId(USER_ID), GATEWAY_TOKEN_SECRET)}` + ); + expect(vi.mocked(getWorkerDb)).not.toHaveBeenCalled(); + expect(vi.mocked(findPepperByUserId)).not.toHaveBeenCalled(); + expect(vi.mocked(validateAndRedeemAccessCode)).not.toHaveBeenCalled(); + }); + it('sets KILOCLAW_ACTIVE_INSTANCE_COOKIE on legacy host (claw.kilosessions.ai)', async () => { const app = buildApp(); const token = await signedAuthCookie(); diff --git a/services/notifications/src/__tests__/auth.test.ts b/services/notifications/src/__tests__/auth.test.ts index c8e847526f..a4de1e9452 100644 --- a/services/notifications/src/__tests__/auth.test.ts +++ b/services/notifications/src/__tests__/auth.test.ts @@ -1,8 +1,13 @@ -import { beforeEach, describe, it, expect, vi } from 'vitest'; import { Hono } from 'hono'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import { signKiloToken } from '@kilocode/worker-utils'; -import { authMiddleware } from '../auth'; -import type { AuthContext } from '../auth'; +import { + EVENT_SERVICE_AUDIENCE, + KILO_CHAT_AUDIENCE, + KILO_GATEWAY_AUDIENCE, + NOTIFICATIONS_AUDIENCE, +} from '@kilocode/worker-utils/internal-service-token-audiences'; +import { authMiddleware, type AuthContext } from '../auth'; type MockEnv = { NEXTAUTH_SECRET: { get: () => Promise }; @@ -11,164 +16,203 @@ type MockEnv = { }; const TEST_JWT_SECRET = 'test-secret-that-is-long-enough-for-hs256'; -const currentPepperByUserId = vi.hoisted(() => new Map()); +const dbState = vi.hoisted(() => ({ lookupCount: 0, fails: false })); +const userRow = vi.hoisted(() => ({ + pepper: 'pepper-current' as string | null, + blockedReason: null as string | null, +})); vi.mock('@kilocode/db/client', () => ({ getWorkerDb: () => ({ select: () => ({ from: () => ({ where: () => ({ - limit: async () => [ - { - api_token_pepper: currentPepperByUserId.get('user-xyz-789'), - blocked_reason: null, - }, - ], + limit: async () => { + dbState.lookupCount++; + if (dbState.fails) throw new Error('connection refused'); + return [{ api_token_pepper: userRow.pepper, blocked_reason: userRow.blockedReason }]; + }, }), }), }), }), })); -function makeApp(_env: MockEnv) { - const app = new Hono<{ Bindings: MockEnv; Variables: AuthContext }>(); - app.use('*', authMiddleware); - app.get('/test', c => c.json({ callerId: c.get('callerId'), callerKind: c.get('callerKind') })); - return app; -} - const defaultEnv: MockEnv = { NEXTAUTH_SECRET: { get: async () => TEST_JWT_SECRET }, HYPERDRIVE: { connectionString: 'postgres://test' }, WORKER_ENV: 'production', }; +function makeApp() { + let downstreamCalls = 0; + const app = new Hono<{ Bindings: MockEnv; Variables: AuthContext }>(); + app.use('*', authMiddleware); + app.get('/test', c => { + downstreamCalls++; + return c.json({ callerId: c.get('callerId'), callerKind: c.get('callerKind') }); + }); + return { app, downstreamCalls: () => downstreamCalls }; +} + +async function signToken( + params: { + audience?: string; + pepper?: string | null; + env?: string; + extra?: { tokenSource?: string; botId?: string; deviceSessionId?: string }; + } = {} +) { + return ( + await signKiloToken({ + userId: 'user-xyz-789', + pepper: 'pepper' in params ? params.pepper : 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'env' in params ? params.env : 'production', + audience: params.audience, + extra: params.extra ?? { tokenSource: 'kilo-chat' }, + }) + ).token; +} + +async function signWithAudience(aud: unknown): Promise { + const now = Math.floor(Date.now() / 1000); + const encode = (bytes: Uint8Array) => + btoa(String.fromCharCode(...bytes)) + .replace(/=/g, '') + .replace(/\+/g, '-') + .replace(/\//g, '_'); + const json = (value: unknown) => encode(new TextEncoder().encode(JSON.stringify(value))); + const input = `${json({ alg: 'HS256', typ: 'JWT' })}.${json({ + version: 3, + kiloUserId: 'user-xyz-789', + apiTokenPepper: 'pepper-current', + env: 'production', + aud, + iat: now, + exp: now + 3600, + })}`; + const key = await crypto.subtle.importKey( + 'raw', + new TextEncoder().encode(TEST_JWT_SECRET), + { name: 'HMAC', hash: 'SHA-256' }, + false, + ['sign'] + ); + const signature = await crypto.subtle.sign('HMAC', key, new TextEncoder().encode(input)); + return `${input}.${encode(new Uint8Array(signature))}`; +} + +async function request(token: string, env = defaultEnv) { + const testApp = makeApp(); + const response = await testApp.app.request( + '/test', + { headers: { authorization: `Bearer ${token}` } }, + env + ); + return { response, ...testApp }; +} + describe('authMiddleware', () => { beforeEach(() => { - currentPepperByUserId.set('user-xyz-789', 'pepper-current'); + dbState.lookupCount = 0; + dbState.fails = false; + userRow.pepper = 'pepper-current'; + userRow.blockedReason = null; }); it('returns 401 with no authorization header', async () => { - const res = await makeApp(defaultEnv).request('/test', {}, defaultEnv); - expect(res.status).toBe(401); - expect(await res.json()).toEqual({ error: 'Unauthorized' }); + const testApp = makeApp(); + const response = await testApp.app.request('/test', {}, defaultEnv); + expect(response.status).toBe(401); + expect(await response.json()).toEqual({ error: 'Unauthorized' }); }); - it('authenticates with a valid JWT and sets user identity', async () => { - const { token } = await signKiloToken({ - userId: 'user-xyz-789', - pepper: 'pepper-current', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'production', - extra: { tokenSource: 'kilo-chat' }, - }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv - ); - expect(res.status).toBe(200); - expect(await res.json()).toEqual({ - callerId: 'user-xyz-789', - callerKind: 'user', - }); + it.each([ + ['a matching string audience', () => signToken({ audience: NOTIFICATIONS_AUDIENCE })], + [ + 'a matching array audience', + () => signWithAudience([NOTIFICATIONS_AUDIENCE, 'other-service']), + ], + [ + 'a one-hour legacy kilo-chat token with bot and device-session claims', + () => + signToken({ + extra: { + tokenSource: 'kilo-chat', + botId: 'bot-123', + deviceSessionId: 'device-123', + }, + }), + ], + ])('authenticates %s', async (_name, createToken) => { + const { response } = await request(await createToken()); + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ callerId: 'user-xyz-789', callerKind: 'user' }); }); - it('authenticates a valid JWT from another token source', async () => { - const { token } = await signKiloToken({ - userId: 'user-xyz-789', - pepper: 'pepper-current', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'production', - extra: { tokenSource: 'cloud-agent' }, - }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv - ); - expect(res.status).toBe(200); - expect(await res.json()).toEqual({ - callerId: 'user-xyz-789', - callerKind: 'user', - }); + it.each([ + ['another batch', KILO_CHAT_AUDIENCE], + ['the gateway', KILO_GATEWAY_AUDIENCE], + ['another service', EVENT_SERVICE_AUDIENCE], + ])('rejects an audience for %s before database access', async (_name, audience) => { + const { response, downstreamCalls } = await request(await signToken({ audience })); + expect(response.status).toBe(401); + expect(dbState.lookupCount).toBe(0); + expect(downstreamCalls()).toBe(0); }); - it('returns 401 when the chat JWT has a stale pepper', async () => { - const { token } = await signKiloToken({ - userId: 'user-xyz-789', - pepper: 'pepper-stale', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'production', - extra: { tokenSource: 'kilo-chat' }, - }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv - ); - expect(res.status).toBe(401); - }); + it.each([null, '', ['notifications', 'notifications'], ['notifications', 1], []])( + 'rejects malformed audience claims before database access', + async audience => { + const { response, downstreamCalls } = await request(await signWithAudience(audience)); + expect(response.status).toBe(401); + expect(dbState.lookupCount).toBe(0); + expect(downstreamCalls()).toBe(0); + } + ); - it('returns 401 when the chat JWT was minted for a different environment', async () => { - const { token } = await signKiloToken({ - userId: 'user-xyz-789', - pepper: 'pepper-current', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'development', - extra: { tokenSource: 'kilo-chat' }, - }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv + it.each([ + ['a stale string pepper', 'pepper-stale', 'pepper-current', 401], + ['a null claim against a non-null pepper', null, 'pepper-current', 401], + ['a null claim against a null pepper', null, null, 200], + ])('preserves pepper semantics for %s', async (_name, pepper, currentPepper, status) => { + userRow.pepper = currentPepper; + const { response } = await request( + await signToken({ audience: NOTIFICATIONS_AUDIENCE, pepper }) ); - expect(res.status).toBe(401); + expect(response.status).toBe(status); }); - it('returns 401 with an expired JWT', async () => { + it('preserves absent pepper claim semantics', async () => { const { token } = await signKiloToken({ userId: 'user-xyz-789', - pepper: null, secret: TEST_JWT_SECRET, - expiresInSeconds: -1, + expiresInSeconds: 3600, env: 'production', - extra: { tokenSource: 'kilo-chat' }, + audience: NOTIFICATIONS_AUDIENCE, }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv - ); - expect(res.status).toBe(401); - expect(await res.json()).toEqual({ error: 'Unauthorized' }); + const { response } = await request(token); + expect(response.status).toBe(200); }); - it('returns 401 for an arbitrary non-JWT bearer', async () => { - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: 'Bearer not-a-jwt' } }, - defaultEnv + it.each([ + ['a blocked account', () => (userRow.blockedReason = 'manual block'), defaultEnv, 'production'], + ['a missing env claim', () => undefined, defaultEnv, undefined], + ['a mismatched environment', () => undefined, defaultEnv, 'development'], + ])('returns 401 for %s', async (_name, arrange, env, tokenEnv) => { + arrange(); + const { response } = await request( + await signToken({ audience: NOTIFICATIONS_AUDIENCE, env: tokenEnv }), + env ); - expect(res.status).toBe(401); + expect(response.status).toBe(401); }); - it('returns 401 when the JWT is signed with a different secret', async () => { - const { token } = await signKiloToken({ - userId: 'user-xyz-789', - pepper: null, - secret: 'a-completely-different-secret-of-correct-length', - expiresInSeconds: 3600, - }); - const res = await makeApp(defaultEnv).request( - '/test', - { headers: { authorization: `Bearer ${token}` } }, - defaultEnv - ); - expect(res.status).toBe(401); + it('maps dependency failures to 401', async () => { + dbState.fails = true; + const { response } = await request(await signToken({ audience: NOTIFICATIONS_AUDIENCE })); + expect(response.status).toBe(401); }); }); diff --git a/services/notifications/src/auth.ts b/services/notifications/src/auth.ts index 5abd95b1a0..acb6d14738 100644 --- a/services/notifications/src/auth.ts +++ b/services/notifications/src/auth.ts @@ -1,6 +1,7 @@ import { createMiddleware } from 'hono/factory'; import { verifyKiloBearerAgainstCurrentPepper } from '@kilocode/worker-utils/kilo-token-auth'; import { extractBearerToken } from '@kilocode/worker-utils'; +import { NOTIFICATIONS_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; import { logger } from './util/logger'; export type AuthContext = { @@ -30,6 +31,7 @@ export const authMiddleware = createMiddleware<{ nextAuthSecret: c.env.NEXTAUTH_SECRET, workerEnv: c.env.WORKER_ENV, connectionString: c.env.HYPERDRIVE.connectionString, + resourceAudience: { audience: NOTIFICATIONS_AUDIENCE, mode: 'allow-legacy' }, }); if (!auth) { return c.json({ error: 'Unauthorized' }, 401); From 20632b66b00de54b48e2b83ecf6c061de39be1db Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 1 Sep 2026 14:52:48 -0500 Subject: [PATCH 3/3] test(auth): restore middleware token rejection coverage --- services/kilo-chat/src/__tests__/auth.test.ts | 37 +++++++++++++++-- services/kiloclaw/src/auth/jwt.ts | 1 + .../notifications/src/__tests__/auth.test.ts | 40 +++++++++++++++++-- 3 files changed, 70 insertions(+), 8 deletions(-) diff --git a/services/kilo-chat/src/__tests__/auth.test.ts b/services/kilo-chat/src/__tests__/auth.test.ts index ab1d326453..700d5b4e9c 100644 --- a/services/kilo-chat/src/__tests__/auth.test.ts +++ b/services/kilo-chat/src/__tests__/auth.test.ts @@ -1,6 +1,7 @@ import { Hono } from 'hono'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { signKiloToken } from '@kilocode/worker-utils'; +import { getWorkerDb } from '@kilocode/db/client'; import { EVENT_SERVICE_AUDIENCE, KILO_CHAT_AUDIENCE, @@ -23,7 +24,7 @@ const userRow = vi.hoisted(() => ({ })); vi.mock('@kilocode/db/client', () => ({ - getWorkerDb: () => ({ + getWorkerDb: vi.fn(() => ({ select: () => ({ from: () => ({ where: () => ({ @@ -35,7 +36,7 @@ vi.mock('@kilocode/db/client', () => ({ }), }), }), - }), + })), })); const defaultEnv: MockEnv = { @@ -58,6 +59,8 @@ function makeApp() { async function signToken( params: { audience?: string; + secret?: string; + expiresInSeconds?: number; pepper?: string | null; env?: string; extra?: { tokenSource?: string; botId?: string; deviceSessionId?: string }; @@ -67,8 +70,8 @@ async function signToken( await signKiloToken({ userId: 'user-xyz-789', pepper: 'pepper' in params ? params.pepper : 'pepper-current', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, + secret: params.secret ?? TEST_JWT_SECRET, + expiresInSeconds: params.expiresInSeconds ?? 3600, env: 'env' in params ? params.env : 'production', audience: params.audience, extra: params.extra, @@ -116,6 +119,7 @@ async function request(token: string, env = defaultEnv) { describe('authMiddleware', () => { beforeEach(() => { + vi.mocked(getWorkerDb).mockClear(); dbState.lookupCount = 0; dbState.fails = false; userRow.pepper = 'pepper-current'; @@ -131,6 +135,10 @@ describe('authMiddleware', () => { it.each([ ['a matching string audience', () => signToken({ audience: KILO_CHAT_AUDIENCE })], + [ + 'a legacy token from another source', + () => signToken({ extra: { tokenSource: 'cloud-agent' } }), + ], ['a matching array audience', () => signWithAudience([KILO_CHAT_AUDIENCE, 'other-service'])], [ 'a legacy kilo-chat token with bot and device-session claims', @@ -149,6 +157,27 @@ describe('authMiddleware', () => { expect(await response.json()).toEqual({ callerId: 'user-xyz-789', callerKind: 'user' }); }); + it.each([ + ['an expired JWT', () => signToken({ audience: KILO_CHAT_AUDIENCE, expiresInSeconds: -60 })], + ['a malformed bearer', () => 'not-a-jwt'], + [ + 'a JWT signed with the wrong secret', + () => + signToken({ + audience: KILO_CHAT_AUDIENCE, + secret: 'wrong-test-secret-at-least-32-characters', + }), + ], + ])('returns 401 for %s before database or downstream access', async (_name, createToken) => { + const { response, downstreamCalls } = await request(await createToken()); + + expect(response.status).toBe(401); + expect(await response.json()).toEqual({ error: 'Unauthorized' }); + expect(getWorkerDb).not.toHaveBeenCalled(); + expect(dbState.lookupCount).toBe(0); + expect(downstreamCalls()).toBe(0); + }); + it.each([ ['another batch', NOTIFICATIONS_AUDIENCE], ['the gateway', KILO_GATEWAY_AUDIENCE], diff --git a/services/kiloclaw/src/auth/jwt.ts b/services/kiloclaw/src/auth/jwt.ts index 0643749778..fee7f10346 100644 --- a/services/kiloclaw/src/auth/jwt.ts +++ b/services/kiloclaw/src/auth/jwt.ts @@ -11,6 +11,7 @@ export type ValidateResult = * Verify a Kilo JWT using HS256 symmetric secret. * * Checks: signature, expiration (built into jose), version === 3 (via shared schema), + * KILOCLAW_AUDIENCE when an audience is present (legacy audience-less tokens are accepted), * and optional env match against the worker's WORKER_ENV. */ export async function validateKiloToken( diff --git a/services/notifications/src/__tests__/auth.test.ts b/services/notifications/src/__tests__/auth.test.ts index a4de1e9452..bcf2402069 100644 --- a/services/notifications/src/__tests__/auth.test.ts +++ b/services/notifications/src/__tests__/auth.test.ts @@ -1,6 +1,7 @@ import { Hono } from 'hono'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { signKiloToken } from '@kilocode/worker-utils'; +import { getWorkerDb } from '@kilocode/db/client'; import { EVENT_SERVICE_AUDIENCE, KILO_CHAT_AUDIENCE, @@ -23,7 +24,7 @@ const userRow = vi.hoisted(() => ({ })); vi.mock('@kilocode/db/client', () => ({ - getWorkerDb: () => ({ + getWorkerDb: vi.fn(() => ({ select: () => ({ from: () => ({ where: () => ({ @@ -35,7 +36,7 @@ vi.mock('@kilocode/db/client', () => ({ }), }), }), - }), + })), })); const defaultEnv: MockEnv = { @@ -58,6 +59,8 @@ function makeApp() { async function signToken( params: { audience?: string; + secret?: string; + expiresInSeconds?: number; pepper?: string | null; env?: string; extra?: { tokenSource?: string; botId?: string; deviceSessionId?: string }; @@ -67,8 +70,8 @@ async function signToken( await signKiloToken({ userId: 'user-xyz-789', pepper: 'pepper' in params ? params.pepper : 'pepper-current', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, + secret: params.secret ?? TEST_JWT_SECRET, + expiresInSeconds: params.expiresInSeconds ?? 3600, env: 'env' in params ? params.env : 'production', audience: params.audience, extra: params.extra ?? { tokenSource: 'kilo-chat' }, @@ -116,6 +119,7 @@ async function request(token: string, env = defaultEnv) { describe('authMiddleware', () => { beforeEach(() => { + vi.mocked(getWorkerDb).mockClear(); dbState.lookupCount = 0; dbState.fails = false; userRow.pepper = 'pepper-current'; @@ -131,6 +135,10 @@ describe('authMiddleware', () => { it.each([ ['a matching string audience', () => signToken({ audience: NOTIFICATIONS_AUDIENCE })], + [ + 'a legacy token from another source', + () => signToken({ extra: { tokenSource: 'cloud-agent' } }), + ], [ 'a matching array audience', () => signWithAudience([NOTIFICATIONS_AUDIENCE, 'other-service']), @@ -152,6 +160,30 @@ describe('authMiddleware', () => { expect(await response.json()).toEqual({ callerId: 'user-xyz-789', callerKind: 'user' }); }); + it.each([ + [ + 'an expired JWT', + () => signToken({ audience: NOTIFICATIONS_AUDIENCE, expiresInSeconds: -60 }), + ], + ['a malformed bearer', () => 'not-a-jwt'], + [ + 'a JWT signed with the wrong secret', + () => + signToken({ + audience: NOTIFICATIONS_AUDIENCE, + secret: 'wrong-test-secret-at-least-32-characters', + }), + ], + ])('returns 401 for %s before database or downstream access', async (_name, createToken) => { + const { response, downstreamCalls } = await request(await createToken()); + + expect(response.status).toBe(401); + expect(await response.json()).toEqual({ error: 'Unauthorized' }); + expect(getWorkerDb).not.toHaveBeenCalled(); + expect(dbState.lookupCount).toBe(0); + expect(downstreamCalls()).toBe(0); + }); + it.each([ ['another batch', KILO_CHAT_AUDIENCE], ['the gateway', KILO_GATEWAY_AUDIENCE],