From d26969c085d064e5c94ae1ccbfb26d6afcbf3f1a Mon Sep 17 00:00:00 2001 From: Bohdan Krasko Date: Wed, 13 May 2026 14:38:30 +0300 Subject: [PATCH] fix: retry installation auth for regional GitHub App events Handle webhook events whose installationId belongs to a different GitHub App by retrying installation authentication after a 404. Apply the fallback in both the shared Octokit helper and the scale-up flow by re-resolving the installation ID for the current app before retrying auth. Preserve existing behavior for installationId=0 and non-404 errors, and add regression tests covering the mismatched installation/app scenario. --- .../control-plane/src/github/octokit.test.ts | 36 +++++++++- .../control-plane/src/github/octokit.ts | 68 +++++++++++++++--- .../src/scale-runners/scale-up.test.ts | 28 ++++++++ .../src/scale-runners/scale-up.ts | 69 ++++++++++++++++--- 4 files changed, 179 insertions(+), 22 deletions(-) diff --git a/lambdas/functions/control-plane/src/github/octokit.test.ts b/lambdas/functions/control-plane/src/github/octokit.test.ts index 3e37d64757..c104b69641 100644 --- a/lambdas/functions/control-plane/src/github/octokit.test.ts +++ b/lambdas/functions/control-plane/src/github/octokit.test.ts @@ -2,6 +2,7 @@ import { Octokit } from '@octokit/rest'; import { ActionRequestMessage } from '../scale-runners/scale-up'; import { getOctokit } from './octokit'; import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { createGithubAppAuth, createGithubInstallationAuth } from '../github/auth'; const mockOctokit = { apps: { @@ -11,7 +12,7 @@ const mockOctokit = { }; vi.mock('../github/auth', async () => ({ - createGithubInstallationAuth: vi.fn().mockImplementation(async (installationId) => { + createGithubInstallationAuth: vi.fn().mockImplementation(async (installationId: number) => { return { token: 'token', type: 'installation', installationId: installationId }; }), createOctokitClient: vi.fn().mockImplementation(() => new Octokit()), @@ -27,7 +28,11 @@ vi.mock('@octokit/rest', async () => ({ // We've already mocked '../github/auth' above describe('Test getOctokit', () => { - const data = [ + const data: Array<{ + description: string; + input: { orgLevelRunner: boolean; installationId: number }; + output: { callReposInstallation: boolean; callOrgInstallation: boolean }; + }> = [ { description: 'Should look-up org installation if installationId is 0.', input: { orgLevelRunner: false, installationId: 0 }, @@ -49,7 +54,7 @@ describe('Test getOctokit', () => { vi.clearAllMocks(); }); - it.each(data)(`$description`, async ({ input, output }) => { + it.each(data)(`$description`, async ({ input, output }: (typeof data)[number]) => { const payload = { eventType: 'workflow_job', id: 0, @@ -74,6 +79,31 @@ describe('Test getOctokit', () => { } else if (output.callReposInstallation) { expect(mockOctokit.apps.getRepoInstallation).toHaveBeenCalled(); expect(mockOctokit.apps.getOrgInstallation).not.toHaveBeenCalled(); + } else { + expect(createGithubAppAuth).not.toHaveBeenCalled(); } }); + + it('Should resolve installation again when event installation belongs to another app', async () => { + const payload = { + eventType: 'workflow_job', + id: 0, + installationId: 999, + repositoryOwner: 'owner', + repositoryName: 'repo', + } as ActionRequestMessage; + + mockOctokit.apps.getOrgInstallation.mockResolvedValue({ data: { id: 123 } }); + + vi.mocked(createGithubInstallationAuth) + .mockRejectedValueOnce({ status: 404 }) + .mockResolvedValueOnce({ token: 'token', type: 'installation', installationId: 123 }); + + await expect(getOctokit('', true, payload)).resolves.toBeDefined(); + + expect(createGithubAppAuth).toHaveBeenCalledTimes(1); + expect(mockOctokit.apps.getOrgInstallation).toHaveBeenCalledWith({ org: 'owner' }); + expect(createGithubInstallationAuth).toHaveBeenNthCalledWith(1, 999, ''); + expect(createGithubInstallationAuth).toHaveBeenNthCalledWith(2, 123, ''); + }); }); diff --git a/lambdas/functions/control-plane/src/github/octokit.ts b/lambdas/functions/control-plane/src/github/octokit.ts index a2cce5f55d..588ec62150 100644 --- a/lambdas/functions/control-plane/src/github/octokit.ts +++ b/lambdas/functions/control-plane/src/github/octokit.ts @@ -2,17 +2,20 @@ import { Octokit } from '@octokit/rest'; import { ActionRequestMessage } from '../scale-runners/scale-up'; import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient } from './auth'; -export async function getInstallationId( - ghesApiUrl: string, +function getErrorStatus(error: unknown): number | undefined { + if (typeof error !== 'object' || error === null) { + return undefined; + } + + const errorWithStatus = error as { status?: number; response?: { status?: number } }; + return errorWithStatus.status ?? errorWithStatus.response?.status; +} + +async function resolveInstallationId( + githubClient: Octokit, enableOrgLevel: boolean, payload: ActionRequestMessage, ): Promise { - if (payload.installationId !== 0) { - return payload.installationId; - } - - const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl); - const githubClient = await createOctokitClient(ghAuth.token, ghesApiUrl); return enableOrgLevel ? ( await githubClient.apps.getOrgInstallation({ @@ -27,6 +30,20 @@ export async function getInstallationId( ).data.id; } +export async function getInstallationId( + ghesApiUrl: string, + enableOrgLevel: boolean, + payload: ActionRequestMessage, +): Promise { + if (payload.installationId !== 0) { + return payload.installationId; + } + + const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl); + const githubClient = await createOctokitClient(ghAuth.token, ghesApiUrl); + return resolveInstallationId(githubClient, enableOrgLevel, payload); +} + /** * * Util method to get an octokit client based on provided installation id. This method should @@ -40,7 +57,36 @@ export async function getOctokit( enableOrgLevel: boolean, payload: ActionRequestMessage, ): Promise { - const installationId = await getInstallationId(ghesApiUrl, enableOrgLevel, payload); - const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl); - return await createOctokitClient(ghAuth.token, ghesApiUrl); + let githubAppClient: Octokit | undefined; + let installationId = payload.installationId !== 0 ? payload.installationId : undefined; + + const getGithubAppClient = async (): Promise => { + if (githubAppClient === undefined) { + const appAuth = await createGithubAppAuth(undefined, ghesApiUrl); + githubAppClient = await createOctokitClient(appAuth.token, ghesApiUrl); + } + + return githubAppClient; + }; + + try { + if (installationId === undefined) { + installationId = await resolveInstallationId(await getGithubAppClient(), enableOrgLevel, payload); + } + + const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl); + return await createOctokitClient(ghAuth.token, ghesApiUrl); + } catch (error) { + if (payload.installationId === 0 || getErrorStatus(error) !== 404) { + throw error; + } + + const resolvedInstallationId = await resolveInstallationId(await getGithubAppClient(), enableOrgLevel, payload); + if (resolvedInstallationId === payload.installationId) { + throw error; + } + + const ghAuth = await createGithubInstallationAuth(resolvedInstallationId, ghesApiUrl); + return await createOctokitClient(ghAuth.token, ghesApiUrl); + } } diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts index 8ac2c14489..0110ec1114 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts @@ -824,6 +824,34 @@ describe('scaleUp with GHES', () => { expect(mockedInstallationAuth).toHaveBeenCalledWith(200, 'https://github.enterprise.something/api/v3'); }); + it('Should resolve installation again when event installation belongs to another app', async () => { + mockOctokit.apps.getOrgInstallation.mockReset(); + mockOctokit.apps.getOrgInstallation.mockImplementation(() => ({ + data: { + id: 123, + }, + })); + + mockedInstallationAuth + .mockRejectedValueOnce({ status: 404 }) + .mockResolvedValueOnce({ + type: 'token', + tokenType: 'installation', + token: 'token', + createdAt: 'some-date', + expiresAt: 'some-date', + permissions: {}, + repositorySelection: 'all', + installationId: 123, + }); + + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockOctokit.apps.getOrgInstallation).toHaveBeenCalledWith({ org: TEST_DATA_SINGLE.repositoryOwner }); + expect(mockedInstallationAuth).toHaveBeenNthCalledWith(1, TEST_DATA_SINGLE.installationId, 'https://github.enterprise.something/api/v3'); + expect(mockedInstallationAuth).toHaveBeenNthCalledWith(2, 123, 'https://github.enterprise.something/api/v3'); + }); + it('Should reuse GitHub clients for same installation', async () => { const messages = createTestMessages(3, [ { repositoryOwner: 'same-org' }, diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts index 395c87e8f8..bee86a404b 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts @@ -149,15 +149,20 @@ function removeTokenFromLogging(config: string[]): string[] { return result; } -export async function getInstallationId( +function getErrorStatus(error: unknown): number | undefined { + if (typeof error !== 'object' || error === null) { + return undefined; + } + + const errorWithStatus = error as { status?: number; response?: { status?: number } }; + return errorWithStatus.status ?? errorWithStatus.response?.status; +} + +async function resolveInstallationId( githubAppClient: Octokit, enableOrgLevel: boolean, payload: ActionRequestMessage, ): Promise { - if (payload.installationId !== 0) { - return payload.installationId; - } - return enableOrgLevel ? ( await githubAppClient.apps.getOrgInstallation({ @@ -172,6 +177,18 @@ export async function getInstallationId( ).data.id; } +export async function getInstallationId( + githubAppClient: Octokit, + enableOrgLevel: boolean, + payload: ActionRequestMessage, +): Promise { + if (payload.installationId !== 0) { + return payload.installationId; + } + + return resolveInstallationId(githubAppClient, enableOrgLevel, payload); +} + export async function isJobQueued(githubInstallationClient: Octokit, payload: ActionRequestMessage): Promise { let isQueued = false; if (payload.eventType === 'workflow_job') { @@ -288,6 +305,39 @@ export async function createRunners( return instances; } +async function createGithubInstallationClient( + githubAppClient: Octokit, + enableOrgLevel: boolean, + payload: ActionRequestMessage, + ghesApiUrl: string, +): Promise { + let installationId = await getInstallationId(githubAppClient, enableOrgLevel, payload); + + try { + const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl); + return await createOctokitClient(ghAuth.token, ghesApiUrl); + } catch (error) { + if (payload.installationId === 0 || getErrorStatus(error) !== 404) { + throw error; + } + + installationId = await resolveInstallationId(githubAppClient, enableOrgLevel, payload); + if (installationId === payload.installationId) { + throw error; + } + + logger.warn('Retrying GitHub installation auth with installation resolved for current app', { + eventInstallationId: payload.installationId, + resolvedInstallationId: installationId, + repositoryOwner: payload.repositoryOwner, + repositoryName: payload.repositoryName, + }); + + const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl); + return await createOctokitClient(ghAuth.token, ghesApiUrl); + } +} + export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise { logger.info('Received scale up requests', { n_requests: payloads.length, @@ -373,9 +423,12 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise