feat(github-app): allow several github apps to be used

Author: thomasnemer

README.md

@@ -104,6 +104,7 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 
 | Name | Description | Type | Default | Required |
 |------|-------------|------|---------|:--------:|
+| <a name="input_additional_github_apps"></a> [additional\_github\_apps](#input\_additional\_github\_apps) | Additional GitHub Apps for distributing API rate limit usage. Each must be installed on the same repos/orgs as the primary app. | <pre>list(object({<br/>    key_base64          = optional(string)<br/>    key_base64_ssm      = optional(object({ arn = string, name = string }))<br/>    id                  = optional(string)<br/>    id_ssm              = optional(object({ arn = string, name = string }))<br/>    installation_id     = optional(string)<br/>    installation_id_ssm = optional(object({ arn = string, name = string }))<br/>  }))</pre> | `[]` | no |
 | <a name="input_ami"></a> [ami](#input\_ami) | AMI configuration for the action runner instances. This object allows you to specify all AMI-related settings in one place.<br/><br/>Parameters:<br/>- `filter`: Map of lists to filter AMIs by various criteria (e.g., { name = ["ubuntu/images/hvm-ssd/ubuntu-jammy-22.04-amd64-*"], state = ["available"] })<br/>- `owners`: List of AMI owners to limit the search. Common values: ["amazon"], ["self"], or specific AWS account IDs<br/>- `id_ssm_parameter_arn`: ARN of an SSM parameter containing the AMI ID. If specified, this overrides both AMI filter and parameter name<br/>- `kms_key_arn`: Optional KMS key ARN if the AMI is encrypted with a customer managed key<br/><br/>Defaults to null, in which case the module falls back to individual AMI variables (deprecated). | <pre>object({<br/>    filter               = optional(map(list(string)), { state = ["available"] })<br/>    owners               = optional(list(string), ["amazon"])<br/>    id_ssm_parameter_arn = optional(string, null)<br/>    kms_key_arn          = optional(string, null)<br/>  })</pre> | `null` | no |
 | <a name="input_ami_housekeeper_cleanup_config"></a> [ami\_housekeeper\_cleanup\_config](#input\_ami\_housekeeper\_cleanup\_config) | Configuration for AMI cleanup.<br/><br/>    `amiFilters` - Filters to use when searching for AMIs to cleanup. Default filter for images owned by the account and that are available.<br/>    `dryRun` - If true, no AMIs will be deregistered. Default false.<br/>    `launchTemplateNames` - Launch template names to use when searching for AMIs to cleanup. Default no launch templates.<br/>    `maxItems` - The maximum number of AMIs that will be queried for cleanup. Default no maximum.<br/>    `minimumDaysOld` - Minimum number of days old an AMI must be to be considered for cleanup. Default 30.<br/>    `ssmParameterNames` - SSM parameter names to use when searching for AMIs to cleanup. This parameter should be set when using SSM to configure the AMI to use. Default no SSM parameters. | <pre>object({<br/>    amiFilters = optional(list(object({<br/>      Name   = string<br/>      Values = list(string)<br/>      })),<br/>      [{<br/>        Name : "state",<br/>        Values : ["available"],<br/>        },<br/>        {<br/>          Name : "image-type",<br/>          Values : ["machine"],<br/>      }]<br/>    )<br/>    dryRun              = optional(bool, false)<br/>    launchTemplateNames = optional(list(string))<br/>    maxItems            = optional(number)<br/>    minimumDaysOld      = optional(number, 30)<br/>    ssmParameterNames   = optional(list(string))<br/>  })</pre> | `{}` | no |
 | <a name="input_ami_housekeeper_lambda_s3_key"></a> [ami\_housekeeper\_lambda\_s3\_key](#input\_ami\_housekeeper\_lambda\_s3\_key) | S3 key for syncer lambda function. Required if using S3 bucket to specify lambdas. | `string` | `null` | no |

examples/multi-runner/README.md

@@ -16,6 +16,10 @@ For exact match, all the labels defined in the workflow should be present in the
 
 For the list of provided runner configurations, there will be a single webhook and only a single GitHub App to receive the notifications for all types of workflow triggers.
 
+## Multiple GitHub Apps (rate limit distribution)
+
+This example also shows how to optionally configure multiple GitHub Apps via the `additional_github_apps` variable. When configured, the control-plane lambdas (scale-up, scale-down, pool, job-retry) randomly select an app for each GitHub API call, spreading the rate limit usage across all apps. Only the primary app needs a webhook URL configured in GitHub.
+
 ## Lambda distribution
 
 Per combination of OS and architecture a lambda distribution syncer will be created. For this example there will be three instances (windows X64, linux X64, linux ARM).

examples/multi-runner/main.tf

@@ -117,6 +117,17 @@ module "runners" {
     webhook_secret = random_id.random.hex
   }
 
+  # Uncomment to distribute GitHub API rate limit usage across multiple GitHub Apps.
+  # Each additional app must be installed on the same repos/orgs as the primary app.
+  # The control-plane lambdas will randomly select an app for each API call.
+  # additional_github_apps = [
+  #   {
+  #     key_base64      = var.additional_github_app_0.key_base64
+  #     id              = var.additional_github_app_0.id
+  #     installation_id = var.additional_github_app_0.installation_id  # optional, avoids an API call
+  #   },
+  # ]
+
   # Deploy webhook using the EventBridge
   eventbridge = {
     enable = true

lambdas/functions/control-plane/src/github/auth.test.ts

@@ -5,7 +5,7 @@ import { RequestInterface, RequestParameters } from '@octokit/types';
 import { getParameter } from '@aws-github-runner/aws-ssm-util';
 import * as nock from 'nock';
 
-import { createGithubAppAuth, createOctokitClient } from './auth';
+import { createGithubAppAuth, createOctokitClient, getStoredInstallationId, resetAppCredentialsCache } from './auth';
 import { describe, it, expect, beforeEach, vi } from 'vitest';
 
 type MockProxy<T> = T & {
@@ -31,6 +31,7 @@ const mockedGet = vi.mocked(getParameter);
 beforeEach(() => {
   vi.resetModules();
   vi.clearAllMocks();
+  resetAppCredentialsCache();
   process.env = { ...cleanEnv };
   process.env.PARAMETER_GITHUB_APP_ID_NAME = PARAMETER_GITHUB_APP_ID_NAME;
   process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME = PARAMETER_GITHUB_APP_KEY_BASE64_NAME;
@@ -207,3 +208,76 @@ ${decryptedValue}`,
     expect(result.token).toBe(token);
   });
 });
+
+describe('Test getStoredInstallationId', () => {
+  const decryptedValue = 'decryptedValue';
+  const b64 = Buffer.from(decryptedValue, 'binary').toString('base64');
+
+  beforeEach(() => {
+    const mockedAuth = vi.fn();
+    mockedAuth.mockResolvedValue({ token: 'token' });
+    const mockWithHook = Object.assign(mockedAuth, { hook: vi.fn() });
+    vi.mocked(createAppAuth).mockReturnValue(mockWithHook);
+  });
+
+  it('returns stored installation ID when configured', async () => {
+    const installationIdParam = `/actions-runner/${ENVIRONMENT}/github_app_installation_id`;
+    process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME = installationIdParam;
+    mockedGet.mockResolvedValueOnce(GITHUB_APP_ID).mockResolvedValueOnce(b64).mockResolvedValueOnce('12345');
+
+    const result = await getStoredInstallationId(0);
+    expect(result).toBe(12345);
+    expect(getParameter).toHaveBeenCalledWith(installationIdParam);
+  });
+
+  it('returns undefined when installation ID param is empty', async () => {
+    process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME = '';
+    mockedGet.mockResolvedValueOnce(GITHUB_APP_ID).mockResolvedValueOnce(b64);
+
+    const result = await getStoredInstallationId(0);
+    expect(result).toBeUndefined();
+  });
+
+  it('returns undefined when env var is not set', async () => {
+    delete process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME;
+    mockedGet.mockResolvedValueOnce(GITHUB_APP_ID).mockResolvedValueOnce(b64);
+
+    const result = await getStoredInstallationId(0);
+    expect(result).toBeUndefined();
+  });
+
+  it('returns undefined for out-of-bounds appIndex', async () => {
+    process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME = '';
+    mockedGet.mockResolvedValueOnce(GITHUB_APP_ID).mockResolvedValueOnce(b64);
+
+    const result = await getStoredInstallationId(99);
+    expect(result).toBeUndefined();
+  });
+
+  it('loads installation IDs for multi-app setup', async () => {
+    const app1IdParam = `/actions-runner/${ENVIRONMENT}/github_app_id`;
+    const app2IdParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_id`;
+    const app1KeyParam = `/actions-runner/${ENVIRONMENT}/github_app_key_base64`;
+    const app2KeyParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_key_base64`;
+    const app2InstallParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_installation_id`;
+
+    process.env.PARAMETER_GITHUB_APP_ID_NAME = `${app1IdParam}:${app2IdParam}`;
+    process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME = `${app1KeyParam}:${app2KeyParam}`;
+    process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME = `:${app2InstallParam}`;
+
+    mockedGet
+      .mockResolvedValueOnce('1') // app1 id
+      .mockResolvedValueOnce(b64) // app1 key
+      .mockResolvedValueOnce('2') // app2 id
+      .mockResolvedValueOnce(b64) // app2 key
+      .mockResolvedValueOnce('67890'); // app2 installation id
+
+    // Primary app (index 0) has no stored installation ID
+    const result0 = await getStoredInstallationId(0);
+    expect(result0).toBeUndefined();
+
+    // Additional app (index 1) has stored installation ID
+    const result1 = await getStoredInstallationId(1);
+    expect(result1).toBe(67890);
+  });
+});

lambdas/functions/control-plane/src/github/auth.ts

@@ -25,6 +25,56 @@ import { EndpointDefaults } from '@octokit/types';
 
 const logger = createChildLogger('gh-auth');
 
+interface GitHubAppCredential {
+  appId: number;
+  privateKey: string;
+  installationId?: number;
+}
+
+let appCredentialsPromise: Promise<GitHubAppCredential[]> | null = null;
+
+async function loadAppCredentials(): Promise<GitHubAppCredential[]> {
+  const idParams = process.env.PARAMETER_GITHUB_APP_ID_NAME.split(':').filter(Boolean);
+  const keyParams = process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME.split(':').filter(Boolean);
+  const installationIdParams = (process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME || '').split(':');
+  if (idParams.length !== keyParams.length) {
+    throw new Error(`GitHub App parameter count mismatch: ${idParams.length} IDs vs ${keyParams.length} keys`);
+  }
+  const credentials: GitHubAppCredential[] = [];
+  for (let i = 0; i < idParams.length; i++) {
+    const appId = parseInt(await getParameter(idParams[i]));
+    const privateKey = Buffer.from(await getParameter(keyParams[i]), 'base64')
+      .toString()
+      .replace('/[\\n]/g', String.fromCharCode(10));
+    const installationIdParam = installationIdParams[i];
+    const installationId =
+      installationIdParam && installationIdParam.length > 0
+        ? parseInt(await getParameter(installationIdParam))
+        : undefined;
+    credentials.push({ appId, privateKey, installationId });
+  }
+  logger.info(`Loaded ${credentials.length} GitHub App credential(s)`);
+  return credentials;
+}
+
+function getAppCredentials(): Promise<GitHubAppCredential[]> {
+  if (!appCredentialsPromise) appCredentialsPromise = loadAppCredentials();
+  return appCredentialsPromise;
+}
+
+export async function getAppCount(): Promise<number> {
+  return (await getAppCredentials()).length;
+}
+
+export function resetAppCredentialsCache(): void {
+  appCredentialsPromise = null;
+}
+
+export async function getStoredInstallationId(appIndex: number): Promise<number | undefined> {
+  const credentials = await getAppCredentials();
+  return credentials[appIndex]?.installationId;
+}
+
 export async function createOctokitClient(token: string, ghesApiUrl = ''): Promise<Octokit> {
   const CustomOctokit = Octokit.plugin(throttling);
   const ocktokitOptions: OctokitOptions = {
@@ -54,35 +104,38 @@ export async function createOctokitClient(token: string, ghesApiUrl = ''): Promi
 export async function createGithubAppAuth(
   installationId: number | undefined,
   ghesApiUrl = '',
-): Promise<AppAuthentication> {
-  const auth = await createAuth(installationId, ghesApiUrl);
-  const appAuthOptions: AppAuthOptions = { type: 'app' };
-  return auth(appAuthOptions);
+  appIndex?: number,
+): Promise<AppAuthentication & { appIndex: number }> {
+  const credentials = await getAppCredentials();
+  const idx = appIndex ?? Math.floor(Math.random() * credentials.length);
+  const auth = await createAuth(installationId, ghesApiUrl, idx);
+  const result = await auth({ type: 'app' });
+  return { ...result, appIndex: idx };
 }
 
 export async function createGithubInstallationAuth(
   installationId: number | undefined,
   ghesApiUrl = '',
+  appIndex?: number,
 ): Promise<InstallationAccessTokenAuthentication> {
-  const auth = await createAuth(installationId, ghesApiUrl);
-  const installationAuthOptions: InstallationAuthOptions = { type: 'installation', installationId };
-  return auth(installationAuthOptions);
+  const credentials = await getAppCredentials();
+  const idx = appIndex ?? Math.floor(Math.random() * credentials.length);
+  const auth = await createAuth(installationId, ghesApiUrl, idx);
+  return auth({ type: 'installation', installationId });
 }
 
-async function createAuth(installationId: number | undefined, ghesApiUrl: string): Promise<AuthInterface> {
-  const appId = parseInt(await getParameter(process.env.PARAMETER_GITHUB_APP_ID_NAME));
-  let authOptions: StrategyOptions = {
-    appId,
-    privateKey: Buffer.from(
-      await getParameter(process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME),
-      'base64',
-      // replace literal \n characters with new lines to allow the key to be stored as a
-      // single line variable. This logic should match how the GitHub Terraform provider
-      // processes private keys to retain compatibility between the projects
-    )
-      .toString()
-      .replace('/[\\n]/g', String.fromCharCode(10)),
-  };
+async function createAuth(
+  installationId: number | undefined,
+  ghesApiUrl: string,
+  appIndex?: number,
+): Promise<AuthInterface> {
+  const credentials = await getAppCredentials();
+  const selected =
+    appIndex !== undefined ? credentials[appIndex] : credentials[Math.floor(Math.random() * credentials.length)];
+
+  logger.debug(`Selected GitHub App ${selected.appId} for authentication`);
+
+  let authOptions: StrategyOptions = { appId: selected.appId, privateKey: selected.privateKey };
   if (installationId) authOptions = { ...authOptions, installationId };
 
   logger.debug(`GHES API URL: ${ghesApiUrl}`);

lambdas/functions/control-plane/src/github/octokit.test.ts

@@ -15,7 +15,9 @@ vi.mock('../github/auth', async () => ({
     return { token: 'token', type: 'installation', installationId: installationId };
   }),
   createOctokitClient: vi.fn().mockImplementation(() => new Octokit()),
-  createGithubAppAuth: vi.fn().mockResolvedValue({ token: 'token' }),
+  createGithubAppAuth: vi.fn().mockResolvedValue({ token: 'token', appIndex: 0 }),
+  getAppCount: vi.fn().mockResolvedValue(1),
+  getStoredInstallationId: vi.fn().mockResolvedValue(undefined),
 }));
 
 vi.mock('@octokit/rest', async () => ({

lambdas/functions/control-plane/src/github/octokit.ts

@@ -1,17 +1,32 @@
 import { Octokit } from '@octokit/rest';
 import { ActionRequestMessage } from '../scale-runners/scale-up';
-import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient } from './auth';
+import {
+  createGithubAppAuth,
+  createGithubInstallationAuth,
+  createOctokitClient,
+  getAppCount,
+  getStoredInstallationId,
+} from './auth';
 
 export async function getInstallationId(
   ghesApiUrl: string,
   enableOrgLevel: boolean,
   payload: ActionRequestMessage,
+  appIndex?: number,
 ): Promise<number> {
-  if (payload.installationId !== 0) {
+  // Use pre-stored installation ID when available (avoids an API call)
+  if (appIndex !== undefined) {
+    const storedId = await getStoredInstallationId(appIndex);
+    if (storedId !== undefined) return storedId;
+  }
+
+  const multiApp = (await getAppCount()) > 1;
+
+  if (!multiApp && payload.installationId !== 0) {
     return payload.installationId;
   }
 
-  const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl);
+  const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl, appIndex);
   const githubClient = await createOctokitClient(ghAuth.token, ghesApiUrl);
   return enableOrgLevel
     ? (
@@ -40,7 +55,11 @@ export async function getOctokit(
   enableOrgLevel: boolean,
   payload: ActionRequestMessage,
 ): Promise<Octokit> {
-  const installationId = await getInstallationId(ghesApiUrl, enableOrgLevel, payload);
-  const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl);
-  return await createOctokitClient(ghAuth.token, ghesApiUrl);
+  // Select one app for this entire auth flow
+  const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl);
+  const appIdx = ghAuth.appIndex;
+
+  const installationId = await getInstallationId(ghesApiUrl, enableOrgLevel, payload, appIdx);
+  const installationAuth = await createGithubInstallationAuth(installationId, ghesApiUrl, appIdx);
+  return await createOctokitClient(installationAuth.token, ghesApiUrl);
 }

lambdas/functions/control-plane/src/github/rate-limit.ts

@@ -9,7 +9,8 @@ let appIdPromise: Promise<string> | null = null;
 
 async function getAppId(): Promise<string> {
   if (!appIdPromise) {
-    appIdPromise = getParameter(process.env.PARAMETER_GITHUB_APP_ID_NAME);
+    const paramName = process.env.PARAMETER_GITHUB_APP_ID_NAME.split(':')[0];
+    appIdPromise = getParameter(paramName);
   }
   return appIdPromise;
 }

lambdas/functions/control-plane/src/pool/pool.test.ts

@@ -34,6 +34,7 @@ vi.mock('./../github/auth', async () => ({
   createGithubAppAuth: vi.fn(),
   createGithubInstallationAuth: vi.fn(),
   createOctokitClient: vi.fn(),
+  getStoredInstallationId: vi.fn().mockResolvedValue(undefined),
 }));
 
 vi.mock('../scale-runners/scale-up', async () => ({
@@ -166,6 +167,7 @@ beforeEach(() => {
     token: 'token',
     appId: 1,
     expiresAt: 'some-date',
+    appIndex: 0,
   });
   mockedInstallationAuth.mockResolvedValue({
     type: 'token',
@@ -336,4 +338,46 @@ describe('Test simple pool.', () => {
       expect(createRunners).toHaveBeenCalledWith(expect.anything(), expect.anything(), 1, expect.anything());
     });
   });
+
+  describe('Multi-app round-robin', () => {
+    beforeEach(() => {
+      (getGitHubEnterpriseApiUrl as ReturnType<typeof vi.fn>).mockReturnValue({
+        ghesApiUrl: '',
+        ghesBaseUrl: '',
+      });
+    });
+
+    it('passes the same appIndex to createGithubInstallationAuth', async () => {
+      mockedAppAuth.mockResolvedValue({
+        type: 'app',
+        token: 'token',
+        appId: 42,
+        expiresAt: 'some-date',
+        appIndex: 1,
+      });
+
+      await adjust({ poolSize: 3 });
+
+      expect(mockedInstallationAuth).toHaveBeenCalledWith(
+        expect.any(Number),
+        expect.any(String),
+        1, // appIndex must match the one from createGithubAppAuth
+      );
+    });
+
+    it('looks up installationId using the selected app JWT', async () => {
+      mockedAppAuth.mockResolvedValue({
+        type: 'app',
+        token: 'app-token-for-selected-app',
+        appId: 42,
+        expiresAt: 'some-date',
+        appIndex: 1,
+      });
+
+      await adjust({ poolSize: 3 });
+
+      // Should look up installationId via the API
+      expect(mockOctokit.apps.getOrgInstallation).toHaveBeenCalledWith({ org: ORG });
+    });
+  });
 });