diff --git a/apps/web/src/lib/integrations/oauth/platforms/slack-callback.db.test.ts b/apps/web/src/lib/integrations/oauth/platforms/slack-callback.db.test.ts new file mode 100644 index 0000000000..33c0508317 --- /dev/null +++ b/apps/web/src/lib/integrations/oauth/platforms/slack-callback.db.test.ts @@ -0,0 +1,200 @@ +import { NextRequest } from 'next/server'; +import { eq } from 'drizzle-orm'; +import { db } from '@kilocode/web-shared/lib/drizzle'; +import { kilocode_users, platform_integrations } from '@kilocode/db/schema'; +import { getUserFromAuth } from '@kilocode/web-shared/lib/user/server'; +import { createOAuthState } from '@/lib/integrations/oauth-state'; +import { beginProviderOAuthAttempt } from '@/lib/integrations/provider-oauth-attempts'; +import { SLACK_SCOPES } from '@/lib/integrations/slack-service'; +import { handleSlackOAuthCallback } from './slack-callback'; + +const mockHandleOAuthCallback = jest.fn(); +const mockSetInstallation = jest.fn(); +const mockDeleteInstallation = jest.fn(); +const mockRevokeToken = jest.fn(); +const mockRevokedTokens: string[] = []; +const mockWriteSlackCredential = jest.fn(); + +jest.mock('@kilocode/web-shared/lib/user/server'); +jest.mock('@/lib/bot', () => ({ + bot: { + initialize: jest.fn(async () => {}), + getAdapter: () => ({ + handleOAuthCallback: (...args: unknown[]) => mockHandleOAuthCallback(...args), + setInstallation: (...args: unknown[]) => mockSetInstallation(...args), + deleteInstallation: (...args: unknown[]) => mockDeleteInstallation(...args), + }), + }, +})); +jest.mock('@slack/web-api', () => ({ + WebClient: function WebClient(this: unknown, token: string) { + mockRevokedTokens.push(token); + return { auth: { revoke: () => mockRevokeToken() } }; + }, +})); +jest.mock('@/lib/integrations/platforms/slack/credential-store', () => ({ + writeSlackCredential: (...args: unknown[]) => mockWriteSlackCredential(...args), +})); +jest.mock('@/lib/slack-bot/model-allow-list', () => ({ + getDefaultAllowedModel: jest.fn(async () => 'test-model'), +})); + +const mockedGetUserFromAuth = jest.mocked(getUserFromAuth); + +const OWNER_USER_ID = crypto.randomUUID(); +const ACTING_USER_ID = crypto.randomUUID(); +const TEAM_ID = 'T_DB_SHARED'; +const TEAM_NAME = 'Shared DB Workspace'; +const RECORDED_TOKEN = 'fixture-recorded-token'; +const INCOMING_TOKEN = 'fixture-incoming-token'; +const REINSTALL_TOKEN = 'fixture-reinstall-token'; + +function buildUser(id: string, name: string) { + return { + id, + google_user_email: `${id}@example.com`, + google_user_name: name, + google_user_image_url: 'https://example.com/avatar.png', + stripe_customer_id: `cus_slack_${id}`, + }; +} + +async function callCallback(ownerUserId: string) { + const owner = `user_${ownerUserId}`; + const state = createOAuthState(owner, ownerUserId, undefined, 'provider_install'); + await beginProviderOAuthAttempt({ + actorUserId: ownerUserId, + owner: { type: 'user', id: ownerUserId }, + provider: 'slack', + state, + purpose: 'provider_install', + }); + const request = new NextRequest( + `http://localhost:3000/api/integrations/slack/callback?code=code-1&state=${encodeURIComponent(state)}` + ); + return handleSlackOAuthCallback(request); +} + +describe('Slack OAuth callback workspace binding', () => { + beforeAll(async () => { + await db + .insert(kilocode_users) + .values([ + buildUser(OWNER_USER_ID, 'Recorded Owner'), + buildUser(ACTING_USER_ID, 'Second Installer'), + ]); + }); + + beforeEach(async () => { + jest.clearAllMocks(); + mockRevokedTokens.length = 0; + mockSetInstallation.mockResolvedValue(undefined); + mockDeleteInstallation.mockResolvedValue(undefined); + mockRevokeToken.mockResolvedValue({ ok: true }); + mockWriteSlackCredential.mockResolvedValue(undefined); + await db + .delete(platform_integrations) + .where(eq(platform_integrations.owned_by_user_id, OWNER_USER_ID)); + await db.insert(platform_integrations).values({ + owned_by_user_id: OWNER_USER_ID, + platform: 'slack', + integration_type: 'oauth', + platform_installation_id: TEAM_ID, + platform_account_id: TEAM_ID, + platform_account_login: TEAM_NAME, + scopes: SLACK_SCOPES, + integration_status: 'active', + metadata: { access_token: RECORDED_TOKEN, bot_user_id: 'U_RECORDED_OWNER' }, + }); + }); + + afterAll(async () => { + await db + .delete(platform_integrations) + .where(eq(platform_integrations.owned_by_user_id, OWNER_USER_ID)); + await db + .delete(platform_integrations) + .where(eq(platform_integrations.owned_by_user_id, ACTING_USER_ID)); + await db.delete(kilocode_users).where(eq(kilocode_users.id, OWNER_USER_ID)); + await db.delete(kilocode_users).where(eq(kilocode_users.id, ACTING_USER_ID)); + }); + + function recordInvocation(accessToken = INCOMING_TOKEN) { + mockHandleOAuthCallback.mockResolvedValue({ + teamId: TEAM_ID, + installation: { + botToken: accessToken, + botUserId: 'U_INCOMING', + teamName: TEAM_NAME, + }, + }); + } + + it("leaves the recorded owner's row and installation in place when another user installs the workspace", async () => { + recordInvocation(); + mockedGetUserFromAuth.mockResolvedValue({ + user: { id: ACTING_USER_ID, is_admin: false }, + authFailedResponse: null, + } as never); + + const [before] = await db + .select() + .from(platform_integrations) + .where(eq(platform_integrations.owned_by_user_id, OWNER_USER_ID)); + + const response = await callCallback(ACTING_USER_ID); + + const location = new URL(response.headers.get('location') ?? ''); + expect(`${location.pathname}${location.search}`).toBe( + '/integrations/slack?error=workspace_already_connected' + ); + + const [after] = await db + .select() + .from(platform_integrations) + .where(eq(platform_integrations.owned_by_user_id, OWNER_USER_ID)); + expect(after).toEqual(before); + expect(mockSetInstallation).toHaveBeenCalledWith(TEAM_ID, { + botToken: RECORDED_TOKEN, + botUserId: 'U_RECORDED_OWNER', + teamName: TEAM_NAME, + }); + expect(mockDeleteInstallation).not.toHaveBeenCalled(); + expect(mockRevokedTokens).toEqual([INCOMING_TOKEN]); + expect(mockWriteSlackCredential).not.toHaveBeenCalled(); + + const actingIntegrations = await db + .select() + .from(platform_integrations) + .where(eq(platform_integrations.owned_by_user_id, ACTING_USER_ID)); + expect(actingIntegrations).toEqual([]); + }); + + it('records a reinstall by the same user and leaves the installation untouched', async () => { + recordInvocation(REINSTALL_TOKEN); + mockedGetUserFromAuth.mockResolvedValue({ + user: { id: OWNER_USER_ID, is_admin: false }, + authFailedResponse: null, + } as never); + + const response = await callCallback(OWNER_USER_ID); + + const location = new URL(response.headers.get('location') ?? ''); + expect(`${location.pathname}${location.search}`).toBe('/integrations/slack?success=installed'); + + const [after] = await db + .select() + .from(platform_integrations) + .where(eq(platform_integrations.owned_by_user_id, OWNER_USER_ID)); + expect(after.platform_installation_id).toBe(TEAM_ID); + expect(after.metadata).toEqual( + expect.objectContaining({ + access_token: REINSTALL_TOKEN, + bot_user_id: 'U_INCOMING', + }) + ); + expect(mockSetInstallation).not.toHaveBeenCalled(); + expect(mockDeleteInstallation).not.toHaveBeenCalled(); + expect(mockRevokedTokens).toEqual([]); + }); +}); diff --git a/apps/web/src/lib/integrations/oauth/platforms/slack-callback.test.ts b/apps/web/src/lib/integrations/oauth/platforms/slack-callback.test.ts index 4edd5eb0b3..168572e7cc 100644 --- a/apps/web/src/lib/integrations/oauth/platforms/slack-callback.test.ts +++ b/apps/web/src/lib/integrations/oauth/platforms/slack-callback.test.ts @@ -2,8 +2,13 @@ process.env.NEXTAUTH_SECRET ||= 'test-nextauth-secret'; const mockEnsureOrganizationAccess = jest.fn(); const mockConsumeProviderOAuthAttempt = jest.fn(); +const mockCancelProviderOAuthAttempt = jest.fn(); const mockUpsertSlackInstallation = jest.fn(); +const mockGetInstallationByTeamId = jest.fn(); +const mockRevokeSlackToken = jest.fn(); const mockHandleOAuthCallback = jest.fn(); +const mockSetInstallation = jest.fn(); +const mockDeleteInstallation = jest.fn(); const mockCaptureException = jest.fn(); const mockCaptureMessage = jest.fn(); @@ -23,12 +28,14 @@ jest.mock('@/lib/bot', () => ({ initialize: jest.fn(async () => {}), getAdapter: () => ({ handleOAuthCallback: (...args: unknown[]) => mockHandleOAuthCallback(...args), + setInstallation: (...args: unknown[]) => mockSetInstallation(...args), + deleteInstallation: (...args: unknown[]) => mockDeleteInstallation(...args), }), }, })); jest.mock('@/lib/integrations/provider-oauth-attempts', () => ({ beginProviderOAuthAttempt: jest.fn(), - cancelProviderOAuthAttempt: jest.fn(), + cancelProviderOAuthAttempt: (...args: unknown[]) => mockCancelProviderOAuthAttempt(...args), consumeProviderOAuthAttempt: (...args: unknown[]) => mockConsumeProviderOAuthAttempt(...args), })); jest.mock('@/lib/integrations/slack-service', () => { @@ -40,6 +47,8 @@ jest.mock('@/lib/integrations/slack-service', () => { } return { SlackWorkspaceAlreadyConnectedError, + getInstallationByTeamId: (...args: unknown[]) => mockGetInstallationByTeamId(...args), + revokeSlackToken: (...args: unknown[]) => mockRevokeSlackToken(...args), upsertSlackInstallation: (...args: unknown[]) => mockUpsertSlackInstallation(...args), }; }); @@ -49,6 +58,7 @@ import { NextRequest } from 'next/server'; import { getUserFromAuth } from '@kilocode/web-shared/lib/user/server'; import { ORGANIZATION_MANAGE_ROLES } from '@kilocode/app-shared/organizations'; import { createOAuthState } from '@/lib/integrations/oauth-state'; +import { SlackWorkspaceAlreadyConnectedError } from '@/lib/integrations/slack-service'; import { handleSlackOAuthCallback } from './slack-callback'; const mockedGetUserFromAuth = jest.mocked(getUserFromAuth); @@ -74,6 +84,13 @@ async function callCallback( return { response, state }; } +function expectRedirectLocation(response: Response, expectedPathWithQuery: string) { + const location = response.headers.get('location'); + expect(location).toBeTruthy(); + const url = new URL(location ?? ''); + expect(`${url.pathname}${url.search}`).toBe(expectedPathWithQuery); +} + describe('handleSlackOAuthCallback organization authority', () => { beforeEach(() => { jest.clearAllMocks(); @@ -165,3 +182,155 @@ describe('handleSlackOAuthCallback organization authority', () => { }); }); }); + +describe('handleSlackOAuthCallback workspace ownership', () => { + const TEAM_ID = 'T_SHARED'; + const TEAM_NAME = 'Shared Workspace'; + const STORED_TOKEN = 'bot-token-existing-owner'; + const REJECTED_TOKEN = 'bot-token-rejected-exchange'; + + function buildStoredIntegration(overrides: Record = {}) { + return { + id: 'integration-existing-owner', + platform: 'slack', + platform_installation_id: TEAM_ID, + platform_account_id: TEAM_ID, + platform_account_login: TEAM_NAME, + integration_status: 'active', + owned_by_user_id: 'other-owner', + owned_by_organization_id: null, + metadata: { access_token: STORED_TOKEN, bot_user_id: 'U_EXISTING_OWNER' }, + ...overrides, + } as never; + } + + function mockRejectedExchange(rejectedToken = REJECTED_TOKEN) { + mockHandleOAuthCallback.mockResolvedValue({ + teamId: TEAM_ID, + installation: { + botToken: rejectedToken, + botUserId: 'U_REJECTED_INSTALL', + teamName: TEAM_NAME, + }, + }); + mockUpsertSlackInstallation.mockRejectedValue( + new SlackWorkspaceAlreadyConnectedError(TEAM_NAME) + ); + } + + beforeEach(() => { + jest.clearAllMocks(); + mockSetInstallation.mockResolvedValue(undefined); + mockDeleteInstallation.mockResolvedValue(undefined); + mockConsumeProviderOAuthAttempt.mockResolvedValue(true); + mockCancelProviderOAuthAttempt.mockResolvedValue(true); + mockEnsureOrganizationAccess.mockResolvedValue('owner'); + mockedGetUserFromAuth.mockResolvedValue({ + user: { id: USER_ID, is_admin: false }, + authFailedResponse: null, + } as never); + mockRevokeSlackToken.mockResolvedValue(true); + mockGetInstallationByTeamId.mockResolvedValue(null); + }); + + it("keeps the recorded owner's installation when another owner installs the workspace", async () => { + mockRejectedExchange(); + mockGetInstallationByTeamId.mockResolvedValue(buildStoredIntegration()); + + const { response } = await callCallback(`user_${USER_ID}`); + + expectRedirectLocation(response, '/integrations/slack?error=workspace_already_connected'); + expect(mockGetInstallationByTeamId).toHaveBeenCalledWith(TEAM_ID); + expect(mockSetInstallation).toHaveBeenCalledTimes(1); + expect(mockSetInstallation).toHaveBeenCalledWith(TEAM_ID, { + botToken: STORED_TOKEN, + botUserId: 'U_EXISTING_OWNER', + teamName: TEAM_NAME, + }); + expect(mockDeleteInstallation).not.toHaveBeenCalled(); + expect(mockRevokeSlackToken).toHaveBeenCalledTimes(1); + expect(mockRevokeSlackToken).toHaveBeenCalledWith(REJECTED_TOKEN); + expect(mockCaptureException).not.toHaveBeenCalled(); + }); + + it('keeps the workspace token when an install returns the recorded token', async () => { + mockRejectedExchange(STORED_TOKEN); + mockGetInstallationByTeamId.mockResolvedValue(buildStoredIntegration()); + + const { response } = await callCallback(`user_${USER_ID}`); + + expectRedirectLocation(response, '/integrations/slack?error=workspace_already_connected'); + expect(mockSetInstallation).toHaveBeenCalledWith(TEAM_ID, { + botToken: STORED_TOKEN, + botUserId: 'U_EXISTING_OWNER', + teamName: TEAM_NAME, + }); + expect(mockRevokeSlackToken).not.toHaveBeenCalled(); + }); + + it('removes the installation when the workspace has no recorded owner', async () => { + mockRejectedExchange(); + + const { response } = await callCallback(`user_${USER_ID}`); + + expectRedirectLocation(response, '/integrations/slack?error=workspace_already_connected'); + expect(mockSetInstallation).not.toHaveBeenCalled(); + expect(mockDeleteInstallation).toHaveBeenCalledWith(TEAM_ID); + expect(mockRevokeSlackToken).toHaveBeenCalledWith(REJECTED_TOKEN); + }); + + it('removes the installation when the recorded integration has no token', async () => { + mockRejectedExchange(); + mockGetInstallationByTeamId.mockResolvedValue( + buildStoredIntegration({ metadata: { bot_user_id: 'U_EXISTING_OWNER' } }) + ); + + const { response } = await callCallback(`user_${USER_ID}`); + + expectRedirectLocation(response, '/integrations/slack?error=workspace_already_connected'); + expect(mockSetInstallation).not.toHaveBeenCalled(); + expect(mockDeleteInstallation).toHaveBeenCalledWith(TEAM_ID); + expect(mockRevokeSlackToken).toHaveBeenCalledWith(REJECTED_TOKEN); + }); + + it('revokes the incoming token when removing the installation fails', async () => { + mockRejectedExchange(); + mockDeleteInstallation.mockRejectedValue(new Error('state unavailable')); + + const { response } = await callCallback(`user_${USER_ID}`); + + expectRedirectLocation(response, '/integrations/slack?error=workspace_already_connected'); + expect(mockRevokeSlackToken).toHaveBeenCalledWith(REJECTED_TOKEN); + expect(mockCaptureException).toHaveBeenCalledTimes(1); + }); + + it('removes the installation when the recorded owner cannot be written', async () => { + mockRejectedExchange(); + mockSetInstallation.mockRejectedValue(new Error('state unavailable')); + mockGetInstallationByTeamId.mockResolvedValue(buildStoredIntegration()); + + const { response } = await callCallback(`user_${USER_ID}`); + + expectRedirectLocation(response, '/integrations/slack?error=workspace_already_connected'); + expect(mockDeleteInstallation).toHaveBeenCalledWith(TEAM_ID); + expect(mockRevokeSlackToken).toHaveBeenCalledWith(REJECTED_TOKEN); + expect(mockCaptureException).toHaveBeenCalledTimes(1); + }); + + it('leaves a same-owner reinstall untouched', async () => { + mockHandleOAuthCallback.mockResolvedValue({ + teamId: TEAM_ID, + installation: { botToken: STORED_TOKEN, botUserId: 'U_EXISTING_OWNER', teamName: TEAM_NAME }, + }); + mockUpsertSlackInstallation.mockResolvedValue(buildStoredIntegration()); + + const { response } = await callCallback(`user_${USER_ID}`); + + expectRedirectLocation(response, '/integrations/slack?success=installed'); + expect(mockUpsertSlackInstallation).toHaveBeenCalledTimes(1); + expect(mockSetInstallation).not.toHaveBeenCalled(); + expect(mockDeleteInstallation).not.toHaveBeenCalled(); + expect(mockRevokeSlackToken).not.toHaveBeenCalled(); + expect(mockGetInstallationByTeamId).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/src/lib/integrations/oauth/platforms/slack-callback.ts b/apps/web/src/lib/integrations/oauth/platforms/slack-callback.ts index 88d4024215..540559005c 100644 --- a/apps/web/src/lib/integrations/oauth/platforms/slack-callback.ts +++ b/apps/web/src/lib/integrations/oauth/platforms/slack-callback.ts @@ -6,9 +6,13 @@ import { ensureOrganizationAccess } from '@kilocode/web-shared/routers/organizat import { captureException, captureMessage } from '@sentry/nextjs'; import { getCallbackReportParams } from '@/lib/integrations/oauth/callback-params'; import { + getInstallationByTeamId, + revokeSlackToken, SlackWorkspaceAlreadyConnectedError, upsertSlackInstallation, } from '@/lib/integrations/slack-service'; +import type { SlackAdapter, SlackInstallation } from '@chat-adapter/slack'; +import type { PlatformIntegration } from '@kilocode/db/schema'; import { isLegacyProviderOAuthState, verifyOAuthState } from '@/lib/integrations/oauth-state'; import { APP_URL } from '@kilocode/web-shared/lib/constants'; import { bot } from '@/lib/bot'; @@ -26,6 +30,95 @@ import { consumeProviderOAuthAttempt } from '@/lib/integrations/provider-oauth-a const SLACK_REDIRECT_URI = getPlatformOAuthCallbackUrl(PLATFORM.SLACK); +/** + * Build the Chat SDK installation for a workspace from its recorded + * integration. + */ +function buildStoredSlackInstallation(integration: PlatformIntegration): SlackInstallation | null { + const metadata = integration.metadata as + | { access_token?: unknown; bot_user_id?: unknown } + | null + | undefined; + const botToken = typeof metadata?.access_token === 'string' ? metadata.access_token : null; + + if (!botToken) return null; + + return { + botToken, + ...(typeof metadata?.bot_user_id === 'string' ? { botUserId: metadata.bot_user_id } : {}), + ...(integration.platform_account_login ? { teamName: integration.platform_account_login } : {}), + }; +} + +/** + * Write the workspace installation recorded for the workspace, or remove the + * workspace installation when nothing is recorded. + */ +async function writeRecordedSlackInstallation({ + slackAdapter, + teamId, + recordedInstallation, +}: { + slackAdapter: SlackAdapter; + teamId: string; + recordedInstallation: SlackInstallation | null; +}): Promise { + if (!recordedInstallation) { + await slackAdapter.deleteInstallation(teamId); + return; + } + + try { + await slackAdapter.setInstallation(teamId, recordedInstallation); + } catch (error) { + captureException(error, { + tags: { endpoint: 'slack/callback', source: 'slack_oauth', op: 'restore-installation' }, + extra: { teamId }, + }); + await slackAdapter.deleteInstallation(teamId); + } +} + +/** + * Keep a workspace's Chat SDK installation attached to the owner recorded for + * it, and revoke a token that ends up attached to no installation. Revocation + * runs whether or not the installation write succeeds. Workspace identity links + * are left as they are. + */ +async function restoreRecordedSlackInstallation({ + slackAdapter, + teamId, + incomingInstallation, +}: { + slackAdapter: SlackAdapter; + teamId: string; + incomingInstallation: SlackInstallation; +}): Promise { + const recordedIntegration = await getInstallationByTeamId(teamId); + const recordedInstallation = recordedIntegration + ? buildStoredSlackInstallation(recordedIntegration) + : null; + const reusesRecordedToken = incomingInstallation.botToken === recordedInstallation?.botToken; + + try { + await writeRecordedSlackInstallation({ slackAdapter, teamId, recordedInstallation }); + } finally { + if (!reusesRecordedToken) { + const revoked = await revokeSlackToken(incomingInstallation.botToken); + if (!revoked) { + captureMessage( + 'Slack token revocation did not complete while settling a workspace install', + { + level: 'warning', + tags: { endpoint: 'slack/callback', source: 'slack_oauth', op: 'revoke-token' }, + extra: { teamId }, + } + ); + } + } + } +} + /** * Slack OAuth Callback * @@ -165,6 +258,23 @@ export async function handleSlackOAuthCallback(request: NextRequest) { }); } catch (error) { if (error instanceof SlackWorkspaceAlreadyConnectedError) { + try { + await restoreRecordedSlackInstallation({ + slackAdapter, + teamId, + incomingInstallation: installation, + }); + } catch (rollbackError) { + captureException(rollbackError, { + tags: { + endpoint: 'slack/callback', + source: 'slack_oauth', + op: 'install-conflict-rollback', + }, + extra: { teamId }, + }); + } + return NextResponse.redirect( new URL( buildIntegrationOAuthRedirectPathFromState(