refactor(storage): resolve providers at operation boundaries

Author: edersonbrilhante

lambdas/functions/control-plane/src/control-plane-providers.ts

@@ -1,8 +1,21 @@
 import { createControlPlaneProviderRegistry } from '@aws-github-runner/compute-providers/control-plane';
 import { computeProviderTypes } from '@aws-github-runner/compute-providers/provider-types';
+import type { CreateStartRunnerConfig } from '@aws-github-runner/compute-providers/core';
+import type { StorageProviderType } from '@aws-github-runner/storage-providers';
 
 import { createStartRunnerConfig } from './scale-runners/github-runner';
 
-export const controlPlaneProviderRegistry = createControlPlaneProviderRegistry(createStartRunnerConfig);
+const providerRegistries = new Map<StorageProviderType, ReturnType<typeof createControlPlaneProviderRegistry>>();
+
+export function getControlPlaneProviderRegistry(storageProviderType: StorageProviderType) {
+  let registry = providerRegistries.get(storageProviderType);
+  if (!registry) {
+    const boundCreateStartRunnerConfig: CreateStartRunnerConfig = (...args) =>
+      createStartRunnerConfig(storageProviderType, ...args);
+    registry = createControlPlaneProviderRegistry(boundCreateStartRunnerConfig);
+    providerRegistries.set(storageProviderType, registry);
+  }
+  return registry;
+}
 
 export { computeProviderTypes };

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

@@ -9,7 +9,7 @@ import * as nock from 'nock';
 import {
   createGithubAppAuth,
   createOctokitClient,
-  getStoredAppId,
+  getAppId,
   getStoredInstallationId,
   onRateLimit,
   onSecondaryRateLimit,
@@ -34,6 +34,7 @@ const ENVIRONMENT = 'dev';
 const GITHUB_APP_ID = '1';
 const PARAMETER_GITHUB_APP_ID_NAME = `/actions-runner/${ENVIRONMENT}/github_app_id`;
 const PARAMETER_GITHUB_APP_KEY_BASE64_NAME = `/actions-runner/${ENVIRONMENT}/github_app_key_base64`;
+const storageProviderType = 'aws_ssm' as const;
 
 const mockedGetParameters = vi.mocked(getParameters);
 
@@ -90,7 +91,7 @@ describe('Test createGithubAppAuth', () => {
   it('Throws early when PARAMETER_GITHUB_APP_ID_NAME is not set', async () => {
     delete process.env.PARAMETER_GITHUB_APP_ID_NAME;
 
-    await expect(createGithubAppAuth(installationId)).rejects.toThrow(
+    await expect(createGithubAppAuth(storageProviderType, installationId)).rejects.toThrow(
       'Environment variable PARAMETER_GITHUB_APP_ID_NAME is not set',
     );
     expect(mockedGetParameters).not.toHaveBeenCalled();
@@ -99,7 +100,7 @@ describe('Test createGithubAppAuth', () => {
   it('Throws early when PARAMETER_GITHUB_APP_KEY_BASE64_NAME is not set', async () => {
     delete process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME;
 
-    await expect(createGithubAppAuth(installationId)).rejects.toThrow(
+    await expect(createGithubAppAuth(storageProviderType, installationId)).rejects.toThrow(
       'Environment variable PARAMETER_GITHUB_APP_KEY_BASE64_NAME is not set',
     );
     expect(mockedGetParameters).not.toHaveBeenCalled();
@@ -120,7 +121,7 @@ describe('Test createGithubAppAuth', () => {
     mockedCreatAppAuth.mockReturnValue(mockWithHook);
 
     // Act
-    await createGithubAppAuth(installationId);
+    await createGithubAppAuth(storageProviderType, installationId);
 
     // Assert
     expect(mockedCreatAppAuth).toBeCalledTimes(1);
@@ -155,7 +156,7 @@ describe('Test createGithubAppAuth', () => {
     });
 
     // Act
-    await createGithubAppAuth(installationId);
+    await createGithubAppAuth(storageProviderType, installationId);
 
     // Generate two JWTs and verify they are different (jti makes them unique)
     const jwt1 = await capturedCreateJwt!(1);
@@ -192,7 +193,7 @@ describe('Test createGithubAppAuth', () => {
     mockedCreatAppAuth.mockReturnValue(mockWithHook);
 
     // Act
-    const result = await createGithubAppAuth(installationId);
+    const result = await createGithubAppAuth(storageProviderType, installationId);
 
     // Assert
     expect(getParameters).toBeCalledWith([PARAMETER_GITHUB_APP_ID_NAME, PARAMETER_GITHUB_APP_KEY_BASE64_NAME]);
@@ -216,7 +217,7 @@ describe('Test createGithubAppAuth', () => {
     mockedCreatAppAuth.mockReturnValue(mockWithHook);
 
     // Act
-    const result = await createGithubAppAuth(installationId);
+    const result = await createGithubAppAuth(storageProviderType, installationId);
 
     // Assert
     expect(getParameters).toBeCalledWith([PARAMETER_GITHUB_APP_ID_NAME, PARAMETER_GITHUB_APP_KEY_BASE64_NAME]);
@@ -253,7 +254,7 @@ describe('Test createGithubAppAuth', () => {
     });
 
     // Act
-    const result = await createGithubAppAuth(installationId, githubServerUrl);
+    const result = await createGithubAppAuth(storageProviderType, installationId, githubServerUrl);
 
     // Assert
     expect(getParameters).toBeCalledWith([PARAMETER_GITHUB_APP_ID_NAME, PARAMETER_GITHUB_APP_KEY_BASE64_NAME]);
@@ -291,7 +292,7 @@ describe('Test createGithubAppAuth', () => {
     mockedCreatAppAuth.mockReturnValue(mockWithHook);
 
     // Act
-    const result = await createGithubAppAuth(installationId, githubServerUrl);
+    const result = await createGithubAppAuth(storageProviderType, installationId, githubServerUrl);
 
     // Assert
     expect(getParameters).toBeCalledWith([PARAMETER_GITHUB_APP_ID_NAME, PARAMETER_GITHUB_APP_KEY_BASE64_NAME]);
@@ -353,7 +354,7 @@ describe('Test getStoredInstallationId', () => {
       ]),
     );
 
-    const result = await getStoredInstallationId(0);
+    const result = await getStoredInstallationId(storageProviderType, 0);
     expect(result).toBe(12345);
   });
 
@@ -368,11 +369,24 @@ describe('Test getStoredInstallationId', () => {
       ]),
     );
 
-    await expect(getStoredAppId(0)).resolves.toBe(GITHUB_APP_ID);
-    await expect(getStoredInstallationId(0)).resolves.toBe(12345);
+    await expect(getAppId(storageProviderType, 0)).resolves.toBe(GITHUB_APP_ID);
+    await expect(getStoredInstallationId(storageProviderType, 0)).resolves.toBe(12345);
     expect(mockedGetParameters).toHaveBeenCalledTimes(1);
   });
 
+  it('retries credential loading after a cached request rejects', async () => {
+    mockedGetParameters.mockRejectedValueOnce(new Error('temporary storage failure')).mockResolvedValueOnce(
+      new Map([
+        [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID],
+        [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64],
+      ]),
+    );
+
+    await expect(getAppId(storageProviderType)).rejects.toThrow('temporary storage failure');
+    await expect(getAppId(storageProviderType)).resolves.toBe(GITHUB_APP_ID);
+    expect(mockedGetParameters).toHaveBeenCalledTimes(2);
+  });
+
   it('returns undefined when installation ID param is empty', async () => {
     process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME = '';
     mockedGetParameters.mockResolvedValueOnce(
@@ -382,7 +396,7 @@ describe('Test getStoredInstallationId', () => {
       ]),
     );
 
-    const result = await getStoredInstallationId(0);
+    const result = await getStoredInstallationId(storageProviderType, 0);
     expect(result).toBeUndefined();
   });
 
@@ -395,7 +409,7 @@ describe('Test getStoredInstallationId', () => {
       ]),
     );
 
-    const result = await getStoredInstallationId(0);
+    const result = await getStoredInstallationId(storageProviderType, 0);
     expect(result).toBeUndefined();
   });
 
@@ -408,7 +422,7 @@ describe('Test getStoredInstallationId', () => {
       ]),
     );
 
-    const result = await getStoredInstallationId(99);
+    const result = await getStoredInstallationId(storageProviderType, 99);
     expect(result).toBeUndefined();
   });
 
@@ -434,11 +448,11 @@ describe('Test getStoredInstallationId', () => {
     );
 
     // Primary app (index 0) has no stored installation ID
-    const result0 = await getStoredInstallationId(0);
+    const result0 = await getStoredInstallationId(storageProviderType, 0);
     expect(result0).toBeUndefined();
 
     // Additional app (index 1) has stored installation ID
-    const result1 = await getStoredInstallationId(1);
+    const result1 = await getStoredInstallationId(storageProviderType, 1);
     expect(result1).toBe(67890);
   });
 });

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

@@ -23,8 +23,7 @@ import { retry } from '@octokit/plugin-retry';
 import { throttling } from '@octokit/plugin-throttling';
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
 import { controlPlaneStorageProviderRegistry } from '@aws-github-runner/storage-providers/control-plane';
-import { resolveStorageProviderType } from '@aws-github-runner/storage-providers/provider-types';
-import type { GitHubAppCredential } from '@aws-github-runner/storage-providers';
+import type { GitHubAppCredential, StorageProviderType } from '@aws-github-runner/storage-providers';
 import { EndpointDefaults } from '@octokit/types';
 
 const logger = createChildLogger('gh-auth');
@@ -71,39 +70,46 @@ export function onSecondaryRateLimit(
   return retryCount < MAX_SECONDARY_RATE_LIMIT_RETRIES;
 }
 
-let appCredentialsPromise: Promise<GitHubAppCredential[]> | null = null;
-
-async function loadAppCredentials(): Promise<GitHubAppCredential[]> {
-  const storageProviderType = resolveStorageProviderType(process.env.RUNNER_STORAGE_PROVIDER);
-  const credentialsReader = controlPlaneStorageProviderRegistry.capability(
-    storageProviderType,
-    'githubAppCredentialsReader',
-  )();
-  const credentials = await credentialsReader.read();
-  logger.info(`Loaded ${credentials.length} GitHub App credential(s)`);
-  return credentials;
-}
-
-function getAppCredentials(): Promise<GitHubAppCredential[]> {
-  if (!appCredentialsPromise) appCredentialsPromise = loadAppCredentials();
-  return appCredentialsPromise;
-}
+const appCredentialsPromises = new Map<StorageProviderType, Promise<GitHubAppCredential[]>>();
+
+async function getAppCredentials(storageProviderType: StorageProviderType): Promise<GitHubAppCredential[]> {
+  let credentialsPromise = appCredentialsPromises.get(storageProviderType);
+  if (!credentialsPromise) {
+    const credentialsReader = controlPlaneStorageProviderRegistry.capability(
+      storageProviderType,
+      'githubAppCredentialsReader',
+    )();
+    credentialsPromise = credentialsReader.read().then((credentials) => {
+      logger.info(`Loaded ${credentials.length} GitHub App credential(s)`);
+      return credentials;
+    });
+    appCredentialsPromises.set(storageProviderType, credentialsPromise);
+  }
 
-export async function getAppCount(): Promise<number> {
-  return (await getAppCredentials()).length;
+  try {
+    return await credentialsPromise;
+  } catch (error) {
+    if (appCredentialsPromises.get(storageProviderType) === credentialsPromise) {
+      appCredentialsPromises.delete(storageProviderType);
+    }
+    throw error;
+  }
 }
 
 export function resetAppCredentialsCache(): void {
-  appCredentialsPromise = null;
+  appCredentialsPromises.clear();
 }
 
-export async function getStoredInstallationId(appIndex: number): Promise<number | undefined> {
-  const credentials = await getAppCredentials();
+export async function getStoredInstallationId(
+  storageProviderType: StorageProviderType,
+  appIndex: number,
+): Promise<number | undefined> {
+  const credentials = await getAppCredentials(storageProviderType);
   return credentials[appIndex]?.installationId;
 }
 
-export async function getStoredAppId(appIndex = 0): Promise<string> {
-  const credential = (await getAppCredentials())[appIndex];
+export async function getAppId(storageProviderType: StorageProviderType, appIndex = 0): Promise<string> {
+  const credential = (await getAppCredentials(storageProviderType))[appIndex];
   if (!credential) {
     throw new Error(`GitHub App credential at index ${appIndex} not found`);
   }
@@ -142,25 +148,35 @@ export async function createOctokitClient(token: string, ghesApiUrl = ''): Promi
 }
 
 export async function createGithubAppAuth(
+  storageProviderType: StorageProviderType,
   installationId: number | undefined,
   ghesApiUrl = '',
   appIndex?: number,
 ): Promise<AppAuthentication & { appIndex: number }> {
-  const credentials = await getAppCredentials();
+  const credentials = await getAppCredentials(storageProviderType);
   const idx = appIndex ?? Math.floor(Math.random() * credentials.length);
-  const auth = await createAuth(installationId, ghesApiUrl, idx);
+  const selected = credentials[idx];
+  if (!selected) {
+    throw new Error(`GitHub App credential at index ${idx} not found`);
+  }
+  const auth = createAuth(selected, installationId, ghesApiUrl);
   const result = await auth({ type: 'app' });
   return { ...result, appIndex: idx };
 }
 
 export async function createGithubInstallationAuth(
+  storageProviderType: StorageProviderType,
   installationId: number | undefined,
   ghesApiUrl = '',
   appIndex?: number,
 ): Promise<InstallationAccessTokenAuthentication> {
-  const credentials = await getAppCredentials();
+  const credentials = await getAppCredentials(storageProviderType);
   const idx = appIndex ?? Math.floor(Math.random() * credentials.length);
-  const auth = await createAuth(installationId, ghesApiUrl, idx);
+  const selected = credentials[idx];
+  if (!selected) {
+    throw new Error(`GitHub App credential at index ${idx} not found`);
+  }
+  const auth = createAuth(selected, installationId, ghesApiUrl);
   return auth({ type: 'installation', installationId });
 }
 
@@ -172,15 +188,11 @@ function signJwt(payload: Record<string, unknown>, privateKey: string): string {
   return `${message}.${signature}`;
 }
 
-async function createAuth(
+function createAuth(
+  selected: GitHubAppCredential,
   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)];
-
+): AuthInterface {
   logger.debug(`Selected GitHub App ${selected.appId} for authentication`);
 
   // Use a custom createJwt callback to include a jti (JWT ID) claim in every token.