diff --git a/apps/web/src/app/api/cloud-agent-next/balance/route.auth.test.ts b/apps/web/src/app/api/cloud-agent-next/balance/route.auth.test.ts new file mode 100644 index 0000000000..d9f92e1a2c --- /dev/null +++ b/apps/web/src/app/api/cloud-agent-next/balance/route.auth.test.ts @@ -0,0 +1,178 @@ +import jwt from 'jsonwebtoken'; +import { buildModernKiloTokenPayload } from '@kilocode/worker-utils/kilo-token-policy'; +import { GET } from './route'; +import { createControlTokenForRequest } from '@/lib/auth/resource-delegation'; +import type { User } from '@kilocode/db/schema'; +import { GET as getGenericBalance } from '@/app/api/profile/balance/route'; + +// Keep getUserFromAuth and validateAuthorizationHeader real; substitute only +// persistence, request context, and unrelated sign-in integrations. +const mockHeaders = jest.fn(); +const mockFindUser = jest.fn(); +const mockMembership = jest.fn(); +const mockBalance = jest.fn(); +jest.mock('next/headers', () => ({ headers: () => mockHeaders(), cookies: jest.fn() })); +jest.mock('next-auth', () => ({ + __esModule: true, + default: jest.fn(), + getServerSession: jest.fn(), +})); +jest.mock('@/lib/user', () => ({ findUserById: (...args: unknown[]) => mockFindUser(...args) })); +jest.mock('@/lib/drizzle', () => ({ + db: { query: { kilocode_users: { findFirst: (...args: unknown[]) => mockFindUser(...args) } } }, + readDb: {}, +})); +jest.mock('@/lib/organizations/organizations', () => ({ + isOrganizationMember: (...args: unknown[]) => mockMembership(...args), +})); +jest.mock('@/lib/organizations/organization-usage', () => ({ + getBalanceAndOrgSettings: (...args: unknown[]) => mockBalance(...args), +})); +jest.mock('@/lib/config.server', () => ({ + NEXTAUTH_SECRET: 'balance-test-secret', + BLACKLIST_TLDS: [], + isResourceTokenIssuanceEnabled: () => true, +})); +jest.mock('@/lib/constants', () => ({ ORGANIZATION_ID_HEADER: 'X-KiloCode-OrganizationId' })); +jest.mock('@/lib/dotenvx', () => ({ getEnvVariable: jest.fn() })); +jest.mock('@/lib/blacklist-domains-config', () => ({ getBlacklistedDomains: async () => [] })); +jest.mock('@/lib/utils.server', () => ({ + warnExceptInTest: jest.fn(), + sentryLogger: () => jest.fn(), +})); +jest.mock('@/lib/posthog', () => ({ __esModule: true, default: jest.fn() })); +jest.mock('@/lib/auth/magic-link-tokens', () => ({})); +jest.mock('@/lib/impact/debug', () => ({})); +jest.mock('@/lib/impact/referral', () => ({})); +jest.mock('@/lib/organizations/trial-utils', () => ({})); +jest.mock('@/lib/organizations/organization-seats', () => ({})); +jest.mock('@/lib/organizations/sales-demo', () => ({})); +jest.mock('@/lib/organizations/organization-sso-policy', () => ({})); +jest.mock('@/lib/organizations/verified-domain-membership', () => ({})); +jest.mock('@/lib/organizations/verified-domain-destination', () => ({})); +jest.mock('@/lib/account-linking-session', () => ({})); +jest.mock('@/lib/admin/admin-access-log', () => ({})); +jest.mock('@/lib/user/sso', () => ({})); +jest.mock('@/lib/web-session-revocation', () => ({})); + +const organizationId = '11111111-1111-4111-8111-111111111111'; +const user = { + id: 'user_123', + api_token_pepper: 'current-pepper', + google_user_email: 'test@example.com', +}; + +function controlToken(overrides: Record = {}, secret = 'balance-test-secret') { + const now = Math.floor(Date.now() / 1000); + // The same builder and claims used by createModernControlToken. + const payload = buildModernKiloTokenPayload({ + userId: user.id, + pepper: user.api_token_pepper, + env: process.env.NODE_ENV, + audience: 'cloud-agent-next', + issuedAt: now, + expiresAt: now + 60, + tokenPurpose: 'human-api', + credentialExchange: false, + extra: { + organizationId, + tokenSource: 'cloud-agent', + runtimeAdmission: { + source: 'user', + authorizationUserId: user.id, + authorizationPepper: user.api_token_pepper, + }, + }, + }); + return jwt.sign({ ...payload, ...overrides }, secret, { algorithm: 'HS256' }); +} + +beforeEach(() => { + jest.clearAllMocks(); + mockFindUser.mockResolvedValue(user); + mockMembership.mockResolvedValue(true); + mockBalance.mockResolvedValue({ balance: 12 }); + jest.spyOn(console, 'warn').mockImplementation(() => {}); +}); +afterEach(() => jest.restoreAllMocks()); + +function requestWith(token: string) { + mockHeaders.mockResolvedValue(new Headers({ Authorization: `Bearer ${token}` })); +} + +it.each(['human-api', 'device-access'])( + 'accepts issuer-shaped %s control only at the dedicated endpoint', + async tokenPurpose => { + requestWith( + controlToken({ + tokenPurpose, + ...(tokenPurpose === 'device-access' ? { deviceSessionId: 'device_123' } : {}), + }) + ); + expect((await GET()).status).toBe(200); + expect(mockMembership).toHaveBeenCalledWith(organizationId, user.id, expect.anything()); + expect(mockBalance).toHaveBeenCalledWith(organizationId, user); + mockBalance.mockClear(); + expect((await getGenericBalance()).status).toBe(401); + expect(mockBalance).not.toHaveBeenCalled(); + } +); + +it.each([ + { name: 'wrong audience', overrides: { aud: 'kilo-api' }, secret: 'balance-test-secret' }, + { name: 'wrong signature', overrides: {}, secret: 'wrong-secret' }, + { + name: 'wrong pepper', + overrides: { apiTokenPepper: 'stale-pepper' }, + secret: 'balance-test-secret', + }, + { name: 'expired', overrides: { exp: 1 }, secret: 'balance-test-secret' }, +])('rejects $name before balance access', async ({ overrides, secret }) => { + requestWith(controlToken(overrides, secret)); + expect((await GET()).status).toBe(401); + expect(mockBalance).not.toHaveBeenCalled(); +}); + +it('rejects an organization the user cannot access', async () => { + mockMembership.mockResolvedValue(false); + requestWith(controlToken()); + expect((await GET()).status).toBe(403); + expect(mockBalance).not.toHaveBeenCalled(); +}); + +it('rejects malformed bearer input', async () => { + requestWith('not-a-jwt'); + expect((await GET()).status).toBe(401); + expect(mockBalance).not.toHaveBeenCalled(); +}); + +it('accepts the actual control issuer output with its environment and pepper claims', async () => { + const now = Math.floor(Date.now() / 1000); + const principal = jwt.sign( + buildModernKiloTokenPayload({ + userId: user.id, + pepper: user.api_token_pepper, + env: process.env.NODE_ENV, + audience: 'kilo-api', + issuedAt: now, + expiresAt: now + 60, + tokenPurpose: 'human-api', + credentialExchange: true, + }), + 'balance-test-secret', + { algorithm: 'HS256' } + ); + const issued = await createControlTokenForRequest(user as User, 'cloud-agent-next', { + headers: new Headers({ Authorization: `Bearer ${principal}` }), + tokenSource: 'cloud-agent', + }); + expect(jwt.verify(issued.token, 'balance-test-secret')).toMatchObject({ + aud: 'cloud-agent-next', + env: process.env.NODE_ENV, + apiTokenPepper: user.api_token_pepper, + credentialExchange: false, + }); + requestWith(issued.token); + expect((await GET()).status).toBe(200); + expect((await getGenericBalance()).status).toBe(401); +}); diff --git a/apps/web/src/app/api/cloud-agent-next/balance/route.test.ts b/apps/web/src/app/api/cloud-agent-next/balance/route.test.ts new file mode 100644 index 0000000000..87e8c64fa7 --- /dev/null +++ b/apps/web/src/app/api/cloud-agent-next/balance/route.test.ts @@ -0,0 +1,41 @@ +import { GET } from './route'; +import { getUserFromAuth } from '@/lib/user/server'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { NextResponse } from 'next/server'; + +jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() })); +jest.mock('@/lib/organizations/organization-usage', () => ({ + getBalanceAndOrgSettings: jest.fn(), +})); + +const auth = jest.mocked(getUserFromAuth); +const balance = jest.mocked(getBalanceAndOrgSettings); + +afterEach(() => jest.resetAllMocks()); + +it('requires Cloud Agent audience authentication and uses the authorized organization', async () => { + const user = { id: 'user_123' }; + auth.mockResolvedValue({ user, organizationId: 'org_123' } as Awaited< + ReturnType + >); + balance.mockResolvedValue({ balance: 12 } as Awaited< + ReturnType + >); + + const response = await GET(); + + expect(auth).toHaveBeenCalledWith({ adminOnly: false, expectedAudience: 'cloud-agent-next' }); + expect(balance).toHaveBeenCalledWith('org_123', user); + expect(await response.json()).toEqual({ balance: 12, isDepleted: false }); +}); + +it.each([401, 403])( + 'does not read balances when authentication or organization access fails (%s)', + async status => { + const authFailedResponse = NextResponse.json({ error: 'Unauthorized' }, { status }); + auth.mockResolvedValue({ authFailedResponse } as Awaited>); + + expect(await GET()).toBe(authFailedResponse); + expect(balance).not.toHaveBeenCalled(); + } +); diff --git a/apps/web/src/app/api/cloud-agent-next/balance/route.ts b/apps/web/src/app/api/cloud-agent-next/balance/route.ts new file mode 100644 index 0000000000..6197c5b076 --- /dev/null +++ b/apps/web/src/app/api/cloud-agent-next/balance/route.ts @@ -0,0 +1,19 @@ +import { CLOUD_AGENT_NEXT_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { getUserFromAuth } from '@/lib/user/server'; +import { NextResponse } from 'next/server'; + +export async function GET(): Promise< + NextResponse<{ error: string } | { balance: number; isDepleted: boolean }> +> { + const { user, authFailedResponse, organizationId } = await getUserFromAuth({ + adminOnly: false, + expectedAudience: CLOUD_AGENT_NEXT_AUDIENCE, + }); + + if (authFailedResponse) return authFailedResponse; + + const { balance } = await getBalanceAndOrgSettings(organizationId, user); + + return NextResponse.json({ balance, isDepleted: balance <= 0 }); +} diff --git a/apps/web/src/lib/auth/cloud-agent-workflow-user.test.ts b/apps/web/src/lib/auth/cloud-agent-workflow-user.test.ts new file mode 100644 index 0000000000..6f7712716d --- /dev/null +++ b/apps/web/src/lib/auth/cloud-agent-workflow-user.test.ts @@ -0,0 +1,76 @@ +import jwt from 'jsonwebtoken'; +import { eq } from 'drizzle-orm'; +import { db } from '@/lib/drizzle'; +import { kilocode_users } from '@kilocode/db/schema'; +import { insertTestUser } from '@/tests/helpers/user.helper'; +import { generateCloudAgentWorkflowToken } from '@/lib/tokens'; +import { prepareCloudAgentWorkflowUser } from './cloud-agent-workflow-user'; + +const issuance = { enabled: true }; +jest.mock('@/lib/config.server', () => ({ + ...jest.requireActual('@/lib/config.server'), + isResourceTokenIssuanceEnabled: () => issuance.enabled, +})); + +beforeEach(() => { + issuance.enabled = true; +}); + +test('initializes a null pepper and issues modern review admission with the persisted value', async () => { + const user = await insertTestUser({ api_token_pepper: null }); + expect(() => + generateCloudAgentWorkflowToken(user, { tokenSource: 'code-review', expiresIn: 3600 }) + ).toThrow('current user pepper'); + const prepared = await prepareCloudAgentWorkflowUser(user); + const [persisted] = await db.select().from(kilocode_users).where(eq(kilocode_users.id, user.id)); + expect(prepared.api_token_pepper).toEqual(expect.any(String)); + expect(prepared.api_token_pepper).toBe(persisted.api_token_pepper); + const claims = jwt.decode( + generateCloudAgentWorkflowToken(prepared, { tokenSource: 'code-review', expiresIn: 3600 }) + ); + expect(claims).toMatchObject({ + aud: 'cloud-agent-next', + tokenPurpose: 'internal-service', + credentialExchange: false, + apiTokenPepper: persisted.api_token_pepper, + runtimeAdmission: { + source: 'automation', + authorizationUserId: user.id, + authorizationPepper: persisted.api_token_pepper, + }, + }); +}); + +test('concurrent initializers return the same persisted pepper', async () => { + const user = await insertTestUser({ api_token_pepper: null }); + const results = await Promise.all( + Array.from({ length: 8 }, () => prepareCloudAgentWorkflowUser(user)) + ); + expect(new Set(results.map(result => result.api_token_pepper)).size).toBe(1); + expect(results[0].api_token_pepper).toEqual(expect.any(String)); +}); + +test('preserves a pepper assigned after the user snapshot was loaded', async () => { + const user = await insertTestUser({ api_token_pepper: null }); + await db + .update(kilocode_users) + .set({ api_token_pepper: 'concurrent-rotation' }) + .where(eq(kilocode_users.id, user.id)); + expect((await prepareCloudAgentWorkflowUser(user)).api_token_pepper).toBe('concurrent-rotation'); +}); + +test('preserves existing peppers and leaves legacy issuance unchanged', async () => { + const existing = await insertTestUser({ api_token_pepper: 'existing-pepper' }); + expect(await prepareCloudAgentWorkflowUser(existing)).toBe(existing); + issuance.enabled = false; + const user = await insertTestUser({ api_token_pepper: null }); + expect(await prepareCloudAgentWorkflowUser(user)).toBe(user); + const [persisted] = await db.select().from(kilocode_users).where(eq(kilocode_users.id, user.id)); + expect(persisted.api_token_pepper).toBeNull(); +}); + +test('fails closed if the user was deleted', async () => { + const user = await insertTestUser({ api_token_pepper: null }); + await db.delete(kilocode_users).where(eq(kilocode_users.id, user.id)); + await expect(prepareCloudAgentWorkflowUser(user)).rejects.toThrow('not found'); +}); diff --git a/apps/web/src/lib/auth/cloud-agent-workflow-user.ts b/apps/web/src/lib/auth/cloud-agent-workflow-user.ts new file mode 100644 index 0000000000..ad18b69078 --- /dev/null +++ b/apps/web/src/lib/auth/cloud-agent-workflow-user.ts @@ -0,0 +1,24 @@ +import { randomUUID } from 'node:crypto'; +import { eq, sql } from 'drizzle-orm'; +import { kilocode_users, type User } from '@kilocode/db/schema'; +import { db } from '@/lib/drizzle'; +import { isResourceTokenIssuanceEnabled, type ResourceTokenFamily } from '@/lib/config.server'; + +export async function prepareCloudAgentWorkflowUser( + user: User, + requiredFamilies: readonly ResourceTokenFamily[] = ['cloud-agent-next'] +): Promise { + if (!requiredFamilies.some(isResourceTokenIssuanceEnabled) || user.api_token_pepper !== null) { + return user; + } + + // Use the primary's persisted value: another issuer or revocation may have + // assigned a pepper since this user was loaded. Never overwrite that value. + const [currentUser] = await db + .update(kilocode_users) + .set({ api_token_pepper: sql`COALESCE(${kilocode_users.api_token_pepper}, ${randomUUID()})` }) + .where(eq(kilocode_users.id, user.id)) + .returning(); + if (!currentUser) throw new Error(`User ${user.id} not found`); + return currentUser; +} diff --git a/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.test.ts b/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.test.ts index 03d7382171..4b3ddf04da 100644 --- a/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.test.ts +++ b/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.test.ts @@ -1,17 +1,15 @@ const mockGetFixTicketById = jest.fn(); -jest.mock('@/lib/drizzle', () => ({ - db: { - select: () => ({ - from: () => ({ - where: () => ({ - limit: () => [{ id: 'user-1', api_token_pepper: 'pepper' }], - }), - }), - }), - }, +jest.mock('@/lib/config.server', () => ({ + ...jest.requireActual('@/lib/config.server'), + isResourceTokenIssuanceEnabled: () => true, })); +import { insertTestUser } from '@/tests/helpers/user.helper'; +import { db } from '@/lib/drizzle'; +import { kilocode_users } from '@kilocode/db/schema'; +import { eq } from 'drizzle-orm'; + jest.mock('@/lib/tokens', () => ({ generateCloudAgentWorkflowToken: jest.fn(() => 'workflow-token'), TOKEN_EXPIRY: { default: 3600 }, @@ -30,7 +28,12 @@ const mockGenerateCloudAgentWorkflowToken = jest.mocked(generateCloudAgentWorkfl const organizationId = 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa'; -beforeEach(() => { +afterEach(async () => { + await db.delete(kilocode_users).where(eq(kilocode_users.id, 'user-1')); +}); + +beforeEach(async () => { + await insertTestUser({ id: 'user-1', api_token_pepper: null }); jest.clearAllMocks(); mockGetFixTicketById.mockResolvedValue({ repo_full_name: 'kilo/repo', @@ -56,8 +59,13 @@ describe('prepareFixPayload workflow token ownership', () => { }, }); + const [persisted] = await db + .select() + .from(kilocode_users) + .where(eq(kilocode_users.id, 'user-1')); + expect(persisted.api_token_pepper).toEqual(expect.any(String)); expect(mockGenerateCloudAgentWorkflowToken).toHaveBeenCalledWith( - expect.objectContaining({ id: 'user-1' }), + expect.objectContaining({ id: 'user-1', api_token_pepper: persisted.api_token_pepper }), expect.objectContaining({ organizationId: expectedOrganizationId }) ); } diff --git a/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.ts b/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.ts index 4260a507e1..6930d4e861 100644 --- a/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.ts +++ b/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.ts @@ -5,6 +5,7 @@ * Returns complete payload ready for cloud agent */ +import { prepareCloudAgentWorkflowUser } from '@/lib/auth/cloud-agent-workflow-user'; import { captureException } from '@sentry/nextjs'; import { db } from '@/lib/drizzle'; import { kilocode_users } from '@kilocode/db/schema'; @@ -51,7 +52,7 @@ export async function prepareFixPayload(params: PreparePayloadParams): Promise ({ - db: { - select: () => ({ - from: () => ({ - where: () => ({ - limit: () => [{ id: 'user-1', api_token_pepper: 'pepper' }], - }), - }), - }), - }, +jest.mock('@/lib/config.server', () => ({ + ...jest.requireActual('@/lib/config.server'), + isResourceTokenIssuanceEnabled: () => true, })); +import { insertTestUser } from '@/tests/helpers/user.helper'; +import { db } from '@/lib/drizzle'; +import { kilocode_users } from '@kilocode/db/schema'; +import { eq } from 'drizzle-orm'; + jest.mock('@/lib/tokens', () => ({ generateCloudAgentWorkflowToken: jest.fn(() => 'workflow-token'), TOKEN_EXPIRY: { default: 3600 }, @@ -30,7 +28,12 @@ const mockGenerateCloudAgentWorkflowToken = jest.mocked(generateCloudAgentWorkfl const organizationId = 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa'; -beforeEach(() => { +afterEach(async () => { + await db.delete(kilocode_users).where(eq(kilocode_users.id, 'user-1')); +}); + +beforeEach(async () => { + await insertTestUser({ id: 'user-1', api_token_pepper: null }); jest.clearAllMocks(); mockGetTriageTicketById.mockResolvedValue({ repo_full_name: 'kilo/repo', @@ -55,8 +58,13 @@ describe('prepareTriagePayload workflow token ownership', () => { }, }); + const [persisted] = await db + .select() + .from(kilocode_users) + .where(eq(kilocode_users.id, 'user-1')); + expect(persisted.api_token_pepper).toEqual(expect.any(String)); expect(mockGenerateCloudAgentWorkflowToken).toHaveBeenCalledWith( - expect.objectContaining({ id: 'user-1' }), + expect.objectContaining({ id: 'user-1', api_token_pepper: persisted.api_token_pepper }), expect.objectContaining({ organizationId: expectedOrganizationId }) ); } diff --git a/apps/web/src/lib/auto-triage/triggers/prepare-triage-payload.ts b/apps/web/src/lib/auto-triage/triggers/prepare-triage-payload.ts index 7eecb6903e..18675be87f 100644 --- a/apps/web/src/lib/auto-triage/triggers/prepare-triage-payload.ts +++ b/apps/web/src/lib/auto-triage/triggers/prepare-triage-payload.ts @@ -5,6 +5,7 @@ * Returns complete payload ready for cloud agent */ +import { prepareCloudAgentWorkflowUser } from '@/lib/auth/cloud-agent-workflow-user'; import { captureException } from '@sentry/nextjs'; import { db } from '@/lib/drizzle'; import { kilocode_users } from '@kilocode/db/schema'; @@ -53,7 +54,7 @@ export async function prepareTriagePayload( } // 3. Generate auth token for cloud agent with bot identifier - const authToken = generateCloudAgentWorkflowToken(user, { + const authToken = generateCloudAgentWorkflowToken(await prepareCloudAgentWorkflowUser(user), { organizationId: owner.type === 'org' ? owner.id : undefined, tokenSource: 'auto-triage', botId: 'auto-triage', diff --git a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts index 2ac2738090..6bd773a3a1 100644 --- a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts +++ b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts @@ -77,6 +77,7 @@ import { } from '@kilocode/worker-utils/bitbucket-workspace-access-token'; import { getGitHubPullRequestCheckoutRef } from '@/lib/integrations/platforms/github/webhook-handlers/pull-request-checkout-ref'; import { getManualCodeReviewConfig } from '../manual-config'; +import { prepareCloudAgentWorkflowUser } from '@/lib/auth/cloud-agent-workflow-user'; const BitbucketWorkspaceSlugSchema = z.string().regex(/^[a-z0-9][a-z0-9_.-]*$/); const BitbucketRepositorySlugSchema = z.string().regex(/^[A-Za-z0-9_.-]+$/); @@ -298,12 +299,15 @@ export async function prepareReviewPayload( expectedHeadSha: expectedHeadSha.data, } ); - const authToken = generateCloudAgentWorkflowToken(user, { - organizationId: owner.type === 'org' ? owner.id : undefined, - tokenSource: 'code-review', - botId: 'reviewer', - expiresIn: TOKEN_EXPIRY.default, - }); + const authToken = generateCloudAgentWorkflowToken( + await prepareCloudAgentWorkflowUser(user), + { + organizationId: owner.type === 'org' ? owner.id : undefined, + tokenSource: 'code-review', + botId: 'reviewer', + expiresIn: TOKEN_EXPIRY.default, + } + ); // Single source for the standard reviewer's model so the session input and the // forward-shaped `reviewAgents[0]` can never drift apart. const standardModel = config.model_slug || DEFAULT_CODE_REVIEW_MODEL; @@ -717,7 +721,7 @@ export async function prepareReviewPayload( ]); // 5. Generate auth token for cloud agent with bot identifier - const authToken = generateCloudAgentWorkflowToken(user, { + const authToken = generateCloudAgentWorkflowToken(await prepareCloudAgentWorkflowUser(user), { organizationId: owner.type === 'org' ? owner.id : undefined, tokenSource: 'code-review', botId: 'reviewer', diff --git a/apps/web/src/lib/security-agent/services/analysis-service.resource-tokens.test.ts b/apps/web/src/lib/security-agent/services/analysis-service.resource-tokens.test.ts new file mode 100644 index 0000000000..5283bab66f --- /dev/null +++ b/apps/web/src/lib/security-agent/services/analysis-service.resource-tokens.test.ts @@ -0,0 +1,156 @@ +import jwt from 'jsonwebtoken'; +import { eq } from 'drizzle-orm'; +import { db } from '@/lib/drizzle'; +import { kilocode_users, type SecurityFinding } from '@kilocode/db/schema'; +import { insertTestUser } from '@/tests/helpers/user.helper'; +import { NEXTAUTH_SECRET } from '@/lib/config.server'; +import { generateWorkflowGatewayToken } from '@/lib/tokens'; +import { prepareCloudAgentWorkflowUser } from '@/lib/auth/cloud-agent-workflow-user'; +import { getSecurityFindingById } from '../db/security-findings'; +import { triageSecurityFinding } from './triage-service'; +import { createCloudAgentNextClient } from '@/lib/cloud-agent-next/cloud-agent-client'; +import { startSecurityAnalysis } from './analysis-service'; + +jest.mock('../db/security-findings', () => ({ getSecurityFindingById: jest.fn() })); +jest.mock('../db/security-analysis', () => ({ + tryAcquireAnalysisStartLease: jest.fn(async () => true), + updateAnalysisStatus: jest.fn(async () => true), + clearAnalysisStatus: jest.fn(), +})); +jest.mock('./triage-service', () => ({ triageSecurityFinding: jest.fn() })); +jest.mock('./extraction-service', () => ({ extractSandboxAnalysis: jest.fn() })); +jest.mock('./auto-dismiss-service', () => ({ maybeAutoDismissAnalysis: jest.fn() })); +jest.mock('../posthog-tracking', () => ({ + trackSecurityAgentAnalysisStarted: jest.fn(), + trackSecurityAgentAnalysisCompleted: jest.fn(), +})); +jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ + createCloudAgentNextClient: jest.fn(() => ({ + prepareSession: jest.fn(async () => ({ + cloudAgentSessionId: 'session', + kiloSessionId: 'kilo', + })), + initiateFromPreparedSession: jest.fn(async () => ({})), + })), + InsufficientCreditsError: class extends Error {}, +})); + +const flags = [ + 'SHARED_RESOURCE_TOKENS_ENABLED', + 'CLOUD_AGENT_RESOURCE_TOKENS_ENABLED', + 'WORKFLOW_GATEWAY_RESOURCE_TOKENS_ENABLED', +] as const; +const saved = flags.map(key => [key, process.env[key]] as const); +beforeEach(() => { + jest.clearAllMocks(); + process.env.SHARED_RESOURCE_TOKENS_ENABLED = 'true'; + process.env.CLOUD_AGENT_RESOURCE_TOKENS_ENABLED = 'false'; + process.env.WORKFLOW_GATEWAY_RESOURCE_TOKENS_ENABLED = 'false'; +}); +afterEach(() => { + for (const [key, value] of saved) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } +}); + +it.each([ + [false, true], + [false, false], + [true, false], +])( + 'starts security analysis with cloud=%s gateway=%s and a null pepper', + async (cloud, gateway) => { + process.env.CLOUD_AGENT_RESOURCE_TOKENS_ENABLED = String(cloud); + process.env.WORKFLOW_GATEWAY_RESOURCE_TOKENS_ENABLED = String(gateway); + const user = await insertTestUser({ api_token_pepper: null }); + if (gateway) { + expect(() => generateWorkflowGatewayToken(user, { tokenSource: 'security-agent' })).toThrow( + 'Workflow gateway tokens require a current user pepper' + ); + } + jest.mocked(getSecurityFindingById).mockResolvedValue({ + id: 'finding', + owned_by_user_id: user.id, + owned_by_organization_id: null, + status: 'open', + source: 'dependabot', + source_id: '42', + repo_full_name: 'acme/repo', + package_name: 'lodash', + package_ecosystem: 'npm', + severity: 'high', + title: 'Vulnerability', + description: 'Test finding', + cwe_ids: [], + raw_data: null, + analysis: null, + analysis_completed_at: null, + analysis_error: null, + analysis_started_at: null, + analysis_status: 'new', + cli_session_id: null, + created_at: new Date().toISOString(), + updated_at: new Date().toISOString(), + first_detected_at: new Date().toISOString(), + last_synced_at: new Date().toISOString(), + cve_id: null, + cvss_score: null, + dependabot_html_url: null, + dependency_scope: 'runtime', + fixed_at: null, + ghsa_id: null, + ignored_by: null, + ignored_reason: null, + manifest_path: 'package.json', + patched_version: null, + platform_integration_id: null, + session_id: null, + sla_due_at: null, + vulnerable_version_range: null, + } satisfies SecurityFinding); + jest.mocked(triageSecurityFinding).mockResolvedValue({ + needsSandboxAnalysis: true, + needsSandboxReasoning: 'Inspect runtime dependency', + suggestedAction: 'analyze_codebase', + confidence: 'high', + triageAt: new Date().toISOString(), + }); + const result = await startSecurityAnalysis({ + findingId: 'finding', + user, + githubRepo: 'acme/repo', + githubToken: 'test-github-token', + }); + expect(result).toMatchObject({ started: true, triageOnly: false }); + const [persisted] = await db + .select() + .from(kilocode_users) + .where(eq(kilocode_users.id, user.id)); + expect(persisted.api_token_pepper).toEqual(cloud || gateway ? expect.any(String) : null); + const gatewayClaims = jwt.verify( + jest.mocked(triageSecurityFinding).mock.calls[0][0].authToken, + NEXTAUTH_SECRET + ) as jwt.JwtPayload; + const cloudClaims = jwt.verify( + jest.mocked(createCloudAgentNextClient).mock.calls[0][0], + NEXTAUTH_SECRET + ) as jwt.JwtPayload; + expect(gatewayClaims.aud).toBe(gateway ? 'kilo-gateway' : undefined); + expect(gatewayClaims.tokenPurpose).toBe(gateway ? 'delegated-workload' : undefined); + expect(cloudClaims.aud).toBe(cloud ? 'cloud-agent-next' : undefined); + expect(cloudClaims.tokenPurpose).toBe(cloud ? 'internal-service' : undefined); + if (gateway) { + expect(gatewayClaims.apiTokenPepper).toBe(persisted.api_token_pepper); + expect(gatewayClaims.exp! - gatewayClaims.iat!).toBe(3600); + } + } +); + +it('does not initialize cloud-only callers when only the gateway gate is enabled', async () => { + process.env.WORKFLOW_GATEWAY_RESOURCE_TOKENS_ENABLED = 'true'; + const user = await insertTestUser({ api_token_pepper: null }); + expect(await prepareCloudAgentWorkflowUser(user)).toBe(user); + const [persisted] = await db.select().from(kilocode_users).where(eq(kilocode_users.id, user.id)); + expect(persisted.api_token_pepper).toBeNull(); +}); diff --git a/apps/web/src/lib/security-agent/services/analysis-service.test.ts b/apps/web/src/lib/security-agent/services/analysis-service.test.ts index 1166278a1b..ada91de427 100644 --- a/apps/web/src/lib/security-agent/services/analysis-service.test.ts +++ b/apps/web/src/lib/security-agent/services/analysis-service.test.ts @@ -55,6 +55,7 @@ jest.mock('@/lib/security-agent/db/security-analysis', () => ({ jest.mock('@/lib/config.server', () => ({ CALLBACK_TOKEN_SECRET: 'test-callback-token-secret', + isResourceTokenIssuanceEnabled: () => false, })); jest.mock('./triage-service', () => ({ diff --git a/apps/web/src/lib/security-agent/services/analysis-service.token-source.test.ts b/apps/web/src/lib/security-agent/services/analysis-service.token-source.test.ts index 4ac31a1071..63569f7528 100644 --- a/apps/web/src/lib/security-agent/services/analysis-service.token-source.test.ts +++ b/apps/web/src/lib/security-agent/services/analysis-service.token-source.test.ts @@ -1,3 +1,6 @@ +import { db } from '@/lib/drizzle'; +import { kilocode_users } from '@kilocode/db/schema'; +import { eq } from 'drizzle-orm'; import { beforeAll, beforeEach, describe, expect, it, jest } from '@jest/globals'; import jwt from 'jsonwebtoken'; import type { SecurityFinding, User } from '@kilocode/db/schema'; @@ -161,59 +164,76 @@ describe('startSecurityAnalysis token source', () => { }); }); - it('uses separate modern gateway and sandbox security-agent credentials', async () => { - const user = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); - const organizationId = crypto.randomUUID(); - const finding = { - ...createFinding(user), - owned_by_organization_id: organizationId, - owned_by_user_id: null, - }; - mockGetSecurityFindingById.mockResolvedValue(finding); + it.each([null, 'existing-pepper'])( + 'uses persisted pepper %s for modern gateway and sandbox credentials', + async api_token_pepper => { + const user = await insertTestUser({ api_token_pepper }); + const organizationId = crypto.randomUUID(); + const finding = { + ...createFinding(user), + owned_by_organization_id: organizationId, + owned_by_user_id: null, + }; + mockGetSecurityFindingById.mockResolvedValue(finding); - const result = await startSecurityAnalysis({ - findingId: finding.id, - user, - githubRepo: finding.repo_full_name, - githubToken: 'github-token', - organizationId, - }); + const result = await startSecurityAnalysis({ + findingId: finding.id, + user, + githubRepo: finding.repo_full_name, + githubToken: 'github-token', + organizationId, + }); - expect(result).toEqual({ started: true, triageOnly: false }); - const triageInput = mockTriageSecurityFinding.mock.calls[0]?.[0]; - const cloudAgentToken = mockCreateCloudAgentNextClient.mock.calls[0]?.[0]; - if (!triageInput) throw new Error('Expected triage to receive an input'); - expect(triageInput.authToken).toEqual(expect.any(String)); - const gatewayClaims = jwt.verify(triageInput.authToken, tokenSecret) as jwt.JwtPayload; - const cloudAgentClaims = jwt.decode(cloudAgentToken); - if (!cloudAgentClaims || typeof cloudAgentClaims === 'string') { - throw new Error('Expected sandbox JWT claims'); + expect(result).toEqual({ started: true, triageOnly: false }); + const [persisted] = await db + .select() + .from(kilocode_users) + .where(eq(kilocode_users.id, user.id)); + expect(persisted.api_token_pepper).toEqual(expect.any(String)); + if (api_token_pepper !== null) expect(persisted.api_token_pepper).toBe(api_token_pepper); + const triageInput = mockTriageSecurityFinding.mock.calls[0]?.[0]; + const cloudAgentToken = mockCreateCloudAgentNextClient.mock.calls[0]?.[0]; + if (!triageInput) throw new Error('Expected triage to receive an input'); + expect(triageInput.authToken).toEqual(expect.any(String)); + const gatewayClaims = jwt.verify(triageInput.authToken, tokenSecret) as jwt.JwtPayload; + const cloudAgentClaims = jwt.decode(cloudAgentToken); + if (!cloudAgentClaims || typeof cloudAgentClaims === 'string') { + throw new Error('Expected sandbox JWT claims'); + } + expect(gatewayClaims).toMatchObject({ + aud: 'kilo-gateway', + apiTokenPepper: persisted.api_token_pepper, + tokenPurpose: 'delegated-workload', + credentialExchange: false, + organizationId, + tokenSource: 'security-agent', + }); + expect(gatewayClaims.exp! - gatewayClaims.iat!).toBeLessThanOrEqual(60 * 60); + expect(cloudAgentClaims).toMatchObject({ + aud: 'cloud-agent-next', + apiTokenPepper: persisted.api_token_pepper, + runtimeAdmission: { + authorizationUserId: user.id, + authorizationPepper: persisted.api_token_pepper, + }, + tokenPurpose: 'internal-service', + credentialExchange: false, + organizationId, + }); + expect(cloudAgentToken).not.toBe(triageInput.authToken); + const tokenPolicy = await verifyKiloTokenForPolicy(triageInput.authToken, tokenSecret, { + audience: 'kilo-gateway', + mode: 'required', + }); + expect(isKiloCredentialExchangeEligible(tokenPolicy, { legacy: 'five-year-api' })).toBe( + false + ); + expect(mockPrepareSession).toHaveBeenCalledTimes(1); + expect(mockInitiateFromPreparedSession).toHaveBeenCalledWith({ + cloudAgentSessionId: 'agent-session-123', + }); } - expect(gatewayClaims).toMatchObject({ - aud: 'kilo-gateway', - tokenPurpose: 'delegated-workload', - credentialExchange: false, - organizationId, - tokenSource: 'security-agent', - }); - expect(gatewayClaims.exp! - gatewayClaims.iat!).toBeLessThanOrEqual(60 * 60); - expect(cloudAgentClaims).toMatchObject({ - aud: 'cloud-agent-next', - tokenPurpose: 'internal-service', - credentialExchange: false, - organizationId, - }); - expect(cloudAgentToken).not.toBe(triageInput.authToken); - const tokenPolicy = await verifyKiloTokenForPolicy(triageInput.authToken, tokenSecret, { - audience: 'kilo-gateway', - mode: 'required', - }); - expect(isKiloCredentialExchangeEligible(tokenPolicy, { legacy: 'five-year-api' })).toBe(false); - expect(mockPrepareSession).toHaveBeenCalledTimes(1); - expect(mockInitiateFromPreparedSession).toHaveBeenCalledWith({ - cloudAgentSessionId: 'agent-session-123', - }); - }); + ); it('preserves the legacy gateway token shape when shared issuance is disabled', async () => { shared.enabled = false; diff --git a/apps/web/src/lib/security-agent/services/analysis-service.ts b/apps/web/src/lib/security-agent/services/analysis-service.ts index 123673e806..ee2173e4f9 100644 --- a/apps/web/src/lib/security-agent/services/analysis-service.ts +++ b/apps/web/src/lib/security-agent/services/analysis-service.ts @@ -1,4 +1,5 @@ import 'server-only'; +import { prepareCloudAgentWorkflowUser } from '@/lib/auth/cloud-agent-workflow-user'; import { randomUUID } from 'crypto'; import { createCloudAgentNextClient, @@ -302,12 +303,16 @@ export async function startSecurityAnalysis(params: { const analysisStartTime = Date.now(); try { - const cloudAgentToken = generateCloudAgentWorkflowToken(user, { + const workflowUser = await prepareCloudAgentWorkflowUser(user, [ + 'cloud-agent-next', + 'workflow-gateway', + ]); + const cloudAgentToken = generateCloudAgentWorkflowToken(workflowUser, { organizationId: findingOrganizationId, tokenSource: 'security-agent', expiresIn: TOKEN_EXPIRY.default, }); - const gatewayToken = generateWorkflowGatewayToken(user, { + const gatewayToken = generateWorkflowGatewayToken(workflowUser, { organizationId: findingOrganizationId, tokenSource: 'security-agent', }); diff --git a/services/cloud-agent-next/src/balance-validation.test.ts b/services/cloud-agent-next/src/balance-validation.test.ts index 4869d5f6ed..f9cc121e58 100644 --- a/services/cloud-agent-next/src/balance-validation.test.ts +++ b/services/cloud-agent-next/src/balance-validation.test.ts @@ -1,3 +1,5 @@ +import jwt from 'jsonwebtoken'; +import { verifyKiloTokenForPolicy } from '@kilocode/worker-utils/kilo-token-policy'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { validateBalanceOnly, @@ -47,6 +49,68 @@ describe('balance-validation', () => { }); describe('validateBalanceOnly', () => { + it('fails closed if bearer classification throws', async () => { + const decode = vi.spyOn(jwt, 'decode').mockImplementationOnce(() => { + throw new Error('Malformed token'); + }); + try { + await expect(validateBalanceOnly('malformed', undefined, mockEnv)).resolves.toEqual({ + success: false, + status: 500, + message: 'Failed to verify balance', + }); + expect(fetchMock).not.toHaveBeenCalled(); + } finally { + decode.mockRestore(); + } + }); + + it.each(['human-api', 'device-access'])( + 'checks balance with a %s control token without forwarding it to the API audience', + async tokenPurpose => { + const token = jwt.sign( + { + version: 3, + kiloUserId: 'user_123', + apiTokenPepper: 'current-pepper', + env: 'test', + aud: 'cloud-agent-next', + tokenPurpose, + credentialExchange: false, + ...(tokenPurpose === 'device-access' ? { deviceSessionId: 'device_123' } : {}), + }, + 'test-secret', + { expiresIn: 60 } + ); + fetchMock.mockImplementation(async (url: string, init: RequestInit) => { + const headers = new Headers(init.headers); + const bearer = headers.get('Authorization')?.slice('Bearer '.length) ?? ''; + try { + await verifyKiloTokenForPolicy(bearer, 'test-secret', { + audience: url.endsWith('/api/cloud-agent-next/balance') + ? 'cloud-agent-next' + : 'kilo-api', + mode: 'required', + }); + return Response.json({ balance: 10, isDepleted: false }); + } catch { + return new Response('Unauthorized', { status: 401 }); + } + }); + + await expect(validateBalanceOnly(token, 'org_123', mockEnv)).resolves.toEqual({ + success: true, + }); + expect(fetchMock).toHaveBeenCalledWith( + 'https://app.kilo.ai/api/cloud-agent-next/balance', + expect.objectContaining({ method: 'GET' }) + ); + const headers = new Headers(fetchMock.mock.calls[0][1].headers); + expect(headers.get('Authorization')).toBe(`Bearer ${token}`); + expect(headers.get('X-KiloCode-OrganizationId')).toBe('org_123'); + } + ); + describe('balance validation', () => { it('returns 402 when balance is depleted', async () => { fetchMock.mockResolvedValue({ diff --git a/services/cloud-agent-next/src/balance-validation.ts b/services/cloud-agent-next/src/balance-validation.ts index 26f913de94..8dc7a03960 100644 --- a/services/cloud-agent-next/src/balance-validation.ts +++ b/services/cloud-agent-next/src/balance-validation.ts @@ -1,3 +1,5 @@ +import jwt from 'jsonwebtoken'; +import { CLOUD_AGENT_NEXT_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; import { DEFAULT_BACKEND_URL } from './constants.js'; import { logger } from './logger.js'; import type { Env } from './types.js'; @@ -36,9 +38,21 @@ export async function validateBalanceOnly( headers.set('X-KiloCode-OrganizationId', orgId); } + // Decode only to select the endpoint; the backend verifies the bearer and + // current user/organization authorization for its exact expected audience. + let claims: ReturnType; + try { + claims = jwt.decode(token); + } catch { + return { success: false, status: 500, message: 'Failed to verify balance' }; + } + const isControlToken = + claims !== null && typeof claims === 'object' && claims.aud === CLOUD_AGENT_NEXT_AUDIENCE; + const balancePath = isControlToken ? '/api/cloud-agent-next/balance' : '/api/profile/balance'; + let response: Response; try { - response = await fetch(`${backendUrl}/api/profile/balance`, { + response = await fetch(`${backendUrl}${balancePath}`, { method: 'GET', headers, });