github-aws-runners · iainlane · Mar 14, 2025
@@ -51,6 +51,7 @@
     "@octokit/rest": "20.1.2",
     "@octokit/types": "^13.8.0",
     "cron-parser": "^4.9.0",
+    "toad-cache": "^3.7.0",
     "typescript": "^5.7.3"
   },
   "nx": {

@@ -7,7 +7,7 @@ import { mocked } from 'jest-mock';
 import { MockProxy, mock } from 'jest-mock-extended';
 import nock from 'nock';
 
-import { createGithubAppAuth, createOctokitClient } from './auth';
+import { createGithubAppAuth, createOctokitClient } from './client';
 
 jest.mock('@aws-github-runner/aws-ssm-util');
 jest.mock('@octokit/auth-app');

@@ -0,0 +1,211 @@
+import { createAppAuth } from '@octokit/auth-app';
+import { StrategyOptions } from '@octokit/auth-app/dist-types/types';
+import { OctokitOptions } from '@octokit/core/dist-types/types';
+import { Octokit } from '@octokit/rest';
+import { throttling } from '@octokit/plugin-throttling';
+import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
+import { getParameter } from '@aws-github-runner/aws-ssm-util';
+import { EndpointDefaults, type OctokitResponse } from '@octokit/types';
+import { RequestError } from '@octokit/request-error';
+import { Lru } from 'toad-cache';
+import type { ActionRequestMessage } from '../scale-runners/scale-up';
+
+const logger = createChildLogger('gh-auth');
+
+interface CacheEntry {
+  etag?: string;
+  lastModified?: string;
+  [key: string]: unknown;
+}
+
+
+// Cache for conditional requests
+// Using 15000 entries with 5 minute TTL
+const cache = new Lru<CacheEntry>(15000, 5 * 60 * 1000);
+
+async function appAuthCommonParameters(): Promise<StrategyOptions> {
+  const appId = parseInt(await getParameter(process.env.PARAMETER_GITHUB_APP_ID_NAME));
+
+  return {
+    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)),
+  };
+}
+
+/**
+ * Handler run before requests are sent. Looks up the URL in the cache, and adds
+ * headers for conditional retrieval if there is an entry.
+ */
+async function beforeRequestHandler(octokit: Octokit, options: Required<EndpointDefaults>): Promise<void> {
+  const { method } = options;
+
+  if (method !== 'GET') {
+    return;
+  }
+
+  const { url } = octokit.request.endpoint.parse(options);
+  const cacheKey = url;
+  const cacheEntry = cache.get(cacheKey);
+
+  if (cacheEntry === undefined) {
+    logger.info('Cache miss', { url });
+    return;
+  }
+
+  const { etag, lastModified } = cacheEntry;
+
+  if (etag !== undefined) {
+    options.headers['If-None-Match'] = etag;
+  }
+
+  if (lastModified !== undefined) {
+    options.headers['If-Modified-Since'] = lastModified;
+  }
+
+  logger.info('Cache hit', { url, etag, lastModified });
+}
+
+/**
+ * Handler run after requests are sent. Caches the response if it has an ETag or
+ * Last-Modified header, so that it can be returned by future conditional
+ * requests if requested again.
+ */
+async function afterRequestHandler(octokit: Octokit, response: OctokitResponse<any, number>, options: Required<EndpointDefaults>): Promise<void> {
+  const { status } = response;
+  const { url } = octokit.request.endpoint.parse(options);
+  logger.info(`Response received`, { status, url });
+
+  const cacheKey = url;
+  const eTag = response.headers.etag;
+  const lastModified = response.headers['last-modified'];
+
+  if (eTag === undefined && lastModified === undefined) {
+    return;
+  }
+
+  logger.info('Caching response', { url, eTag, lastModified });
+
+  cache.set(cacheKey, {
+    ...(eTag !== undefined ? { etag: eTag } : {}),
+    ...(lastModified !== undefined ? { lastModified } : {}),
+    ...response,
+  });
+}
+
+/**
+ * Handler run if a request fails. This handler is called for any non-2xx
+ * response. We will get "304 Not Modified" responses when the conditional
+ * request is satisfied, and we should return the cached data in that case.
+ */
+async function errorRequestHandler(octokit: Octokit, error: Error, options: Required<EndpointDefaults>): Promise<CacheEntry> {
+  if (!(error instanceof RequestError)) {
+    throw error;
+  }
+
+  const { status } = error;
+
+  if (status != 304) {
+    throw error;
+  }
+
+  const { url } = octokit.request.endpoint.parse(options);
+
+  const entry = cache.get(url);
+
+  if (entry === undefined) {
+      throw new Error(`Received 304 Not Modified response for ${url}, but it wasn't found in the cache.`);
+  }
+
+  return entry;
+}
+
+export async function createAppAuthClient(ghesApiUrl: string = ''): Promise<Octokit> {
+  const CustomOctokit = Octokit.plugin(throttling);
+
+  const octokit = new CustomOctokit({
+    authStrategy: createAppAuth,
+    auth: await appAuthCommonParameters(),
+    baseUrl: ghesApiUrl || undefined,
+    previews: ghesApiUrl ? ['antiope'] : undefined,
+    throttle: {
+      onRateLimit: (retryAfter: number, options: Required<EndpointDefaults>) => {
+        logger.warn(`GitHub rate limit: Request quota exhausted for request ${options.method} ${options.url}.`, {
+          retryAfter,
+        });
+      },
+      onSecondaryRateLimit: (retryAfter: number, options: Required<EndpointDefaults>) => {
+        logger.warn(`GitHub rate limit: SecondaryRateLimit detected for request ${options.method} ${options.url}`, {
+          retryAfter,
+        });
+      },
+    },
+    userAgent: process.env.USER_AGENT || 'github-aws-runners',
+  });
+
+  octokit.hook.before('request', async (options) => beforeRequestHandler(octokit, options));
+  octokit.hook.after('request', async (response, options) => afterRequestHandler(octokit, response, options));
+  octokit.hook.error('request', async (error, options) => errorRequestHandler(octokit, error, options));
+
+  return octokit;
+}
+
+async function getInstallationId(
+  appClient: Octokit,
+  enableOrgLevel: boolean,
+  payload: ActionRequestMessage,
+): Promise<number> {
+  if (payload.installationId !== 0) {
+    return payload.installationId;
+  }
+
+  return (
+    enableOrgLevel
+      ? await appClient.apps.getOrgInstallation({
+          org: payload.repositoryOwner,
+        })
+      : await appClient.apps.getRepoInstallation({
+          owner: payload.repositoryOwner,
+          repo: payload.repositoryName,
+        })
+  ).data.id;
+}
+
+export async function createAppInstallationClient(appOctokit: Octokit, enableOrgLevel: boolean, payload: ActionRequestMessage): Promise<Octokit> {
+  const installationId = await getInstallationId(appOctokit, enableOrgLevel, payload);
+
+  return appOctokit.auth({
+    type: 'installation',
+    installationId,
+    factory: ({ octokitOptions, ...auth }: { octokitOptions: OctokitOptions }) =>
+      new Octokit({
+        ...octokitOptions,
+        auth: auth,
+      }),
+  }) as Promise<Octokit>;
+}
+
+export function getGitHubEnterpriseApiUrl() {
+  const ghesBaseUrl = process.env.GHES_URL;
+  let ghesApiUrl = '';
+  if (ghesBaseUrl) {
+    const url = new URL(ghesBaseUrl);
+    const domain = url.hostname;
+    if (domain.endsWith('.ghe.com')) {
+      // Data residency: Prepend 'api.'
+      ghesApiUrl = `https://api.${domain}`;
+    } else {
+      // GitHub Enterprise Server: Append '/api/v3'
+      ghesApiUrl = `${ghesBaseUrl}/api/v3`;
+    }
+  }
+  logger.debug(`Github Enterprise URLs: api_url - ${ghesApiUrl}; base_url - ${ghesBaseUrl}`);
+  return { ghesApiUrl, ghesBaseUrl };
+}
@@ -9,6 +9,11 @@ import { scaleDown } from './scale-runners/scale-down';
 import { scaleUp } from './scale-runners/scale-up';
 import { SSMCleanupOptions, cleanSSMTokens } from './scale-runners/ssm-housekeeper';
 import { checkAndRetryJob } from './scale-runners/job-retry';
+import { createAppAuthClient, getGitHubEnterpriseApiUrl } from './github/client';
+
+const { ghesApiUrl } = getGitHubEnterpriseApiUrl();
+// TODO: needs to be ESM for top-level await, or we create this lazily.
+const ghAppClient = await createAppAuthClient(ghesApiUrl);
 
 export async function scaleUpHandler(event: SQSEvent, context: Context): Promise<void> {
   setContext(context, 'lambda.ts');
@@ -20,7 +25,7 @@ export async function scaleUpHandler(event: SQSEvent, context: Context): Promise
   }
 
   try {
-    await scaleUp(event.Records[0].eventSource, JSON.parse(event.Records[0].body));
+    await scaleUp(ghAppClient, event.Records[0].eventSource, JSON.parse(event.Records[0].body));
   } catch (e) {
     if (e instanceof ScaleError) {
       throw e;

@@ -4,7 +4,7 @@ import moment from 'moment-timezone';
 import nock from 'nock';
 
 import { listEC2Runners } from '../aws/runners';
-import * as ghAuth from '../github/auth';
+import * as ghAuth from '../github/client';
 import { createRunners, getGitHubEnterpriseApiUrl } from '../scale-runners/scale-up';
 import { adjust } from './pool';