Skip to content

feat(termination-watcher): support additional GitHub Apps #5467

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
guicaulada wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(termination-watcher): support additional GitHub Apps #5467

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions lambdas/functions/termination-watcher/src/deregister.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,11 @@ const mockApps = {
getRepoInstallation: vi.fn(),
};

const mockHookAfter = vi.fn();

function MockOctokit() {
return {
hook: { after: mockHookAfter },
actions: mockActions,
apps: mockApps,
paginate: mockPaginate,
Expand Down Expand Up @@ -114,6 +117,28 @@ describe('deregisterRunner', () => {
setupAuthMocks();
});

it('deregisters through an additional App when it is selected', async () => {
mockGetCredentials.mockResolvedValue([
{ appId: 1, privateKey: 'one' },
{ appId: 2, privateKey: 'two' },
]);
mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 222 } });
mockPaginate.iterator.mockReturnValue(
(async function* () {
yield { data: [{ id: 42, name: `runner-${orgInstance.InstanceId}` }] };
})(),
);
mockActions.deleteSelfHostedRunnerFromOrg.mockResolvedValue({});
const random = vi.spyOn(Math, 'random').mockReturnValue(0.99);
try {
await deregisterRunner(orgInstance, baseConfig);
expect(mockCreateAppAuth).toHaveBeenCalledWith({ appId: 2, privateKey: 'two', installationId: 222 });
expect(mockActions.deleteSelfHostedRunnerFromOrg).toHaveBeenCalledWith({ org: 'test-org', runner_id: 42 });
} finally {
random.mockRestore();
}
});

it('should skip deregistration when disabled', async () => {
await deregisterRunner(orgInstance, { ...baseConfig, enableRunnerDeregistration: false });
expect(mockGetCredentials).not.toHaveBeenCalled();
Expand Down
120 changes: 5 additions & 115 deletions lambdas/functions/termination-watcher/src/deregister.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,10 @@
import { createAppAuth } from '@octokit/auth-app';
import { Octokit } from '@octokit/rest';
import { throttling } from '@octokit/plugin-throttling';
import { request } from '@octokit/request';
import type { Octokit } from '@octokit/rest';
import { Instance } from '@aws-sdk/client-ec2';
import { SQSClient, SendMessageCommand } from '@aws-sdk/client-sqs';
import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
import { createCommonStorage, type GitHubAppCredential } from '@aws-github-runner/storage-providers';
import type { EndpointDefaults } from '@octokit/types';
import type { Config } from './ConfigResolver';
import { createRunnerInstallationClient } from './github-app-client';
export { createThrottleOptions, resetAppCredentialsCache } from './github-app-client';

export interface DeregisterRetryMessage {
instanceId: string;
Expand All @@ -21,73 +18,6 @@ const sqsClient = new SQSClient({ region: process.env.AWS_REGION });

const logger = createChildLogger('deregister');

let appCredentialsPromise: Promise<GitHubAppCredential> | undefined;

export function createThrottleOptions() {
return {
onRateLimit: (_retryAfter: number, options: Required<EndpointDefaults>) => {
logger.warn(`Rate limit hit for ${options.method} ${options.url}`);
return false;
},
onSecondaryRateLimit: (_retryAfter: number, options: Required<EndpointDefaults>) => {
logger.warn(`Secondary rate limit hit for ${options.method} ${options.url}`);
return false;
},
};
}

async function loadAppCredentials(): Promise<GitHubAppCredential> {
const credentials = await createCommonStorage().githubAppCredentials.get();
const credential = credentials[0];
if (!credential) {
throw new Error('No GitHub App credentials found');
}
return credential;
}

function getAppCredentials(): Promise<GitHubAppCredential> {
if (!appCredentialsPromise) {
appCredentialsPromise = loadAppCredentials().catch((error: unknown) => {
appCredentialsPromise = undefined;
throw error;
});
}
return appCredentialsPromise;
}

export function resetAppCredentialsCache(): void {
appCredentialsPromise = undefined;
}

function createOctokitInstance(token: string, ghesApiUrl: string): Octokit {
const CustomOctokit = Octokit.plugin(throttling);
const octokitOptions: ConstructorParameters<typeof Octokit>[0] = {
auth: token,
};
if (ghesApiUrl) {
octokitOptions.baseUrl = ghesApiUrl;
}
return new CustomOctokit({
...octokitOptions,
userAgent: 'github-aws-runners-termination-watcher',
throttle: createThrottleOptions(),
});
}

async function createAuthenticatedClient(ghesApiUrl: string): Promise<Octokit> {
const { appId, privateKey } = await getAppCredentials();
const authOptions: { appId: number; privateKey: string; request?: typeof request } = {
appId,
privateKey,
};
if (ghesApiUrl) {
authOptions.request = request.defaults({ baseUrl: ghesApiUrl });
}
const auth = createAppAuth(authOptions);
const appAuth = await auth({ type: 'app' });
return createOctokitInstance(appAuth.token, ghesApiUrl);
}

function getOwnerFromTags(instance: Instance): string | undefined {
return instance.Tags?.find((tag) => tag.Key === 'ghr:Owner')?.Value;
}
Expand All @@ -96,44 +26,6 @@ function getRunnerTypeFromTags(instance: Instance): string | undefined {
return instance.Tags?.find((tag) => tag.Key === 'ghr:Type')?.Value;
}

async function getInstallationId(octokit: Octokit, owner: string): Promise<number> {
const { data: installation } = await octokit.apps.getOrgInstallation({ org: owner });
return installation.id;
}

async function getInstallationIdForRepo(octokit: Octokit, owner: string, repo: string): Promise<number> {
const { data: installation } = await octokit.apps.getRepoInstallation({ owner, repo });
return installation.id;
}

async function createInstallationClient(
appOctokit: Octokit,
owner: string,
runnerType: string,
ghesApiUrl: string,
): Promise<Octokit> {
let installationId: number;
if (runnerType === 'Repo') {
const [repoOwner, repo] = owner.split('/');
installationId = await getInstallationIdForRepo(appOctokit, repoOwner, repo);
} else {
installationId = await getInstallationId(appOctokit, owner);
}

const { appId, privateKey } = await getAppCredentials();
const authOptions: { appId: number; privateKey: string; installationId: number; request?: typeof request } = {
appId,
privateKey,
installationId,
};
if (ghesApiUrl) {
authOptions.request = request.defaults({ baseUrl: ghesApiUrl });
}
const auth = createAppAuth(authOptions);
const installationAuth = await auth({ type: 'installation' });
return createOctokitInstance(installationAuth.token, ghesApiUrl);
}

async function findRunnerByInstanceId(
octokit: Octokit,
owner: string,
Expand Down Expand Up @@ -206,8 +98,7 @@ export async function deregisterRunner(instance: Instance, config: Config): Prom
try {
logger.info('Attempting to deregister runner from GitHub', { instanceId, owner, runnerType });

const appOctokit = await createAuthenticatedClient(config.ghesApiUrl);
const installationOctokit = await createInstallationClient(appOctokit, owner, runnerType, config.ghesApiUrl);
const installationOctokit = await createRunnerInstallationClient(owner, runnerType, config.ghesApiUrl);

const runner = await findRunnerByInstanceId(installationOctokit, owner, instanceId, runnerType);
if (!runner) {
Expand Down Expand Up @@ -261,8 +152,7 @@ export async function handleDeregisterRetry(queueUrl: string, message: Deregiste
logger.info('Processing deregistration retry from SQS', { instanceId, owner, runnerType, retryCount });

try {
const appOctokit = await createAuthenticatedClient('');
const installationOctokit = await createInstallationClient(appOctokit, owner, runnerType, '');
const installationOctokit = await createRunnerInstallationClient(owner, runnerType, process.env.GHES_URL ?? '');

const runner = await findRunnerByInstanceId(installationOctokit, owner, instanceId, runnerType);
if (!runner) {
Expand Down
161 changes: 161 additions & 0 deletions lambdas/functions/termination-watcher/src/github-app-client.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,161 @@
import { createCommonStorage, type GitHubAppCredentialsStore } from '@aws-github-runner/storage-providers';
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { createThrottleOptions, resetAppCredentialsCache, createRunnerInstallationClient } from './github-app-client';
import type { EndpointDefaults } from '@octokit/types';

vi.mock('@aws-github-runner/storage-providers', () => ({
createCommonStorage: vi.fn(),
}));

const mockedCreateCommonStorage = vi.mocked(createCommonStorage);
const mockGetCredentials = vi.fn<GitHubAppCredentialsStore['get']>();
const credentialsStore = { get: mockGetCredentials } satisfies GitHubAppCredentialsStore;

const mockCreateAppAuth = vi.fn();
vi.mock('@octokit/auth-app', () => ({
createAppAuth: (...args: unknown[]) => mockCreateAppAuth(...args),
}));

const mockPaginate = {
iterator: vi.fn(),
};

const mockActions = {
listSelfHostedRunnersForOrg: vi.fn(),
listSelfHostedRunnersForRepo: vi.fn(),
deleteSelfHostedRunnerFromOrg: vi.fn(),
deleteSelfHostedRunnerFromRepo: vi.fn(),
};

const mockApps = {
getOrgInstallation: vi.fn(),
getRepoInstallation: vi.fn(),
};

const mockHookAfter = vi.fn();

function MockOctokit() {
return {
hook: { after: mockHookAfter },
actions: mockActions,
apps: mockApps,
paginate: mockPaginate,
};
}
MockOctokit.plugin = vi.fn().mockReturnValue(MockOctokit);

vi.mock('@octokit/rest', () => ({
Octokit: MockOctokit,
}));

vi.mock('@octokit/plugin-throttling', () => ({
throttling: vi.fn(),
}));

vi.mock('@octokit/request', () => ({
request: {
defaults: vi.fn().mockReturnValue(vi.fn()),
},
}));

function setupAuthMocks() {
mockGetCredentials.mockResolvedValue([{ appId: 12345, privateKey: 'fake-private-key' }]);

// App auth returns app token
const mockAuth = vi.fn();
mockAuth.mockImplementation((opts: { type: string }) => {
if (opts.type === 'app') {
return Promise.resolve({ token: 'app-token' });
}
return Promise.resolve({ token: 'installation-token' });
});
mockCreateAppAuth.mockReturnValue(mockAuth);
}

describe('multi-App installation clients', () => {
beforeEach(() => {
vi.clearAllMocks();
resetAppCredentialsCache();
mockedCreateCommonStorage.mockReturnValue({ githubAppCredentials: credentialsStore });
setupAuthMocks();
});

it('creates an organization installation client for scheduled registration cleanup', async () => {
mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 999 } });
const client = await createRunnerInstallationClient('test-org', 'Org', '');
expect(client.actions).toBe(mockActions);
expect(mockApps.getOrgInstallation).toHaveBeenCalledWith({ org: 'test-org' });
expect(mockCreateAppAuth).toHaveBeenCalledWith({
appId: 12345,
privateKey: 'fake-private-key',
installationId: 999,
});
});

it('uses additional credentials and keeps App and installation auth paired', async () => {
mockGetCredentials.mockResolvedValue([
{ appId: 1, privateKey: 'one' },
{ appId: 2, privateKey: 'two' },
]);
mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 222 } });
const random = vi.spyOn(Math, 'random').mockReturnValue(0.99);
try {
await createRunnerInstallationClient('test-org', 'Org', '');
expect(mockCreateAppAuth).toHaveBeenNthCalledWith(1, { appId: 2, privateKey: 'two' });
expect(mockCreateAppAuth).toHaveBeenNthCalledWith(2, { appId: 2, privateKey: 'two', installationId: 222 });
} finally {
random.mockRestore();
}
});

it('selects another App after an exhausted installation response', async () => {
mockGetCredentials.mockResolvedValue([
{ appId: 1, privateKey: 'one' },
{ appId: 2, privateKey: 'two' },
]);
mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 222 } });
const random = vi.spyOn(Math, 'random').mockReturnValue(0);
try {
await createRunnerInstallationClient('test-org', 'Org', '');
mockHookAfter.mock.calls[0][1]({ headers: { 'x-ratelimit-remaining': '0' } });
mockCreateAppAuth.mockClear();
await createRunnerInstallationClient('test-org', 'Org', '');
expect(mockCreateAppAuth).toHaveBeenNthCalledWith(1, { appId: 2, privateKey: 'two' });
} finally {
random.mockRestore();
}
});

it('skips an App in secondary-limit cooldown', async () => {
mockGetCredentials.mockResolvedValue([
{ appId: 1, privateKey: 'one' },
{ appId: 2, privateKey: 'two' },
]);
mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 222 } });
createThrottleOptions(1).onSecondaryRateLimit(60, { method: 'GET', url: '/runners' } as Required<EndpointDefaults>);
const random = vi.spyOn(Math, 'random').mockReturnValue(0);
try {
await createRunnerInstallationClient('test-org', 'Org', '');
expect(mockCreateAppAuth).toHaveBeenNthCalledWith(1, { appId: 2, privateKey: 'two' });
} finally {
random.mockRestore();
}
});

it('tries another configured App when installation authentication fails', async () => {
mockGetCredentials.mockResolvedValue([
{ appId: 1, privateKey: 'one' },
{ appId: 2, privateKey: 'two' },
]);
mockApps.getOrgInstallation
.mockRejectedValueOnce(new Error('not installed'))
.mockResolvedValue({ data: { id: 222 } });
const random = vi.spyOn(Math, 'random').mockReturnValue(0);
try {
await createRunnerInstallationClient('test-org', 'Org', '');
expect(mockCreateAppAuth).toHaveBeenLastCalledWith({ appId: 2, privateKey: 'two', installationId: 222 });
} finally {
random.mockRestore();
}
});
});
Loading
Loading