From 83c90a58f246122b2720da33a650b3765a02e35f Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Tue, 18 Aug 2026 22:24:00 +0200 Subject: [PATCH 1/6] refactor(storage): extract runner config store --- lambdas/functions/control-plane/package.json | 1 + .../functions/control-plane/src/modules.d.ts | 1 - .../src/pool/pool-contract.test.ts | 3 + .../control-plane/src/pool/pool.test.ts | 13 +++ .../functions/control-plane/src/pool/pool.ts | 4 +- .../src/scale-runners/github-runner.ts | 50 ++++++----- .../scale-runners/scale-up-contract.test.ts | 3 + .../src/scale-runners/scale-up.test.ts | 18 +++- .../src/scale-runners/scale-up.ts | 4 +- .../ec2/src/control-plane/runner-config.ts | 2 +- .../ec2/src/control-plane/scale-up.test.ts | 3 +- lambdas/libs/compute-providers/core/index.ts | 3 +- .../aws/ssm/environment.d.ts | 10 +++ .../aws/ssm/parameter-store-tags.ts | 42 +++++++++ .../aws/ssm/runner-config-store.test.ts | 88 +++++++++++++++++++ .../aws/ssm/runner-config-store.ts | 37 ++++++++ lambdas/libs/storage-providers/core/index.ts | 14 +++ .../libs/storage-providers/environment.d.ts | 9 ++ lambdas/libs/storage-providers/index.ts | 2 + lambdas/libs/storage-providers/package.json | 29 ++++++ .../storage-providers/runner-config.test.ts | 84 ++++++++++++++++++ .../libs/storage-providers/runner-config.ts | 46 ++++++++++ lambdas/libs/storage-providers/tsconfig.json | 5 ++ .../libs/storage-providers/vitest.config.ts | 14 +++ lambdas/yarn.lock | 9 ++ 25 files changed, 463 insertions(+), 31 deletions(-) create mode 100644 lambdas/libs/storage-providers/aws/ssm/environment.d.ts create mode 100644 lambdas/libs/storage-providers/aws/ssm/parameter-store-tags.ts create mode 100644 lambdas/libs/storage-providers/aws/ssm/runner-config-store.test.ts create mode 100644 lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts create mode 100644 lambdas/libs/storage-providers/core/index.ts create mode 100644 lambdas/libs/storage-providers/environment.d.ts create mode 100644 lambdas/libs/storage-providers/index.ts create mode 100644 lambdas/libs/storage-providers/package.json create mode 100644 lambdas/libs/storage-providers/runner-config.test.ts create mode 100644 lambdas/libs/storage-providers/runner-config.ts create mode 100644 lambdas/libs/storage-providers/tsconfig.json create mode 100644 lambdas/libs/storage-providers/vitest.config.ts diff --git a/lambdas/functions/control-plane/package.json b/lambdas/functions/control-plane/package.json index dc74e770a9..0f443fc849 100644 --- a/lambdas/functions/control-plane/package.json +++ b/lambdas/functions/control-plane/package.json @@ -33,6 +33,7 @@ "@aws-github-runner/aws-powertools-util": "*", "@aws-github-runner/aws-ssm-util": "*", "@aws-github-runner/compute-providers": "*", + "@aws-github-runner/storage-providers": "*", "@aws-lambda-powertools/parameters": "^2.31.0", "@aws-sdk/client-ec2": "^3.1009.0", "@aws-sdk/client-sqs": "^3.1009.0", diff --git a/lambdas/functions/control-plane/src/modules.d.ts b/lambdas/functions/control-plane/src/modules.d.ts index d5157ccb37..84b0d23a02 100644 --- a/lambdas/functions/control-plane/src/modules.d.ts +++ b/lambdas/functions/control-plane/src/modules.d.ts @@ -18,7 +18,6 @@ declare namespace NodeJS { RUNNER_OWNER: string; COMPUTE_PROVIDER_TYPE?: string; SCALE_DOWN_CONFIG: string; - SSM_TOKEN_PATH: string; SSM_CLEANUP_CONFIG: string; SUBNET_IDS: string; INSTANCE_TYPES: string; diff --git a/lambdas/functions/control-plane/src/pool/pool-contract.test.ts b/lambdas/functions/control-plane/src/pool/pool-contract.test.ts index e519e412e4..28e38c6a77 100644 --- a/lambdas/functions/control-plane/src/pool/pool-contract.test.ts +++ b/lambdas/functions/control-plane/src/pool/pool-contract.test.ts @@ -1,5 +1,6 @@ import type { Octokit } from '@octokit/rest'; import type { ComputeProviderType } from '@aws-github-runner/compute-providers/provider-types'; +import { resetRunnerConfigStore } from '@aws-github-runner/storage-providers'; import { beforeEach, vi } from 'vitest'; import { definePoolContractTests } from '../test/compute-provider-contracts/pool'; @@ -48,6 +49,8 @@ const computeProviders = providerTypes.map((type) => ({ beforeEach(() => { vi.clearAllMocks(); process.env = { ...cleanEnv }; + process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/tokens'; + resetRunnerConfigStore(); mockedAppAuth.mockResolvedValue({ type: 'app', token: 'app-token', appId: 1, expiresAt: 'some-date' }); mockedInstallationAuth.mockResolvedValue({ diff --git a/lambdas/functions/control-plane/src/pool/pool.test.ts b/lambdas/functions/control-plane/src/pool/pool.test.ts index 568403c3be..5b7d5bc5e9 100644 --- a/lambdas/functions/control-plane/src/pool/pool.test.ts +++ b/lambdas/functions/control-plane/src/pool/pool.test.ts @@ -4,6 +4,7 @@ import * as nock from 'nock'; import { createRunners } from '@aws-github-runner/compute-providers/aws/ec2/control-plane/runner-config'; import { listEC2Runners } from '@aws-github-runner/compute-providers/aws/ec2/control-plane/runners'; +import { resetRunnerConfigStore } from '@aws-github-runner/storage-providers'; import * as ghAuth from '../github/auth'; import { getGitHubEnterpriseApiUrl } from '../scale-runners/github-runner'; import { adjust } from './pool'; @@ -134,6 +135,7 @@ beforeEach(() => { vi.resetModules(); vi.clearAllMocks(); process.env = { ...cleanEnv }; + resetRunnerConfigStore(); process.env.GITHUB_APP_KEY_BASE64 = 'TEST_CERTIFICATE_DATA'; process.env.GITHUB_APP_ID = '1337'; process.env.GITHUB_APP_CLIENT_ID = 'TEST_CLIENT_ID'; @@ -253,6 +255,17 @@ describe('Test simple pool.', () => { expect(mockListRunners).not.toHaveBeenCalled(); }); + it('Rejects an unsupported runner config store before GitHub or runner lookups.', async () => { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'unsupported-provider'; + + await expect(adjust({ poolSize: 10, type: 'ec2' })).rejects.toThrow( + "Unsupported runner config storage provider 'unsupported-provider'", + ); + + expect(mockedAppAuth).not.toHaveBeenCalled(); + expect(mockListRunners).not.toHaveBeenCalled(); + }); + it('Should not top up if pool size is reached.', async () => { await adjust({ poolSize: 1, type: 'ec2' }); expect(createRunners).not.toHaveBeenCalled(); diff --git a/lambdas/functions/control-plane/src/pool/pool.ts b/lambdas/functions/control-plane/src/pool/pool.ts index da5d2ee9b1..ab7f5a7c94 100644 --- a/lambdas/functions/control-plane/src/pool/pool.ts +++ b/lambdas/functions/control-plane/src/pool/pool.ts @@ -1,6 +1,7 @@ import { Octokit } from '@octokit/rest'; import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; import { resolveComputeProviderType } from '@aws-github-runner/compute-providers/provider-types'; +import { getRunnerConfigStore } from '@aws-github-runner/storage-providers'; import yn from 'yn'; import { @@ -31,7 +32,6 @@ export async function adjust(event: PoolEvent): Promise { const runnerGroup = process.env.RUNNER_GROUP_NAME || ''; const runnerNamePrefix = process.env.RUNNER_NAME_PREFIX || ''; const environment = process.env.ENVIRONMENT; - const ssmTokenPath = process.env.SSM_TOKEN_PATH; const ssmConfigPath = process.env.SSM_CONFIG_PATH || ''; const ephemeral = yn(process.env.ENABLE_EPHEMERAL_RUNNERS, { default: false }); const enableJitConfig = yn(process.env.ENABLE_JIT_CONFIG, { default: ephemeral }); @@ -41,6 +41,7 @@ export async function adjust(event: PoolEvent): Promise { process.env.SSM_PARAMETER_STORE_TAGS && process.env.SSM_PARAMETER_STORE_TAGS.trim() !== '' ? validateSsmParameterStoreTags(process.env.SSM_PARAMETER_STORE_TAGS) : []; + getRunnerConfigStore(); // -1 disables the maximum check, matching the scale-up lambda's semantics. Defaults to unlimited // when unset so the pool keeps its previous behavior on stacks that do not provide the variable. const maximumRunners = parseInt(process.env.RUNNERS_MAXIMUM_COUNT || '-1'); @@ -103,7 +104,6 @@ export async function adjust(event: PoolEvent): Promise { runnerNamePrefix, runnerType: 'Org', disableAutoUpdate: disableAutoUpdate, - ssmTokenPath, ssmConfigPath, ssmParameterStoreTags, }, diff --git a/lambdas/functions/control-plane/src/scale-runners/github-runner.ts b/lambdas/functions/control-plane/src/scale-runners/github-runner.ts index 745c66770a..79c78608dc 100644 --- a/lambdas/functions/control-plane/src/scale-runners/github-runner.ts +++ b/lambdas/functions/control-plane/src/scale-runners/github-runner.ts @@ -1,5 +1,10 @@ import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; import { getParameter, putParameter } from '@aws-github-runner/aws-ssm-util'; +import { + getRunnerConfigStore, + type RunnerConfigMetadataTag, + type RunnerConfigStore, +} from '@aws-github-runner/storage-providers'; import { Octokit } from '@octokit/rest'; import { getStoredInstallationId } from '../github/auth'; @@ -14,7 +19,7 @@ export interface GitHubRunnerMetadata { } export interface StartRunnerConfigOptions { - getSsmParameterTags?: (runnerId: string) => { Key: string; Value: string }[]; + getRunnerConfigMetadataTags?: (runnerId: string) => RunnerConfigMetadataTag[]; onJitConfigCreated?: (runnerId: string, metadata: GitHubRunnerMetadata) => Promise; } @@ -250,18 +255,20 @@ export async function createStartRunnerConfig( ghClient: Octokit, options: StartRunnerConfigOptions = {}, ): Promise { + const runnerConfigStore = getRunnerConfigStore(); if (githubRunnerConfig.enableJitConfig && githubRunnerConfig.ephemeral) { - return await createJitConfig(githubRunnerConfig, runnerIds, ghClient, options); + return await createJitConfig(githubRunnerConfig, runnerIds, ghClient, runnerConfigStore, options); } else { - return await createRegistrationTokenConfig(githubRunnerConfig, runnerIds, ghClient, options); + return await createRegistrationTokenConfig(githubRunnerConfig, runnerIds, ghClient, runnerConfigStore, options); } } -function addDelay(runnerIds: string[]) { +function addDelay(runnerIds: string[], runnerConfigStore: RunnerConfigStore) { const delay = async (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); - const ssmParameterStoreMaxThroughput = 40; - const isDelay = runnerIds.length >= ssmParameterStoreMaxThroughput; - return { isDelay, delay }; + const maxWritesPerSecond = runnerConfigStore.maxWritesPerSecond; + const isDelay = maxWritesPerSecond !== undefined && runnerIds.length >= maxWritesPerSecond; + const delayMilliseconds = maxWritesPerSecond === undefined ? 0 : 1000 / maxWritesPerSecond; + return { isDelay, delay, delayMilliseconds }; } /** @@ -273,9 +280,10 @@ async function createRegistrationTokenConfig( githubRunnerConfig: CreateGitHubRunnerConfig, runnerIds: string[], ghClient: Octokit, + runnerConfigStore: RunnerConfigStore, options: StartRunnerConfigOptions, ): Promise { - const { isDelay, delay } = addDelay(runnerIds); + const { isDelay, delay, delayMilliseconds } = addDelay(runnerIds, runnerConfigStore); const token = await getGithubRunnerRegistrationToken(githubRunnerConfig, ghClient); const runnerServiceConfig = generateRunnerServiceConfig(githubRunnerConfig, token); @@ -284,12 +292,13 @@ async function createRegistrationTokenConfig( }); for (const runnerId of runnerIds) { - await putParameter(`${githubRunnerConfig.ssmTokenPath}/${runnerId}`, runnerServiceConfig.join(' '), true, { - tags: [...(options.getSsmParameterTags?.(runnerId) ?? []), ...githubRunnerConfig.ssmParameterStoreTags], - }); + await runnerConfigStore.create( + { runnerId, value: runnerServiceConfig.join(' ') }, + { metadataTags: options.getRunnerConfigMetadataTags?.(runnerId) }, + ); if (isDelay) { - // Delay to prevent AWS ssm rate limits by being within the max throughput limit - await delay(25); + // Delay to stay within the selected store's maximum write throughput. + await delay(delayMilliseconds); } } @@ -306,10 +315,11 @@ async function createJitConfig( githubRunnerConfig: CreateGitHubRunnerConfig, runnerIds: string[], ghClient: Octokit, + runnerConfigStore: RunnerConfigStore, options: StartRunnerConfigOptions, ): Promise { const runnerGroupId = await getRunnerGroupId(githubRunnerConfig, ghClient); - const { isDelay, delay } = addDelay(runnerIds); + const { isDelay, delay, delayMilliseconds } = addDelay(runnerIds, runnerConfigStore); const runnerLabels = githubRunnerConfig.runnerLabels.split(','); const failedRunnerIds: string[] = []; @@ -347,16 +357,16 @@ async function createJitConfig( runnerLabels, }); - // store jit config in ssm parameter store logger.debug('Runner JIT config for ephemeral runner generated.', { instance: runnerId, }); - await putParameter(`${githubRunnerConfig.ssmTokenPath}/${runnerId}`, runnerConfig.data.encoded_jit_config, true, { - tags: [...(options.getSsmParameterTags?.(runnerId) ?? []), ...githubRunnerConfig.ssmParameterStoreTags], - }); + await runnerConfigStore.create( + { runnerId, value: runnerConfig.data.encoded_jit_config }, + { metadataTags: options.getRunnerConfigMetadataTags?.(runnerId) }, + ); if (isDelay) { - // Delay to prevent AWS ssm rate limits by being within the max throughput limit - await delay(25); + // Delay to stay within the selected store's maximum write throughput. + await delay(delayMilliseconds); } } catch (error) { failedRunnerIds.push(runnerId); diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up-contract.test.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up-contract.test.ts index 3c1a0362bb..5d190f3e5f 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up-contract.test.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up-contract.test.ts @@ -1,4 +1,5 @@ import type { Octokit } from '@octokit/rest'; +import { resetRunnerConfigStore } from '@aws-github-runner/storage-providers'; import { beforeEach, vi } from 'vitest'; import { providerTypes } from '../test/compute-provider-contracts/provider-types'; @@ -56,6 +57,8 @@ const computeProviders = providerTypes.map((type) => ({ beforeEach(() => { vi.clearAllMocks(); process.env = { ...cleanEnv }; + process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/tokens'; + resetRunnerConfigStore(); mockedAppAuth.mockResolvedValue({ type: 'app', token: 'app-token', appId: 1, expiresAt: 'some-date' }); mockedInstallationAuth.mockResolvedValue({ 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 eb6899ff79..a95941bf30 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 @@ -17,6 +17,7 @@ import type { ScaleUpComputeProvider, } from './types'; import { getParameter } from '@aws-github-runner/aws-ssm-util'; +import { resetRunnerConfigStore } from '@aws-github-runner/storage-providers'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { Octokit } from '@octokit/rest'; @@ -147,6 +148,7 @@ function setDefaults() { process.env.GITHUB_APP_CLIENT_SECRET = 'TEST_CLIENT_SECRET'; process.env.RUNNERS_MAXIMUM_COUNT = '3'; process.env.ENVIRONMENT = EXPECTED_RUNNER_PARAMS.environment; + process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; } async function createTestProviderRunners(input: CreateScaleUpRunnersInput): Promise { @@ -168,7 +170,7 @@ async function createTestProviderRunners(input: CreateScaleUpRunnersInput [{ Key: 'RunnerId', Value: runnerId }], + getRunnerConfigMetadataTags: (runnerId) => [{ key: 'RunnerId', value: runnerId }], }, ); } catch { @@ -187,6 +189,7 @@ beforeEach(() => { vi.resetModules(); vi.clearAllMocks(); setDefaults(); + resetRunnerConfigStore(); defaultSSMGetParameterMockImpl(); defaultOctokitMockImpl(); @@ -2164,6 +2167,19 @@ describe('compute provider selection', () => { }); }); +describe('runner config store preflight', () => { + it('rejects an unsupported store before resolving compute or GitHub providers', async () => { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'unsupported-provider'; + + await expect(scaleUpModule.scaleUp(TEST_DATA)).rejects.toThrow( + "Unsupported runner config storage provider 'unsupported-provider'", + ); + + expect(mockedResolveCapability).not.toHaveBeenCalled(); + expect(mockedAppAuth).not.toHaveBeenCalled(); + }); +}); + describe('Multi-app round-robin', () => { const mockedGetAppCount = vi.mocked(ghAuth.getAppCount); const mockedGetStoredInstallationId = vi.mocked(ghAuth.getStoredInstallationId); 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 44d522a1f0..a33a8c9705 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts @@ -1,5 +1,6 @@ import { addPersistentContextToChildLogger, createChildLogger } from '@aws-github-runner/aws-powertools-util'; import { resolveComputeProviderType } from '@aws-github-runner/compute-providers/provider-types'; +import { getRunnerConfigStore } from '@aws-github-runner/storage-providers'; import { Octokit } from '@octokit/rest'; import yn from 'yn'; @@ -80,7 +81,6 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise { function createEc2StartRunnerConfigOptions(): StartRunnerConfigOptions { return { - getSsmParameterTags: (instanceId) => [{ Key: 'InstanceId', Value: instanceId }], + getRunnerConfigMetadataTags: (instanceId) => [{ key: 'InstanceId', value: instanceId }], onJitConfigCreated: async (instanceId, metadata) => await tagEc2RunnerMetadata(instanceId, metadata), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts index 0c6bac69f4..7695839ba4 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts @@ -51,7 +51,6 @@ function runnerConfig(overrides: Partial = {}): Create runnerOwner, runnerType: 'Org', disableAutoUpdate: false, - ssmTokenPath: '/github-action-runners/default/runners/config', ssmConfigPath: '/github-action-runners/default/runners/config', ssmParameterStoreTags: [], ...overrides, @@ -182,7 +181,7 @@ describe('scaleUp with GHES', () => { { Key: 'ghr:runner_labels', Value: 'label1,label2' }, ]); const [, , , options] = mockCreateStartRunnerConfig.mock.calls[0]; - expect(options?.getSsmParameterTags?.('i-12345')).toEqual([{ Key: 'InstanceId', Value: 'i-12345' }]); + expect(options?.getRunnerConfigMetadataTags?.('i-12345')).toEqual([{ key: 'InstanceId', value: 'i-12345' }]); }); it('chunks comma-joined GitHub runner labels by the EC2 tag value max length', async () => { diff --git a/lambdas/libs/compute-providers/core/index.ts b/lambdas/libs/compute-providers/core/index.ts index 2b5f937f36..9125dd9b90 100644 --- a/lambdas/libs/compute-providers/core/index.ts +++ b/lambdas/libs/compute-providers/core/index.ts @@ -21,7 +21,6 @@ export interface CreateGitHubRunnerConfig { runnerOwner: string; runnerType: RunnerType; disableAutoUpdate: boolean; - ssmTokenPath: string; ssmConfigPath: string; ssmParameterStoreTags: { Key: string; Value: string }[]; } @@ -32,7 +31,7 @@ export interface GitHubRunnerMetadata { } export interface StartRunnerConfigOptions { - getSsmParameterTags?: (runnerId: string) => { Key: string; Value: string }[]; + getRunnerConfigMetadataTags?: (runnerId: string) => { key: string; value: string }[]; onJitConfigCreated?: (runnerId: string, metadata: GitHubRunnerMetadata) => Promise; } diff --git a/lambdas/libs/storage-providers/aws/ssm/environment.d.ts b/lambdas/libs/storage-providers/aws/ssm/environment.d.ts new file mode 100644 index 0000000000..c6dd725742 --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/environment.d.ts @@ -0,0 +1,10 @@ +export {}; + +declare global { + namespace NodeJS { + interface ProcessEnv { + SSM_PARAMETER_STORE_TAGS?: string; + SSM_TOKEN_PATH?: string; + } + } +} diff --git a/lambdas/libs/storage-providers/aws/ssm/parameter-store-tags.ts b/lambdas/libs/storage-providers/aws/ssm/parameter-store-tags.ts new file mode 100644 index 0000000000..d35150e10a --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/parameter-store-tags.ts @@ -0,0 +1,42 @@ +interface SsmParameterStoreTag { + Key: string; + Value: string; +} + +export function loadSsmParameterStoreTagsFromEnvironment(): SsmParameterStoreTag[] { + return process.env.SSM_PARAMETER_STORE_TAGS && process.env.SSM_PARAMETER_STORE_TAGS.trim() !== '' + ? validateSsmParameterStoreTags(process.env.SSM_PARAMETER_STORE_TAGS) + : []; +} + +function validateSsmParameterStoreTags(tagsJson: string): SsmParameterStoreTag[] { + try { + const tags: unknown = JSON.parse(tagsJson); + + if (!Array.isArray(tags)) { + throw new Error('Tags must be an array'); + } + + if (tags.length === 0) { + return []; + } + + tags.forEach((tag: unknown, index: number) => { + if (typeof tag !== 'object' || tag === null) { + throw new Error(`Tag at index ${index} must be an object`); + } + + const candidate = tag as Record; + if (!candidate.Key || typeof candidate.Key !== 'string' || candidate.Key.trim() === '') { + throw new Error(`Tag at index ${index} has missing or invalid 'Key' property`); + } + if (!Object.prototype.hasOwnProperty.call(candidate, 'Value') || typeof candidate.Value !== 'string') { + throw new Error(`Tag at index ${index} has missing or invalid 'Value' property`); + } + }); + + return tags as SsmParameterStoreTag[]; + } catch (error) { + throw new Error(`Failed to parse SSM_PARAMETER_STORE_TAGS: ${(error as Error).message}`); + } +} diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-store.test.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.test.ts new file mode 100644 index 0000000000..eaacff2718 --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.test.ts @@ -0,0 +1,88 @@ +import { putParameter } from '@aws-github-runner/aws-ssm-util'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { createAwsSsmRunnerConfigStore } from './runner-config-store'; + +vi.mock('@aws-github-runner/aws-ssm-util', () => ({ + putParameter: vi.fn(), +})); + +const putParameterMock = vi.mocked(putParameter); +const cleanEnv = process.env; + +describe('aws_ssm runner config store', () => { + beforeEach(() => { + vi.clearAllMocks(); + process.env = { ...cleanEnv }; + delete process.env.SSM_PARAMETER_STORE_TAGS; + process.env.SSM_TOKEN_PATH = '/runner/tokens'; + }); + + it('creates a secure parameter at the legacy path with metadata tags before configured tags', async () => { + process.env.SSM_PARAMETER_STORE_TAGS = JSON.stringify([ + { Key: 'Environment', Value: 'test' }, + { Key: 'Team', Value: 'actions' }, + ]); + const store = createAwsSsmRunnerConfigStore(); + + await store.create( + { runnerId: 'i-123', value: 'encoded-jit-config' }, + { metadataTags: [{ key: 'InstanceId', value: 'i-123' }] }, + ); + + expect(store.maxWritesPerSecond).toBe(40); + expect(putParameterMock).toHaveBeenCalledWith('/runner/tokens/i-123', 'encoded-jit-config', true, { + tags: [ + { Key: 'InstanceId', Value: 'i-123' }, + { Key: 'Environment', Value: 'test' }, + { Key: 'Team', Value: 'actions' }, + ], + }); + }); + + it('uses an empty tag list when no tags are configured', async () => { + const store = createAwsSsmRunnerConfigStore(); + + await store.create({ runnerId: 'runner-1', value: 'registration-config' }); + + expect(putParameterMock).toHaveBeenCalledWith('/runner/tokens/runner-1', 'registration-config', true, { + tags: [], + }); + }); + + it.each([undefined, '', ' '])('rejects missing or blank SSM_TOKEN_PATH %j before writing', (tokenPath) => { + setTokenPath(tokenPath); + + expect(() => createAwsSsmRunnerConfigStore()).toThrow('Environment variable SSM_TOKEN_PATH is not set'); + expect(putParameterMock).not.toHaveBeenCalled(); + }); + + it.each([ + ['{}', 'Tags must be an array'], + ['[null]', 'Tag at index 0 must be an object'], + [JSON.stringify([{ Key: '', Value: 'test' }]), "Tag at index 0 has missing or invalid 'Key' property"], + [JSON.stringify([{ Key: 'Environment' }]), "Tag at index 0 has missing or invalid 'Value' property"], + ])('rejects invalid legacy SSM parameter tags', (tags, reason) => { + process.env.SSM_PARAMETER_STORE_TAGS = tags; + + expect(() => createAwsSsmRunnerConfigStore()).toThrow(`Failed to parse SSM_PARAMETER_STORE_TAGS: ${reason}`); + expect(putParameterMock).not.toHaveBeenCalled(); + }); + + it('treats a blank legacy tag value as no configured tags', async () => { + process.env.SSM_PARAMETER_STORE_TAGS = ' '; + const store = createAwsSsmRunnerConfigStore(); + + await store.create({ runnerId: 'runner-1', value: 'jit-config' }); + + expect(putParameterMock).toHaveBeenCalledWith('/runner/tokens/runner-1', 'jit-config', true, { tags: [] }); + }); +}); + +function setTokenPath(tokenPath: string | undefined): void { + if (tokenPath === undefined) { + delete process.env.SSM_TOKEN_PATH; + } else { + process.env.SSM_TOKEN_PATH = tokenPath; + } +} diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts new file mode 100644 index 0000000000..179bcfa87c --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts @@ -0,0 +1,37 @@ +import { putParameter } from '@aws-github-runner/aws-ssm-util'; + +import type { RunnerConfigMetadataTag, RunnerConfigRecord, RunnerConfigStore } from '../../core'; +import type {} from './environment'; +import { loadSsmParameterStoreTagsFromEnvironment } from './parameter-store-tags'; + +interface AwsSsmRunnerConfigStoreConfig { + tokenPath: string; + parameterStoreTags: { Key: string; Value: string }[]; +} + +export function createAwsSsmRunnerConfigStore(): RunnerConfigStore { + const tokenPath = process.env.SSM_TOKEN_PATH; + if (!tokenPath || tokenPath.trim() === '') { + throw new Error('Environment variable SSM_TOKEN_PATH is not set'); + } + + return new AwsSsmRunnerConfigStore({ + tokenPath, + parameterStoreTags: loadSsmParameterStoreTagsFromEnvironment(), + }); +} + +class AwsSsmRunnerConfigStore implements RunnerConfigStore { + readonly maxWritesPerSecond = 40; + + constructor(private readonly config: AwsSsmRunnerConfigStoreConfig) {} + + async create(record: RunnerConfigRecord, options: { metadataTags?: RunnerConfigMetadataTag[] } = {}): Promise { + await putParameter(`${this.config.tokenPath}/${record.runnerId}`, record.value, true, { + tags: [ + ...(options.metadataTags ?? []).map(({ key, value }) => ({ Key: key, Value: value })), + ...this.config.parameterStoreTags, + ], + }); + } +} diff --git a/lambdas/libs/storage-providers/core/index.ts b/lambdas/libs/storage-providers/core/index.ts new file mode 100644 index 0000000000..7f9df413c0 --- /dev/null +++ b/lambdas/libs/storage-providers/core/index.ts @@ -0,0 +1,14 @@ +export interface RunnerConfigMetadataTag { + key: string; + value: string; +} + +export interface RunnerConfigRecord { + runnerId: string; + value: string; +} + +export interface RunnerConfigStore { + readonly maxWritesPerSecond?: number; + create(record: RunnerConfigRecord, options?: { metadataTags?: RunnerConfigMetadataTag[] }): Promise; +} diff --git a/lambdas/libs/storage-providers/environment.d.ts b/lambdas/libs/storage-providers/environment.d.ts new file mode 100644 index 0000000000..0f7ade9095 --- /dev/null +++ b/lambdas/libs/storage-providers/environment.d.ts @@ -0,0 +1,9 @@ +export {}; + +declare global { + namespace NodeJS { + interface ProcessEnv { + RUNNER_CONFIG_STORAGE_PROVIDER?: string; + } + } +} diff --git a/lambdas/libs/storage-providers/index.ts b/lambdas/libs/storage-providers/index.ts new file mode 100644 index 0000000000..862924118b --- /dev/null +++ b/lambdas/libs/storage-providers/index.ts @@ -0,0 +1,2 @@ +export type { RunnerConfigMetadataTag, RunnerConfigRecord, RunnerConfigStore } from './core'; +export { getRunnerConfigStore, resetRunnerConfigStore } from './runner-config'; diff --git a/lambdas/libs/storage-providers/package.json b/lambdas/libs/storage-providers/package.json new file mode 100644 index 0000000000..2b753fc436 --- /dev/null +++ b/lambdas/libs/storage-providers/package.json @@ -0,0 +1,29 @@ +{ + "name": "@aws-github-runner/storage-providers", + "version": "1.0.0", + "main": "index.ts", + "exports": { + ".": "./index.ts" + }, + "type": "module", + "license": "MIT", + "scripts": { + "test": "NODE_ENV=test nx test", + "test:watch": "NODE_ENV=test nx test --watch", + "lint": "eslint .", + "format": "prettier --write \"**/*.ts\"", + "format-check": "prettier --check \"**/*.ts\"", + "all": "yarn format && yarn lint && yarn test" + }, + "dependencies": { + "@aws-github-runner/aws-ssm-util": "*" + }, + "nx": { + "includedScripts": [ + "format", + "format-check", + "lint", + "all" + ] + } +} diff --git a/lambdas/libs/storage-providers/runner-config.test.ts b/lambdas/libs/storage-providers/runner-config.test.ts new file mode 100644 index 0000000000..10467ebdb4 --- /dev/null +++ b/lambdas/libs/storage-providers/runner-config.test.ts @@ -0,0 +1,84 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { createAwsSsmRunnerConfigStore } from './aws/ssm/runner-config-store'; +import type { RunnerConfigStore } from './core'; +import { getRunnerConfigStore, resetRunnerConfigStore } from './runner-config'; + +vi.mock('./aws/ssm/runner-config-store', () => ({ + createAwsSsmRunnerConfigStore: vi.fn(), +})); + +const createAwsSsmRunnerConfigStoreMock = vi.mocked(createAwsSsmRunnerConfigStore); +const cleanEnv = process.env; + +describe('runner config store selection', () => { + beforeEach(() => { + vi.clearAllMocks(); + process.env = { ...cleanEnv }; + delete process.env.RUNNER_CONFIG_STORAGE_PROVIDER; + resetRunnerConfigStore(); + }); + + it.each([undefined, '', ' '])('uses aws_ssm for default selector input %j', (provider) => { + setProvider(provider); + const store = stubStore(); + + expect(getRunnerConfigStore()).toBe(store); + expect(createAwsSsmRunnerConfigStoreMock).toHaveBeenCalledOnce(); + }); + + it.each(['aws_ssm', ' AWS_SSM '])('uses aws_ssm for explicit selector input %j', (provider) => { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = provider; + const store = stubStore(); + + expect(getRunnerConfigStore()).toBe(store); + expect(createAwsSsmRunnerConfigStoreMock).toHaveBeenCalledOnce(); + }); + + it('rejects an unsupported provider on first use', () => { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'not-registered'; + + expect(createAwsSsmRunnerConfigStoreMock).not.toHaveBeenCalled(); + expect(() => getRunnerConfigStore()).toThrow("Unsupported runner config storage provider 'not-registered'"); + expect(createAwsSsmRunnerConfigStoreMock).not.toHaveBeenCalled(); + }); + + it('selects lazily and caches the created store', () => { + const store = stubStore(); + + expect(createAwsSsmRunnerConfigStoreMock).not.toHaveBeenCalled(); + const first = getRunnerConfigStore(); + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'not-registered'; + const second = getRunnerConfigStore(); + + expect(first).toBe(store); + expect(second).toBe(store); + expect(createAwsSsmRunnerConfigStoreMock).toHaveBeenCalledOnce(); + }); + + it('selects again after the test reset', () => { + const firstStore = stubStore(); + expect(getRunnerConfigStore()).toBe(firstStore); + + const secondStore = { create: vi.fn() } satisfies RunnerConfigStore; + createAwsSsmRunnerConfigStoreMock.mockReturnValue(secondStore); + resetRunnerConfigStore(); + + expect(getRunnerConfigStore()).toBe(secondStore); + expect(createAwsSsmRunnerConfigStoreMock).toHaveBeenCalledTimes(2); + }); +}); + +function setProvider(provider: string | undefined): void { + if (provider === undefined) { + delete process.env.RUNNER_CONFIG_STORAGE_PROVIDER; + } else { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = provider; + } +} + +function stubStore(): RunnerConfigStore { + const store = { create: vi.fn() } satisfies RunnerConfigStore; + createAwsSsmRunnerConfigStoreMock.mockReturnValue(store); + return store; +} diff --git a/lambdas/libs/storage-providers/runner-config.ts b/lambdas/libs/storage-providers/runner-config.ts new file mode 100644 index 0000000000..dfac2fcfb1 --- /dev/null +++ b/lambdas/libs/storage-providers/runner-config.ts @@ -0,0 +1,46 @@ +import { createAwsSsmRunnerConfigStore } from './aws/ssm/runner-config-store'; +import type { RunnerConfigStore } from './core'; +import type {} from './environment'; + +type RunnerConfigStoreFactory = () => RunnerConfigStore; + +const providerFactories = { + aws_ssm: createAwsSsmRunnerConfigStore, +} as const satisfies Record; + +type RunnerConfigStorageProvider = keyof typeof providerFactories; + +const defaultProvider = 'aws_ssm' satisfies RunnerConfigStorageProvider; + +let runnerConfigStore: RunnerConfigStore | undefined; + +export function getRunnerConfigStore(): RunnerConfigStore { + runnerConfigStore ??= providerFactories[resolveProvider(process.env.RUNNER_CONFIG_STORAGE_PROVIDER)](); + return runnerConfigStore; +} + +// Test-only reset for cases that need to exercise first-use environment selection. +export function resetRunnerConfigStore(): void { + runnerConfigStore = undefined; +} + +function resolveProvider(provider: unknown): RunnerConfigStorageProvider { + if (provider === undefined) { + return defaultProvider; + } + + if (typeof provider !== 'string') { + throw new Error(`Unsupported runner config storage provider '${String(provider)}'`); + } + + const normalizedProvider = provider.trim().toLowerCase(); + if (normalizedProvider === '') { + return defaultProvider; + } + + if (!Object.prototype.hasOwnProperty.call(providerFactories, normalizedProvider)) { + throw new Error(`Unsupported runner config storage provider '${String(provider)}'`); + } + + return normalizedProvider as RunnerConfigStorageProvider; +} diff --git a/lambdas/libs/storage-providers/tsconfig.json b/lambdas/libs/storage-providers/tsconfig.json new file mode 100644 index 0000000000..139069a7cf --- /dev/null +++ b/lambdas/libs/storage-providers/tsconfig.json @@ -0,0 +1,5 @@ +{ + "extends": "../../tsconfig.json", + "include": ["**/*.ts"], + "exclude": ["**/*.test.ts", "vitest.config.ts"] +} diff --git a/lambdas/libs/storage-providers/vitest.config.ts b/lambdas/libs/storage-providers/vitest.config.ts new file mode 100644 index 0000000000..a5812ad13e --- /dev/null +++ b/lambdas/libs/storage-providers/vitest.config.ts @@ -0,0 +1,14 @@ +import { resolve } from 'path'; + +import { mergeConfig } from 'vitest/config'; +import defaultConfig from '../../vitest.base.config'; + +export default mergeConfig(defaultConfig, { + test: { + setupFiles: [resolve(__dirname, '../../aws-vitest-setup.ts')], + coverage: { + include: ['index.ts', 'runner-config.ts', 'core/**/*.ts', 'aws/**/*.ts'], + exclude: ['**/*.test.ts', '**/*.d.ts'], + }, + }, +}); diff --git a/lambdas/yarn.lock b/lambdas/yarn.lock index 56ae435c2c..84a1a7069c 100644 --- a/lambdas/yarn.lock +++ b/lambdas/yarn.lock @@ -164,6 +164,7 @@ __metadata: "@aws-github-runner/aws-powertools-util": "npm:*" "@aws-github-runner/aws-ssm-util": "npm:*" "@aws-github-runner/compute-providers": "npm:*" + "@aws-github-runner/storage-providers": "npm:*" "@aws-lambda-powertools/parameters": "npm:^2.31.0" "@aws-sdk/client-ec2": "npm:^3.1009.0" "@aws-sdk/client-sqs": "npm:^3.1009.0" @@ -209,6 +210,14 @@ __metadata: languageName: unknown linkType: soft +"@aws-github-runner/storage-providers@npm:*, @aws-github-runner/storage-providers@workspace:libs/storage-providers": + version: 0.0.0-use.local + resolution: "@aws-github-runner/storage-providers@workspace:libs/storage-providers" + dependencies: + "@aws-github-runner/aws-ssm-util": "npm:*" + languageName: unknown + linkType: soft + "@aws-github-runner/termination-watcher@workspace:functions/termination-watcher": version: 0.0.0-use.local resolution: "@aws-github-runner/termination-watcher@workspace:functions/termination-watcher" From 3eec22dfb85c5e706fb5c67debd097afb3284483 Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Tue, 18 Aug 2026 23:32:22 +0200 Subject: [PATCH 2/6] refactor(storage): extract group cache and cleanup --- .../control-plane/src/lambda.test.ts | 36 ++++-- lambdas/functions/control-plane/src/lambda.ts | 6 +- .../src/local-ssm-housekeeper.ts | 9 +- .../functions/control-plane/src/modules.d.ts | 1 - .../src/pool/pool-contract.test.ts | 6 - .../control-plane/src/pool/pool.test.ts | 15 --- .../functions/control-plane/src/pool/pool.ts | 11 +- .../src/scale-runners/github-runner.ts | 76 +++-------- .../scale-runners/scale-up-contract.test.ts | 3 - .../src/scale-runners/scale-up.test.ts | 21 +--- .../src/scale-runners/scale-up.ts | 10 -- .../src/scale-runners/ssm-housekeeper.test.ts | 118 ------------------ .../ec2/src/control-plane/runner-config.ts | 2 +- .../ec2/src/control-plane/scale-up.test.ts | 4 +- lambdas/libs/compute-providers/core/index.ts | 4 +- .../aws/ssm/environment.d.ts | 2 + .../aws/ssm/runner-config-housekeeper.test.ts | 97 ++++++++++++++ .../aws/ssm/runner-config-housekeeper.ts} | 12 +- .../aws/ssm/runner-config-store.test.ts | 11 +- .../aws/ssm/runner-config-store.ts | 35 ++++-- .../aws/ssm/runner-group-cache-store.test.ts | 72 +++++++++++ .../aws/ssm/runner-group-cache-store.ts | 41 ++++++ lambdas/libs/storage-providers/core/index.ts | 15 ++- lambdas/libs/storage-providers/index.ts | 9 +- lambdas/libs/storage-providers/package.json | 8 +- lambdas/libs/storage-providers/provider.ts | 26 ++++ .../storage-providers/runner-config.test.ts | 4 +- .../libs/storage-providers/runner-config.ts | 31 +---- .../runner-group-cache.test.ts | 83 ++++++++++++ .../storage-providers/runner-group-cache.ts | 23 ++++ .../libs/storage-providers/vitest.config.ts | 2 +- lambdas/yarn.lock | 4 + 32 files changed, 488 insertions(+), 309 deletions(-) delete mode 100644 lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.test.ts create mode 100644 lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.test.ts rename lambdas/{functions/control-plane/src/scale-runners/ssm-housekeeper.ts => libs/storage-providers/aws/ssm/runner-config-housekeeper.ts} (84%) create mode 100644 lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.test.ts create mode 100644 lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.ts create mode 100644 lambdas/libs/storage-providers/provider.ts create mode 100644 lambdas/libs/storage-providers/runner-group-cache.test.ts create mode 100644 lambdas/libs/storage-providers/runner-group-cache.ts diff --git a/lambdas/functions/control-plane/src/lambda.test.ts b/lambdas/functions/control-plane/src/lambda.test.ts index 26b130ffe1..f93b4eac49 100644 --- a/lambdas/functions/control-plane/src/lambda.test.ts +++ b/lambdas/functions/control-plane/src/lambda.test.ts @@ -1,4 +1,5 @@ import { captureLambdaHandler, logger } from '@aws-github-runner/aws-powertools-util'; +import { getRunnerConfigStore, type RunnerConfigStore } from '@aws-github-runner/storage-providers'; import { Context, SQSEvent, SQSRecord } from 'aws-lambda'; import { addMiddleware, adjustPool, scaleDownHandler, scaleUpHandler, ssmHousekeeper, jobRetryCheck } from './lambda'; @@ -6,7 +7,6 @@ import { adjust } from './pool/pool'; import { scaleDown } from './scale-runners/scale-down'; import { scaleUp } from './scale-runners/scale-up'; import type { ActionRequestMessage } from './scale-runners/types'; -import { cleanSSMTokens } from './scale-runners/ssm-housekeeper'; import { checkAndRetryJob } from './scale-runners/job-retry'; import { describe, it, expect, vi, MockedFunction, beforeEach } from 'vitest'; @@ -64,10 +64,17 @@ const context: Context = { vi.mock('./pool/pool'); vi.mock('./scale-runners/scale-down'); vi.mock('./scale-runners/scale-up'); -vi.mock('./scale-runners/ssm-housekeeper'); vi.mock('./scale-runners/job-retry'); vi.mock('@aws-github-runner/aws-powertools-util'); vi.mock('@aws-github-runner/aws-ssm-util'); +vi.mock('@aws-github-runner/storage-providers', () => ({ + getRunnerConfigStore: vi.fn(), +})); + +const runnerConfigStore = { + create: vi.fn(), + houseKeeper: vi.fn(), +} satisfies RunnerConfigStore; describe('Test scale up lambda wrapper.', () => { it('Do not handle empty record sets.', async () => { @@ -297,22 +304,31 @@ describe('Test middleware', () => { }); describe('Test ssm housekeeper lambda wrapper.', () => { - it('Invoke without errors.', async () => { - vi.mocked(cleanSSMTokens).mockResolvedValue(); + beforeEach(() => { + vi.mocked(getRunnerConfigStore).mockReturnValue(runnerConfigStore); + }); - process.env.SSM_CLEANUP_CONFIG = JSON.stringify({ - dryRun: false, - minimumDaysOld: 1, - tokenPath: '/path/to/tokens/', - }); + it('Invoke without errors.', async () => { + runnerConfigStore.houseKeeper.mockResolvedValue(); await expect(ssmHousekeeper({}, context)).resolves.not.toThrow(); + expect(getRunnerConfigStore).toHaveBeenCalledOnce(); + expect(runnerConfigStore.houseKeeper).toHaveBeenCalledOnce(); }); it('Errors not throws.', async () => { - vi.mocked(cleanSSMTokens).mockRejectedValue(new Error()); + runnerConfigStore.houseKeeper.mockRejectedValue(new Error()); await expect(ssmHousekeeper({}, context)).resolves.not.toThrow(); }); + + it('does not catch provider construction errors', async () => { + const error = new Error('Invalid provider configuration'); + vi.mocked(getRunnerConfigStore).mockImplementation(() => { + throw error; + }); + + await expect(ssmHousekeeper({}, context)).rejects.toBe(error); + }); }); describe('Test job retry check wrapper', () => { diff --git a/lambdas/functions/control-plane/src/lambda.ts b/lambdas/functions/control-plane/src/lambda.ts index d229a0350e..270980c3bc 100644 --- a/lambdas/functions/control-plane/src/lambda.ts +++ b/lambdas/functions/control-plane/src/lambda.ts @@ -1,13 +1,13 @@ import middy from '@middy/core'; import { logger, setContext } from '@aws-github-runner/aws-powertools-util'; import { captureLambdaHandler, tracer } from '@aws-github-runner/aws-powertools-util'; +import { getRunnerConfigStore } from '@aws-github-runner/storage-providers'; import { Context, type SQSBatchItemFailure, type SQSBatchResponse, SQSEvent } from 'aws-lambda'; import { PoolEvent, adjust } from './pool/pool'; import { scaleDown } from './scale-runners/scale-down'; import { scaleUp } from './scale-runners/scale-up'; import type { ActionRequestMessage, ActionRequestMessageSQS } from './scale-runners/types'; -import { SSMCleanupOptions, cleanSSMTokens } from './scale-runners/ssm-housekeeper'; import { checkAndRetryJob } from './scale-runners/job-retry'; export async function scaleUpHandler(event: SQSEvent, context: Context): Promise { @@ -121,10 +121,10 @@ addMiddleware(); export async function ssmHousekeeper(event: unknown, context: Context): Promise { setContext(context, 'lambda.ts'); logger.logEventIfEnabled(event); - const config = JSON.parse(process.env.SSM_CLEANUP_CONFIG) as SSMCleanupOptions; + const runnerConfigStore = getRunnerConfigStore(); try { - await cleanSSMTokens(config); + await runnerConfigStore.houseKeeper(); } catch (e) { logger.error(`${(e as Error).message}`, { error: e as Error }); } diff --git a/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts b/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts index ec635b13ad..08b062a193 100644 --- a/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts +++ b/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts @@ -1,11 +1,14 @@ -import { cleanSSMTokens } from './scale-runners/ssm-housekeeper'; +import { getRunnerConfigStore } from '@aws-github-runner/storage-providers'; export function run(): void { - cleanSSMTokens({ + process.env.SSM_CLEANUP_CONFIG = JSON.stringify({ dryRun: true, minimumDaysOld: 3, tokenPath: '/ghr/my-env/runners/tokens', - }) + }); + + getRunnerConfigStore() + .houseKeeper() .then() .catch((e) => { console.log(e); diff --git a/lambdas/functions/control-plane/src/modules.d.ts b/lambdas/functions/control-plane/src/modules.d.ts index 84b0d23a02..d32f8431e0 100644 --- a/lambdas/functions/control-plane/src/modules.d.ts +++ b/lambdas/functions/control-plane/src/modules.d.ts @@ -18,7 +18,6 @@ declare namespace NodeJS { RUNNER_OWNER: string; COMPUTE_PROVIDER_TYPE?: string; SCALE_DOWN_CONFIG: string; - SSM_CLEANUP_CONFIG: string; SUBNET_IDS: string; INSTANCE_TYPES: string; INSTANCE_TARGET_CAPACITY_TYPE: 'on-demand' | 'spot'; diff --git a/lambdas/functions/control-plane/src/pool/pool-contract.test.ts b/lambdas/functions/control-plane/src/pool/pool-contract.test.ts index 28e38c6a77..a949afae39 100644 --- a/lambdas/functions/control-plane/src/pool/pool-contract.test.ts +++ b/lambdas/functions/control-plane/src/pool/pool-contract.test.ts @@ -1,6 +1,5 @@ import type { Octokit } from '@octokit/rest'; import type { ComputeProviderType } from '@aws-github-runner/compute-providers/provider-types'; -import { resetRunnerConfigStore } from '@aws-github-runner/storage-providers'; import { beforeEach, vi } from 'vitest'; import { definePoolContractTests } from '../test/compute-provider-contracts/pool'; @@ -21,7 +20,6 @@ vi.mock('../github/auth', () => ({ vi.mock('../scale-runners/github-runner', () => ({ createStartRunnerConfig: vi.fn(), getGitHubEnterpriseApiUrl: vi.fn(), - validateSsmParameterStoreTags: vi.fn(), })); const mockedAppAuth = vi.mocked(ghAuth.createGithubAppAuth); @@ -49,9 +47,6 @@ const computeProviders = providerTypes.map((type) => ({ beforeEach(() => { vi.clearAllMocks(); process.env = { ...cleanEnv }; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/tokens'; - resetRunnerConfigStore(); - mockedAppAuth.mockResolvedValue({ type: 'app', token: 'app-token', appId: 1, expiresAt: 'some-date' }); mockedInstallationAuth.mockResolvedValue({ type: 'token', @@ -65,7 +60,6 @@ beforeEach(() => { }); mockedCreateClient.mockResolvedValue(githubClient); vi.mocked(githubRunner.getGitHubEnterpriseApiUrl).mockReturnValue({ ghesApiUrl: '', ghesBaseUrl: '' }); - vi.mocked(githubRunner.validateSsmParameterStoreTags).mockReturnValue([]); vi.mocked(githubClient.apps.getOrgInstallation).mockResolvedValue({ data: { id: 2 } } as never); vi.mocked(githubClient.paginate).mockResolvedValue([]); }); diff --git a/lambdas/functions/control-plane/src/pool/pool.test.ts b/lambdas/functions/control-plane/src/pool/pool.test.ts index 5b7d5bc5e9..42a7467e0f 100644 --- a/lambdas/functions/control-plane/src/pool/pool.test.ts +++ b/lambdas/functions/control-plane/src/pool/pool.test.ts @@ -4,7 +4,6 @@ import * as nock from 'nock'; import { createRunners } from '@aws-github-runner/compute-providers/aws/ec2/control-plane/runner-config'; import { listEC2Runners } from '@aws-github-runner/compute-providers/aws/ec2/control-plane/runners'; -import { resetRunnerConfigStore } from '@aws-github-runner/storage-providers'; import * as ghAuth from '../github/auth'; import { getGitHubEnterpriseApiUrl } from '../scale-runners/github-runner'; import { adjust } from './pool'; @@ -52,7 +51,6 @@ vi.mock('../scale-runners/github-runner', async () => ({ ghesApiUrl: '', ghesBaseUrl: '', }), - validateSsmParameterStoreTags: vi.fn().mockReturnValue([]), })); const mocktokit = Octokit as MockedClass; @@ -135,7 +133,6 @@ beforeEach(() => { vi.resetModules(); vi.clearAllMocks(); process.env = { ...cleanEnv }; - resetRunnerConfigStore(); process.env.GITHUB_APP_KEY_BASE64 = 'TEST_CERTIFICATE_DATA'; process.env.GITHUB_APP_ID = '1337'; process.env.GITHUB_APP_CLIENT_ID = 'TEST_CLIENT_ID'; @@ -145,7 +142,6 @@ beforeEach(() => { process.env.ENABLE_ORGANIZATION_RUNNERS = 'true'; process.env.LAUNCH_TEMPLATE_NAME = 'lt-1'; process.env.SUBNET_IDS = 'subnet-123'; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/tokens'; process.env.INSTANCE_TYPES = 'm5.large'; process.env.INSTANCE_TARGET_CAPACITY_TYPE = 'spot'; process.env.RUNNER_OWNER = ORG; @@ -255,17 +251,6 @@ describe('Test simple pool.', () => { expect(mockListRunners).not.toHaveBeenCalled(); }); - it('Rejects an unsupported runner config store before GitHub or runner lookups.', async () => { - process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'unsupported-provider'; - - await expect(adjust({ poolSize: 10, type: 'ec2' })).rejects.toThrow( - "Unsupported runner config storage provider 'unsupported-provider'", - ); - - expect(mockedAppAuth).not.toHaveBeenCalled(); - expect(mockListRunners).not.toHaveBeenCalled(); - }); - it('Should not top up if pool size is reached.', async () => { await adjust({ poolSize: 1, type: 'ec2' }); expect(createRunners).not.toHaveBeenCalled(); diff --git a/lambdas/functions/control-plane/src/pool/pool.ts b/lambdas/functions/control-plane/src/pool/pool.ts index ab7f5a7c94..21e91adebc 100644 --- a/lambdas/functions/control-plane/src/pool/pool.ts +++ b/lambdas/functions/control-plane/src/pool/pool.ts @@ -1,7 +1,6 @@ import { Octokit } from '@octokit/rest'; import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; import { resolveComputeProviderType } from '@aws-github-runner/compute-providers/provider-types'; -import { getRunnerConfigStore } from '@aws-github-runner/storage-providers'; import yn from 'yn'; import { @@ -11,7 +10,7 @@ import { getStoredInstallationId, } from '../github/auth'; import { controlPlaneProviderRegistry } from '../control-plane-providers'; -import { getGitHubEnterpriseApiUrl, validateSsmParameterStoreTags } from '../scale-runners/github-runner'; +import { getGitHubEnterpriseApiUrl } from '../scale-runners/github-runner'; import type { RunnerStatus } from './pool-provider'; const logger = createChildLogger('pool'); @@ -32,16 +31,10 @@ export async function adjust(event: PoolEvent): Promise { const runnerGroup = process.env.RUNNER_GROUP_NAME || ''; const runnerNamePrefix = process.env.RUNNER_NAME_PREFIX || ''; const environment = process.env.ENVIRONMENT; - const ssmConfigPath = process.env.SSM_CONFIG_PATH || ''; const ephemeral = yn(process.env.ENABLE_EPHEMERAL_RUNNERS, { default: false }); const enableJitConfig = yn(process.env.ENABLE_JIT_CONFIG, { default: ephemeral }); const disableAutoUpdate = yn(process.env.DISABLE_RUNNER_AUTOUPDATE, { default: false }); const runnerOwner = process.env.RUNNER_OWNER; - const ssmParameterStoreTags: { Key: string; Value: string }[] = - process.env.SSM_PARAMETER_STORE_TAGS && process.env.SSM_PARAMETER_STORE_TAGS.trim() !== '' - ? validateSsmParameterStoreTags(process.env.SSM_PARAMETER_STORE_TAGS) - : []; - getRunnerConfigStore(); // -1 disables the maximum check, matching the scale-up lambda's semantics. Defaults to unlimited // when unset so the pool keeps its previous behavior on stacks that do not provide the variable. const maximumRunners = parseInt(process.env.RUNNERS_MAXIMUM_COUNT || '-1'); @@ -104,8 +97,6 @@ export async function adjust(event: PoolEvent): Promise { runnerNamePrefix, runnerType: 'Org', disableAutoUpdate: disableAutoUpdate, - ssmConfigPath, - ssmParameterStoreTags, }, numberOfRunners: topUp, githubInstallationClient, diff --git a/lambdas/functions/control-plane/src/scale-runners/github-runner.ts b/lambdas/functions/control-plane/src/scale-runners/github-runner.ts index 79c78608dc..ec78a2a2b5 100644 --- a/lambdas/functions/control-plane/src/scale-runners/github-runner.ts +++ b/lambdas/functions/control-plane/src/scale-runners/github-runner.ts @@ -1,8 +1,8 @@ import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; -import { getParameter, putParameter } from '@aws-github-runner/aws-ssm-util'; import { + getRunnerGroupCacheStore, getRunnerConfigStore, - type RunnerConfigMetadataTag, + type RunnerConfigMetadata, type RunnerConfigStore, } from '@aws-github-runner/storage-providers'; import { Octokit } from '@octokit/rest'; @@ -19,7 +19,7 @@ export interface GitHubRunnerMetadata { } export interface StartRunnerConfigOptions { - getRunnerConfigMetadataTags?: (runnerId: string) => RunnerConfigMetadataTag[]; + getRunnerConfigMetadata?: (runnerId: string) => RunnerConfigMetadata[]; onJitConfigCreated?: (runnerId: string, metadata: GitHubRunnerMetadata) => Promise; } @@ -56,37 +56,6 @@ function quoteShellArg(value: string): string { return `'${value.replace(/'/g, `'\\''`)}'`; } -export function validateSsmParameterStoreTags(tagsJson: string): { Key: string; Value: string }[] { - try { - const tags = JSON.parse(tagsJson); - - if (!Array.isArray(tags)) { - throw new Error('Tags must be an array'); - } - - if (tags.length === 0) { - return []; - } - - tags.forEach((tag, index) => { - if (typeof tag !== 'object' || tag === null) { - throw new Error(`Tag at index ${index} must be an object`); - } - if (!tag.Key || typeof tag.Key !== 'string' || tag.Key.trim() === '') { - throw new Error(`Tag at index ${index} has missing or invalid 'Key' property`); - } - if (!Object.prototype.hasOwnProperty.call(tag, 'Value') || typeof tag.Value !== 'string') { - throw new Error(`Tag at index ${index} has missing or invalid 'Value' property`); - } - }); - - return tags; - } catch (err) { - logger.error('Invalid SSM_PARAMETER_STORE_TAGS format', { error: err }); - throw new Error(`Failed to parse SSM_PARAMETER_STORE_TAGS: ${(err as Error).message}`); - } -} - async function getGithubRunnerRegistrationToken(githubRunnerConfig: CreateGitHubRunnerConfig, ghClient: Octokit) { const registrationToken = githubRunnerConfig.runnerType === 'Org' @@ -191,39 +160,30 @@ export async function getRunnerGroupId( // if the runnerType is Repo, then runnerGroupId is default to 1 let runnerGroupId: number | undefined = 1; if (githubRunnerConfig.runnerType === 'Org' && githubRunnerConfig.runnerGroup !== undefined) { - let runnerGroup: string | undefined; - // check if runner group id is already stored in SSM Parameter Store and - // use it if it exists to avoid API call to GitHub + const runnerGroupCacheStore = getRunnerGroupCacheStore(); + let cachedRunnerGroupId: number | undefined; + // Use a cached runner group id when available to avoid an API call to GitHub. try { - runnerGroup = await getParameter( - `${githubRunnerConfig.ssmConfigPath}/runner-group/${githubRunnerConfig.runnerGroup}`, - ); + cachedRunnerGroupId = await runnerGroupCacheStore.get(githubRunnerConfig.runnerGroup); } catch (err) { logger.debug('Handling error:', err as Error); - logger.warn( - `SSM Parameter "${githubRunnerConfig.ssmConfigPath}/runner-group/${githubRunnerConfig.runnerGroup}" - for Runner group ${githubRunnerConfig.runnerGroup} does not exist`, - ); + logger.warn(`Cached id for runner group ${githubRunnerConfig.runnerGroup} does not exist`); } - if (runnerGroup === undefined) { + if (cachedRunnerGroupId === undefined) { // get runner group id from GitHub runnerGroupId = await getRunnerGroupByName(ghClient, githubRunnerConfig); - // store runner group id in SSM + // cache the runner group id try { - await putParameter( - `${githubRunnerConfig.ssmConfigPath}/runner-group/${githubRunnerConfig.runnerGroup}`, - runnerGroupId.toString(), - false, - { - tags: githubRunnerConfig.ssmParameterStoreTags, - }, - ); + await runnerGroupCacheStore.create({ + runnerGroupName: githubRunnerConfig.runnerGroup, + runnerGroupId, + }); } catch (err) { - logger.debug('Error storing runner group id in SSM Parameter Store', err as Error); + logger.debug('Error storing runner group id in cache', err as Error); throw err; } } else { - runnerGroupId = parseInt(runnerGroup); + runnerGroupId = cachedRunnerGroupId; } } return runnerGroupId; @@ -294,7 +254,7 @@ async function createRegistrationTokenConfig( for (const runnerId of runnerIds) { await runnerConfigStore.create( { runnerId, value: runnerServiceConfig.join(' ') }, - { metadataTags: options.getRunnerConfigMetadataTags?.(runnerId) }, + { metadata: options.getRunnerConfigMetadata?.(runnerId) }, ); if (isDelay) { // Delay to stay within the selected store's maximum write throughput. @@ -362,7 +322,7 @@ async function createJitConfig( }); await runnerConfigStore.create( { runnerId, value: runnerConfig.data.encoded_jit_config }, - { metadataTags: options.getRunnerConfigMetadataTags?.(runnerId) }, + { metadata: options.getRunnerConfigMetadata?.(runnerId) }, ); if (isDelay) { // Delay to stay within the selected store's maximum write throughput. diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up-contract.test.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up-contract.test.ts index 5d190f3e5f..3c1a0362bb 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up-contract.test.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up-contract.test.ts @@ -1,5 +1,4 @@ import type { Octokit } from '@octokit/rest'; -import { resetRunnerConfigStore } from '@aws-github-runner/storage-providers'; import { beforeEach, vi } from 'vitest'; import { providerTypes } from '../test/compute-provider-contracts/provider-types'; @@ -57,8 +56,6 @@ const computeProviders = providerTypes.map((type) => ({ beforeEach(() => { vi.clearAllMocks(); process.env = { ...cleanEnv }; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/tokens'; - resetRunnerConfigStore(); mockedAppAuth.mockResolvedValue({ type: 'app', token: 'app-token', appId: 1, expiresAt: 'some-date' }); mockedInstallationAuth.mockResolvedValue({ 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 a95941bf30..9ae7e6a5da 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 @@ -17,7 +17,7 @@ import type { ScaleUpComputeProvider, } from './types'; import { getParameter } from '@aws-github-runner/aws-ssm-util'; -import { resetRunnerConfigStore } from '@aws-github-runner/storage-providers'; +import { resetRunnerConfigStore, resetRunnerGroupCacheStore } from '@aws-github-runner/storage-providers'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { Octokit } from '@octokit/rest'; @@ -148,6 +148,7 @@ function setDefaults() { process.env.GITHUB_APP_CLIENT_SECRET = 'TEST_CLIENT_SECRET'; process.env.RUNNERS_MAXIMUM_COUNT = '3'; process.env.ENVIRONMENT = EXPECTED_RUNNER_PARAMS.environment; + process.env.SSM_CONFIG_PATH = '/github-action-runners/default/runners/config'; process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; } @@ -170,7 +171,7 @@ async function createTestProviderRunners(input: CreateScaleUpRunnersInput [{ key: 'RunnerId', value: runnerId }], + getRunnerConfigMetadata: (runnerId) => [{ key: 'RunnerId', value: runnerId }], }, ); } catch { @@ -190,6 +191,7 @@ beforeEach(() => { vi.clearAllMocks(); setDefaults(); resetRunnerConfigStore(); + resetRunnerGroupCacheStore(); defaultSSMGetParameterMockImpl(); defaultOctokitMockImpl(); @@ -1198,6 +1200,7 @@ describe('scaleUp with public GH', () => { it('creates a ephemeral runner with JIT config.', async () => { process.env.ENABLE_EPHEMERAL_RUNNERS = 'true'; process.env.ENABLE_JOB_QUEUED_CHECK = 'false'; + delete process.env.SSM_CONFIG_PATH; process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; await scaleUpModule.scaleUp(TEST_DATA); expect(mockOctokit.actions.getJobForWorkflowRun).not.toBeCalled(); @@ -1243,6 +1246,7 @@ describe('scaleUp with public GH', () => { process.env.ENABLE_JIT_CONFIG = 'true'; process.env.ENABLE_JOB_QUEUED_CHECK = 'false'; process.env.RUNNER_LABELS = 'jit'; + delete process.env.SSM_CONFIG_PATH; process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; await scaleUpModule.scaleUp(TEST_DATA); expect(mockOctokit.actions.getJobForWorkflowRun).not.toBeCalled(); @@ -2167,19 +2171,6 @@ describe('compute provider selection', () => { }); }); -describe('runner config store preflight', () => { - it('rejects an unsupported store before resolving compute or GitHub providers', async () => { - process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'unsupported-provider'; - - await expect(scaleUpModule.scaleUp(TEST_DATA)).rejects.toThrow( - "Unsupported runner config storage provider 'unsupported-provider'", - ); - - expect(mockedResolveCapability).not.toHaveBeenCalled(); - expect(mockedAppAuth).not.toHaveBeenCalled(); - }); -}); - describe('Multi-app round-robin', () => { const mockedGetAppCount = vi.mocked(ghAuth.getAppCount); const mockedGetStoredInstallationId = vi.mocked(ghAuth.getStoredInstallationId); 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 a33a8c9705..48733c13c9 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts @@ -1,6 +1,5 @@ import { addPersistentContextToChildLogger, createChildLogger } from '@aws-github-runner/aws-powertools-util'; import { resolveComputeProviderType } from '@aws-github-runner/compute-providers/provider-types'; -import { getRunnerConfigStore } from '@aws-github-runner/storage-providers'; import { Octokit } from '@octokit/rest'; import yn from 'yn'; @@ -12,7 +11,6 @@ import { resolveInstallationId, isJobQueued, UnsupportedEventError, - validateSsmParameterStoreTags, } from './github-runner'; import { publishRetryMessage } from './job-retry'; import type { @@ -86,12 +84,6 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise { - beforeEach(() => { - mockSSMClient.reset(); - mockSSMClient.on(GetParametersByPathCommand).resolves({ - Parameters: undefined, - }); - mockSSMClient.on(GetParametersByPathCommand, { Path: tokenPath }).resolves({ - Parameters: [ - { - Name: tokenPath + 'i-old-01', - LastModifiedDate: dateOld, - }, - ], - NextToken: 'next', - }); - mockSSMClient.on(GetParametersByPathCommand, { Path: tokenPath, NextToken: 'next' }).resolves({ - Parameters: [ - { - Name: tokenPath + 'i-new-01', - LastModifiedDate: now, - }, - ], - NextToken: undefined, - }); - }); - - it('should delete parameters older then minimumDaysOld', async () => { - await cleanSSMTokens({ - dryRun: false, - minimumDaysOld: deleteAmisOlderThenDays, - tokenPath: tokenPath, - }); - - expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersByPathCommand, { Path: tokenPath }); - expect(mockSSMClient).toHaveReceivedCommandWith(DeleteParameterCommand, { Name: tokenPath + 'i-old-01' }); - expect(mockSSMClient).not.toHaveReceivedCommandWith(DeleteParameterCommand, { Name: tokenPath + 'i-new-01' }); - }); - - it('should not delete when dry run is activated', async () => { - await cleanSSMTokens({ - dryRun: true, - minimumDaysOld: deleteAmisOlderThenDays, - tokenPath: tokenPath, - }); - - expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersByPathCommand, { Path: tokenPath }); - expect(mockSSMClient).not.toHaveReceivedCommandWith(DeleteParameterCommand, { Name: tokenPath + 'i-old-01' }); - expect(mockSSMClient).not.toHaveReceivedCommandWith(DeleteParameterCommand, { Name: tokenPath + 'i-new-01' }); - }); - - it('should not call delete when no parameters are found.', async () => { - await expect( - cleanSSMTokens({ - dryRun: false, - minimumDaysOld: deleteAmisOlderThenDays, - tokenPath: 'no-exist', - }), - ).resolves.not.toThrow(); - - expect(mockSSMClient).not.toHaveReceivedCommandWith(DeleteParameterCommand, { Name: tokenPath + 'i-old-01' }); - expect(mockSSMClient).not.toHaveReceivedCommandWith(DeleteParameterCommand, { Name: tokenPath + 'i-new-01' }); - }); - - it('should not error on delete failure.', async () => { - mockSSMClient.on(DeleteParameterCommand).rejects(new Error('ParameterNotFound')); - - await expect( - cleanSSMTokens({ - dryRun: false, - minimumDaysOld: deleteAmisOlderThenDays, - tokenPath: tokenPath, - }), - ).resolves.not.toThrow(); - }); - - it('should only accept valid options.', async () => { - await expect( - cleanSSMTokens({ - dryRun: false, - minimumDaysOld: undefined as unknown as number, - tokenPath: tokenPath, - }), - ).rejects.toBeInstanceOf(Error); - - await expect( - cleanSSMTokens({ - dryRun: false, - minimumDaysOld: 0, - tokenPath: tokenPath, - }), - ).rejects.toBeInstanceOf(Error); - - await expect( - cleanSSMTokens({ - dryRun: false, - minimumDaysOld: 1, - tokenPath: undefined as unknown as string, - }), - ).rejects.toBeInstanceOf(Error); - }); -}); diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts index 6e6946c7e7..6d195d958b 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts @@ -141,7 +141,7 @@ async function terminateFailedInstances(instanceIds: string[]): Promise { function createEc2StartRunnerConfigOptions(): StartRunnerConfigOptions { return { - getRunnerConfigMetadataTags: (instanceId) => [{ key: 'InstanceId', value: instanceId }], + getRunnerConfigMetadata: (instanceId) => [{ key: 'InstanceId', value: instanceId }], onJitConfigCreated: async (instanceId, metadata) => await tagEc2RunnerMetadata(instanceId, metadata), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts index 7695839ba4..9018e08ff9 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts @@ -51,8 +51,6 @@ function runnerConfig(overrides: Partial = {}): Create runnerOwner, runnerType: 'Org', disableAutoUpdate: false, - ssmConfigPath: '/github-action-runners/default/runners/config', - ssmParameterStoreTags: [], ...overrides, }; } @@ -181,7 +179,7 @@ describe('scaleUp with GHES', () => { { Key: 'ghr:runner_labels', Value: 'label1,label2' }, ]); const [, , , options] = mockCreateStartRunnerConfig.mock.calls[0]; - expect(options?.getRunnerConfigMetadataTags?.('i-12345')).toEqual([{ key: 'InstanceId', value: 'i-12345' }]); + expect(options?.getRunnerConfigMetadata?.('i-12345')).toEqual([{ key: 'InstanceId', value: 'i-12345' }]); }); it('chunks comma-joined GitHub runner labels by the EC2 tag value max length', async () => { diff --git a/lambdas/libs/compute-providers/core/index.ts b/lambdas/libs/compute-providers/core/index.ts index 9125dd9b90..908b67fb54 100644 --- a/lambdas/libs/compute-providers/core/index.ts +++ b/lambdas/libs/compute-providers/core/index.ts @@ -21,8 +21,6 @@ export interface CreateGitHubRunnerConfig { runnerOwner: string; runnerType: RunnerType; disableAutoUpdate: boolean; - ssmConfigPath: string; - ssmParameterStoreTags: { Key: string; Value: string }[]; } export interface GitHubRunnerMetadata { @@ -31,7 +29,7 @@ export interface GitHubRunnerMetadata { } export interface StartRunnerConfigOptions { - getRunnerConfigMetadataTags?: (runnerId: string) => { key: string; value: string }[]; + getRunnerConfigMetadata?: (runnerId: string) => { key: string; value: string }[]; onJitConfigCreated?: (runnerId: string, metadata: GitHubRunnerMetadata) => Promise; } diff --git a/lambdas/libs/storage-providers/aws/ssm/environment.d.ts b/lambdas/libs/storage-providers/aws/ssm/environment.d.ts index c6dd725742..b3fb63cdbe 100644 --- a/lambdas/libs/storage-providers/aws/ssm/environment.d.ts +++ b/lambdas/libs/storage-providers/aws/ssm/environment.d.ts @@ -3,6 +3,8 @@ export {}; declare global { namespace NodeJS { interface ProcessEnv { + SSM_CONFIG_PATH?: string; + SSM_CLEANUP_CONFIG?: string; SSM_PARAMETER_STORE_TAGS?: string; SSM_TOKEN_PATH?: string; } diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.test.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.test.ts new file mode 100644 index 0000000000..9c837c7fb7 --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.test.ts @@ -0,0 +1,97 @@ +import { DeleteParameterCommand, GetParametersByPathCommand, SSMClient } from '@aws-sdk/client-ssm'; +import { mockClient } from 'aws-sdk-client-mock'; +import 'aws-sdk-client-mock-jest/vitest'; +import { beforeEach, describe, expect, it } from 'vitest'; + +import { createAwsSsmRunnerConfigStore } from './runner-config-store'; + +const mockSSMClient = mockClient(SSMClient); +const cleanEnv = process.env; +const minimumDaysOld = 1; +const now = new Date(); +const oldDate = new Date(); +oldDate.setDate(oldDate.getDate() - minimumDaysOld - 1); +const tokenPath = '/path/to/tokens/'; + +describe('aws_ssm runner config housekeeper', () => { + beforeEach(() => { + mockSSMClient.reset(); + process.env = { ...cleanEnv }; + delete process.env.SSM_TOKEN_PATH; + process.env.AWS_REGION = 'eu-east-1'; + setCleanupOptions({ dryRun: false, minimumDaysOld, tokenPath }); + + mockSSMClient.on(GetParametersByPathCommand).resolves({ + Parameters: undefined, + }); + mockSSMClient.on(GetParametersByPathCommand, { Path: tokenPath }).resolves({ + Parameters: [ + { + Name: `${tokenPath}i-old-01`, + LastModifiedDate: oldDate, + }, + ], + NextToken: 'next', + }); + mockSSMClient.on(GetParametersByPathCommand, { Path: tokenPath, NextToken: 'next' }).resolves({ + Parameters: [ + { + Name: `${tokenPath}i-new-01`, + LastModifiedDate: now, + }, + ], + NextToken: undefined, + }); + }); + + it('constructs without writer configuration and deletes expired records across pages', async () => { + const store = createAwsSsmRunnerConfigStore(); + + await store.houseKeeper(); + + expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersByPathCommand, { Path: tokenPath }); + expect(mockSSMClient).toHaveReceivedCommandWith(DeleteParameterCommand, { Name: `${tokenPath}i-old-01` }); + expect(mockSSMClient).not.toHaveReceivedCommandWith(DeleteParameterCommand, { Name: `${tokenPath}i-new-01` }); + }); + + it('does not delete records during a dry run', async () => { + setCleanupOptions({ dryRun: true, minimumDaysOld, tokenPath }); + const store = createAwsSsmRunnerConfigStore(); + + await store.houseKeeper(); + + expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersByPathCommand, { Path: tokenPath }); + expect(mockSSMClient).not.toHaveReceivedCommand(DeleteParameterCommand); + }); + + it('does not delete when no records are found', async () => { + setCleanupOptions({ dryRun: false, minimumDaysOld, tokenPath: 'does-not-exist' }); + const store = createAwsSsmRunnerConfigStore(); + + await expect(store.houseKeeper()).resolves.not.toThrow(); + + expect(mockSSMClient).not.toHaveReceivedCommand(DeleteParameterCommand); + }); + + it('continues when deleting an expired record fails', async () => { + mockSSMClient.on(DeleteParameterCommand).rejects(new Error('ParameterNotFound')); + const store = createAwsSsmRunnerConfigStore(); + + await expect(store.houseKeeper()).resolves.not.toThrow(); + }); + + it.each([ + { dryRun: false, minimumDaysOld: undefined as unknown as number, tokenPath }, + { dryRun: false, minimumDaysOld: 0, tokenPath }, + { dryRun: false, minimumDaysOld, tokenPath: undefined as unknown as string }, + ])('rejects invalid cleanup options %#', async (options) => { + setCleanupOptions(options); + const store = createAwsSsmRunnerConfigStore(); + + await expect(store.houseKeeper()).rejects.toBeInstanceOf(Error); + }); +}); + +function setCleanupOptions(options: { dryRun: boolean; minimumDaysOld: number; tokenPath: string }): void { + process.env.SSM_CLEANUP_CONFIG = JSON.stringify(options); +} diff --git a/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts similarity index 84% rename from lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts rename to lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts index 857b974a9d..2c0c22359c 100644 --- a/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts @@ -1,14 +1,13 @@ import { DeleteParameterCommand, GetParametersByPathCommand, SSMClient } from '@aws-sdk/client-ssm'; -import { logger } from '@aws-github-runner/aws-powertools-util'; -import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util'; +import { getTracedAWSV3Client, logger } from '@aws-github-runner/aws-powertools-util'; -export interface SSMCleanupOptions { +export interface SsmRunnerConfigCleanupOptions { dryRun: boolean; minimumDaysOld: number; tokenPath: string; } -function validateOptions(options: SSMCleanupOptions): void { +function validateOptions(options: SsmRunnerConfigCleanupOptions): void { const errorMessages: string[] = []; if (!options.minimumDaysOld || options.minimumDaysOld < 1) { errorMessages.push(`minimumDaysOld must be greater then 0, value is set to "${options.minimumDaysOld}"`); @@ -21,7 +20,7 @@ function validateOptions(options: SSMCleanupOptions): void { } } -export async function cleanSSMTokens(options: SSMCleanupOptions): Promise { +export async function cleanSsmRunnerConfigs(options: SsmRunnerConfigCleanupOptions): Promise { logger.info(`Cleaning tokens / JIT config older then ${options.minimumDaysOld} days, dryRun: ${options.dryRun}`); logger.debug('Cleaning with options', { options }); validateOptions(options); @@ -36,7 +35,6 @@ export async function cleanSSMTokens(options: SSMCleanupOptions): Promise parameters.NextToken = nextParameters.NextToken; } logger.info(`Found #${parameters.Parameters?.length} parameters in path ${options.tokenPath}`); - logger.debug('Found parameters', { parameters }); // minimumDate = today - minimumDaysOld const minimumDate = new Date(); @@ -47,7 +45,7 @@ export async function cleanSSMTokens(options: SSMCleanupOptions): Promise logger.info(`Deleting parameter ${parameter.Name} with last modified date ${parameter.LastModifiedDate}`); try { if (!options.dryRun) { - // sleep 50ms to avoid rait limit + // sleep 50ms to avoid rate limit await new Promise((resolve) => setTimeout(resolve, 50)); await client.send(new DeleteParameterCommand({ Name: parameter.Name })); } diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-store.test.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.test.ts index eaacff2718..04ecdda05d 100644 --- a/lambdas/libs/storage-providers/aws/ssm/runner-config-store.test.ts +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.test.ts @@ -14,11 +14,12 @@ describe('aws_ssm runner config store', () => { beforeEach(() => { vi.clearAllMocks(); process.env = { ...cleanEnv }; + delete process.env.SSM_CLEANUP_CONFIG; delete process.env.SSM_PARAMETER_STORE_TAGS; process.env.SSM_TOKEN_PATH = '/runner/tokens'; }); - it('creates a secure parameter at the legacy path with metadata tags before configured tags', async () => { + it('maps metadata to tags before configured SSM tags', async () => { process.env.SSM_PARAMETER_STORE_TAGS = JSON.stringify([ { Key: 'Environment', Value: 'test' }, { Key: 'Team', Value: 'actions' }, @@ -27,7 +28,7 @@ describe('aws_ssm runner config store', () => { await store.create( { runnerId: 'i-123', value: 'encoded-jit-config' }, - { metadataTags: [{ key: 'InstanceId', value: 'i-123' }] }, + { metadata: [{ key: 'InstanceId', value: 'i-123' }] }, ); expect(store.maxWritesPerSecond).toBe(40); @@ -77,6 +78,12 @@ describe('aws_ssm runner config store', () => { expect(putParameterMock).toHaveBeenCalledWith('/runner/tokens/runner-1', 'jit-config', true, { tags: [] }); }); + + it.each(['', '{invalid-json'])('parses cleanup configuration %j during provider construction', (config) => { + process.env.SSM_CLEANUP_CONFIG = config; + + expect(() => createAwsSsmRunnerConfigStore()).toThrow(); + }); }); function setTokenPath(tokenPath: string | undefined): void { diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts index 179bcfa87c..dec2241086 100644 --- a/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts @@ -1,23 +1,32 @@ import { putParameter } from '@aws-github-runner/aws-ssm-util'; -import type { RunnerConfigMetadataTag, RunnerConfigRecord, RunnerConfigStore } from '../../core'; +import type { RunnerConfigMetadata, RunnerConfigRecord, RunnerConfigStore } from '../../core'; import type {} from './environment'; import { loadSsmParameterStoreTagsFromEnvironment } from './parameter-store-tags'; +import { cleanSsmRunnerConfigs, type SsmRunnerConfigCleanupOptions } from './runner-config-housekeeper'; interface AwsSsmRunnerConfigStoreConfig { - tokenPath: string; + tokenPath?: string; parameterStoreTags: { Key: string; Value: string }[]; + cleanupOptions?: SsmRunnerConfigCleanupOptions; } export function createAwsSsmRunnerConfigStore(): RunnerConfigStore { const tokenPath = process.env.SSM_TOKEN_PATH; - if (!tokenPath || tokenPath.trim() === '') { + const cleanupOptions = + process.env.SSM_CLEANUP_CONFIG !== undefined + ? (JSON.parse(process.env.SSM_CLEANUP_CONFIG) as SsmRunnerConfigCleanupOptions) + : undefined; + const hasWriterConfig = tokenPath !== undefined && tokenPath.trim() !== ''; + + if (!hasWriterConfig && cleanupOptions === undefined) { throw new Error('Environment variable SSM_TOKEN_PATH is not set'); } return new AwsSsmRunnerConfigStore({ - tokenPath, - parameterStoreTags: loadSsmParameterStoreTagsFromEnvironment(), + tokenPath: hasWriterConfig ? tokenPath : undefined, + parameterStoreTags: hasWriterConfig ? loadSsmParameterStoreTagsFromEnvironment() : [], + cleanupOptions, }); } @@ -26,12 +35,24 @@ class AwsSsmRunnerConfigStore implements RunnerConfigStore { constructor(private readonly config: AwsSsmRunnerConfigStoreConfig) {} - async create(record: RunnerConfigRecord, options: { metadataTags?: RunnerConfigMetadataTag[] } = {}): Promise { + async create(record: RunnerConfigRecord, options: { metadata?: RunnerConfigMetadata[] } = {}): Promise { + if (!this.config.tokenPath) { + throw new Error('Environment variable SSM_TOKEN_PATH is not set'); + } + await putParameter(`${this.config.tokenPath}/${record.runnerId}`, record.value, true, { tags: [ - ...(options.metadataTags ?? []).map(({ key, value }) => ({ Key: key, Value: value })), + ...(options.metadata ?? []).map(({ key, value }) => ({ Key: key, Value: value })), ...this.config.parameterStoreTags, ], }); } + + async houseKeeper(): Promise { + if (!this.config.cleanupOptions) { + throw new Error('Environment variable SSM_CLEANUP_CONFIG is not set'); + } + + await cleanSsmRunnerConfigs(this.config.cleanupOptions); + } } diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.test.ts b/lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.test.ts new file mode 100644 index 0000000000..3943d5a401 --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.test.ts @@ -0,0 +1,72 @@ +import { getParameter, putParameter } from '@aws-github-runner/aws-ssm-util'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { createAwsSsmRunnerGroupCacheStore } from './runner-group-cache-store'; + +vi.mock('@aws-github-runner/aws-ssm-util', () => ({ + getParameter: vi.fn(), + putParameter: vi.fn(), +})); + +const getParameterMock = vi.mocked(getParameter); +const putParameterMock = vi.mocked(putParameter); +const cleanEnv = process.env; + +describe('aws_ssm runner group cache store', () => { + beforeEach(() => { + vi.clearAllMocks(); + process.env = { ...cleanEnv }; + delete process.env.SSM_PARAMETER_STORE_TAGS; + process.env.SSM_CONFIG_PATH = '/runner/config'; + }); + + it('gets and parses a runner group id from the legacy path', async () => { + getParameterMock.mockResolvedValue('42'); + const store = createAwsSsmRunnerGroupCacheStore(); + + await expect(store.get('Default')).resolves.toBe(42); + expect(getParameterMock).toHaveBeenCalledWith('/runner/config/runner-group/Default'); + }); + + it('preserves the previous parseInt behavior for cached values', async () => { + getParameterMock.mockResolvedValue('42cached'); + const store = createAwsSsmRunnerGroupCacheStore(); + + await expect(store.get('Default')).resolves.toBe(42); + }); + + it('propagates cache read errors', async () => { + const error = new Error('not found'); + getParameterMock.mockRejectedValue(error); + const store = createAwsSsmRunnerGroupCacheStore(); + + await expect(store.get('Default')).rejects.toBe(error); + }); + + it('creates a plaintext parameter at the legacy path with configured tags', async () => { + process.env.SSM_PARAMETER_STORE_TAGS = JSON.stringify([{ Key: 'Environment', Value: 'test' }]); + const store = createAwsSsmRunnerGroupCacheStore(); + + await store.create({ runnerGroupName: 'Default', runnerGroupId: 42 }); + + expect(putParameterMock).toHaveBeenCalledWith('/runner/config/runner-group/Default', '42', false, { + tags: [{ Key: 'Environment', Value: 'test' }], + }); + }); + + it.each([undefined, '', ' '])('rejects missing or blank SSM_CONFIG_PATH %j', (configPath) => { + setConfigPath(configPath); + + expect(() => createAwsSsmRunnerGroupCacheStore()).toThrow('Environment variable SSM_CONFIG_PATH is not set'); + expect(getParameterMock).not.toHaveBeenCalled(); + expect(putParameterMock).not.toHaveBeenCalled(); + }); +}); + +function setConfigPath(configPath: string | undefined): void { + if (configPath === undefined) { + delete process.env.SSM_CONFIG_PATH; + } else { + process.env.SSM_CONFIG_PATH = configPath; + } +} diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.ts b/lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.ts new file mode 100644 index 0000000000..f79ee1b9ed --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.ts @@ -0,0 +1,41 @@ +import { getParameter, putParameter } from '@aws-github-runner/aws-ssm-util'; + +import type { RunnerGroupCacheRecord, RunnerGroupCacheStore } from '../../core'; +import type {} from './environment'; +import { loadSsmParameterStoreTagsFromEnvironment } from './parameter-store-tags'; + +interface AwsSsmRunnerGroupCacheStoreConfig { + configPath: string; + parameterStoreTags: { Key: string; Value: string }[]; +} + +export function createAwsSsmRunnerGroupCacheStore(): RunnerGroupCacheStore { + const configPath = process.env.SSM_CONFIG_PATH; + if (!configPath || configPath.trim() === '') { + throw new Error('Environment variable SSM_CONFIG_PATH is not set'); + } + + return new AwsSsmRunnerGroupCacheStore({ + configPath, + parameterStoreTags: loadSsmParameterStoreTagsFromEnvironment(), + }); +} + +class AwsSsmRunnerGroupCacheStore implements RunnerGroupCacheStore { + constructor(private readonly config: AwsSsmRunnerGroupCacheStoreConfig) {} + + async get(runnerGroupName: string): Promise { + const runnerGroupId = await getParameter(this.parameterName(runnerGroupName)); + return parseInt(runnerGroupId); + } + + async create(record: RunnerGroupCacheRecord): Promise { + await putParameter(this.parameterName(record.runnerGroupName), record.runnerGroupId.toString(), false, { + tags: this.config.parameterStoreTags, + }); + } + + private parameterName(runnerGroupName: string): string { + return `${this.config.configPath}/runner-group/${runnerGroupName}`; + } +} diff --git a/lambdas/libs/storage-providers/core/index.ts b/lambdas/libs/storage-providers/core/index.ts index 7f9df413c0..6fba06f035 100644 --- a/lambdas/libs/storage-providers/core/index.ts +++ b/lambdas/libs/storage-providers/core/index.ts @@ -1,4 +1,4 @@ -export interface RunnerConfigMetadataTag { +export interface RunnerConfigMetadata { key: string; value: string; } @@ -10,5 +10,16 @@ export interface RunnerConfigRecord { export interface RunnerConfigStore { readonly maxWritesPerSecond?: number; - create(record: RunnerConfigRecord, options?: { metadataTags?: RunnerConfigMetadataTag[] }): Promise; + create(record: RunnerConfigRecord, options?: { metadata?: RunnerConfigMetadata[] }): Promise; + houseKeeper(): Promise; +} + +export interface RunnerGroupCacheRecord { + runnerGroupName: string; + runnerGroupId: number; +} + +export interface RunnerGroupCacheStore { + get(runnerGroupName: string): Promise; + create(record: RunnerGroupCacheRecord): Promise; } diff --git a/lambdas/libs/storage-providers/index.ts b/lambdas/libs/storage-providers/index.ts index 862924118b..8001457df5 100644 --- a/lambdas/libs/storage-providers/index.ts +++ b/lambdas/libs/storage-providers/index.ts @@ -1,2 +1,9 @@ -export type { RunnerConfigMetadataTag, RunnerConfigRecord, RunnerConfigStore } from './core'; +export type { + RunnerConfigMetadata, + RunnerConfigRecord, + RunnerConfigStore, + RunnerGroupCacheRecord, + RunnerGroupCacheStore, +} from './core'; export { getRunnerConfigStore, resetRunnerConfigStore } from './runner-config'; +export { getRunnerGroupCacheStore, resetRunnerGroupCacheStore } from './runner-group-cache'; diff --git a/lambdas/libs/storage-providers/package.json b/lambdas/libs/storage-providers/package.json index 2b753fc436..745d026851 100644 --- a/lambdas/libs/storage-providers/package.json +++ b/lambdas/libs/storage-providers/package.json @@ -15,8 +15,14 @@ "format-check": "prettier --check \"**/*.ts\"", "all": "yarn format && yarn lint && yarn test" }, + "devDependencies": { + "aws-sdk-client-mock": "^4.1.0", + "aws-sdk-client-mock-jest": "^4.1.0" + }, "dependencies": { - "@aws-github-runner/aws-ssm-util": "*" + "@aws-github-runner/aws-powertools-util": "*", + "@aws-github-runner/aws-ssm-util": "*", + "@aws-sdk/client-ssm": "^3.1009.0" }, "nx": { "includedScripts": [ diff --git a/lambdas/libs/storage-providers/provider.ts b/lambdas/libs/storage-providers/provider.ts new file mode 100644 index 0000000000..82b2c7895b --- /dev/null +++ b/lambdas/libs/storage-providers/provider.ts @@ -0,0 +1,26 @@ +export const runnerConfigStorageProviders = ['aws_ssm'] as const; + +export type RunnerConfigStorageProvider = (typeof runnerConfigStorageProviders)[number]; + +const defaultProvider = 'aws_ssm' satisfies RunnerConfigStorageProvider; + +export function resolveRunnerConfigStorageProvider(provider: unknown): RunnerConfigStorageProvider { + if (provider === undefined) { + return defaultProvider; + } + + if (typeof provider !== 'string') { + throw new Error(`Unsupported runner config storage provider '${String(provider)}'`); + } + + const normalizedProvider = provider.trim().toLowerCase(); + if (normalizedProvider === '') { + return defaultProvider; + } + + if (!runnerConfigStorageProviders.includes(normalizedProvider as RunnerConfigStorageProvider)) { + throw new Error(`Unsupported runner config storage provider '${String(provider)}'`); + } + + return normalizedProvider as RunnerConfigStorageProvider; +} diff --git a/lambdas/libs/storage-providers/runner-config.test.ts b/lambdas/libs/storage-providers/runner-config.test.ts index 10467ebdb4..95ba3787dd 100644 --- a/lambdas/libs/storage-providers/runner-config.test.ts +++ b/lambdas/libs/storage-providers/runner-config.test.ts @@ -60,7 +60,7 @@ describe('runner config store selection', () => { const firstStore = stubStore(); expect(getRunnerConfigStore()).toBe(firstStore); - const secondStore = { create: vi.fn() } satisfies RunnerConfigStore; + const secondStore = { create: vi.fn(), houseKeeper: vi.fn() } satisfies RunnerConfigStore; createAwsSsmRunnerConfigStoreMock.mockReturnValue(secondStore); resetRunnerConfigStore(); @@ -78,7 +78,7 @@ function setProvider(provider: string | undefined): void { } function stubStore(): RunnerConfigStore { - const store = { create: vi.fn() } satisfies RunnerConfigStore; + const store = { create: vi.fn(), houseKeeper: vi.fn() } satisfies RunnerConfigStore; createAwsSsmRunnerConfigStoreMock.mockReturnValue(store); return store; } diff --git a/lambdas/libs/storage-providers/runner-config.ts b/lambdas/libs/storage-providers/runner-config.ts index dfac2fcfb1..180c808d78 100644 --- a/lambdas/libs/storage-providers/runner-config.ts +++ b/lambdas/libs/storage-providers/runner-config.ts @@ -1,21 +1,19 @@ import { createAwsSsmRunnerConfigStore } from './aws/ssm/runner-config-store'; import type { RunnerConfigStore } from './core'; import type {} from './environment'; +import { resolveRunnerConfigStorageProvider, type RunnerConfigStorageProvider } from './provider'; type RunnerConfigStoreFactory = () => RunnerConfigStore; const providerFactories = { aws_ssm: createAwsSsmRunnerConfigStore, -} as const satisfies Record; - -type RunnerConfigStorageProvider = keyof typeof providerFactories; - -const defaultProvider = 'aws_ssm' satisfies RunnerConfigStorageProvider; +} as const satisfies Record; let runnerConfigStore: RunnerConfigStore | undefined; export function getRunnerConfigStore(): RunnerConfigStore { - runnerConfigStore ??= providerFactories[resolveProvider(process.env.RUNNER_CONFIG_STORAGE_PROVIDER)](); + runnerConfigStore ??= + providerFactories[resolveRunnerConfigStorageProvider(process.env.RUNNER_CONFIG_STORAGE_PROVIDER)](); return runnerConfigStore; } @@ -23,24 +21,3 @@ export function getRunnerConfigStore(): RunnerConfigStore { export function resetRunnerConfigStore(): void { runnerConfigStore = undefined; } - -function resolveProvider(provider: unknown): RunnerConfigStorageProvider { - if (provider === undefined) { - return defaultProvider; - } - - if (typeof provider !== 'string') { - throw new Error(`Unsupported runner config storage provider '${String(provider)}'`); - } - - const normalizedProvider = provider.trim().toLowerCase(); - if (normalizedProvider === '') { - return defaultProvider; - } - - if (!Object.prototype.hasOwnProperty.call(providerFactories, normalizedProvider)) { - throw new Error(`Unsupported runner config storage provider '${String(provider)}'`); - } - - return normalizedProvider as RunnerConfigStorageProvider; -} diff --git a/lambdas/libs/storage-providers/runner-group-cache.test.ts b/lambdas/libs/storage-providers/runner-group-cache.test.ts new file mode 100644 index 0000000000..e67e307905 --- /dev/null +++ b/lambdas/libs/storage-providers/runner-group-cache.test.ts @@ -0,0 +1,83 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { createAwsSsmRunnerGroupCacheStore } from './aws/ssm/runner-group-cache-store'; +import type { RunnerGroupCacheStore } from './core'; +import { getRunnerGroupCacheStore, resetRunnerGroupCacheStore } from './runner-group-cache'; + +vi.mock('./aws/ssm/runner-group-cache-store', () => ({ + createAwsSsmRunnerGroupCacheStore: vi.fn(), +})); + +const createAwsSsmRunnerGroupCacheStoreMock = vi.mocked(createAwsSsmRunnerGroupCacheStore); +const cleanEnv = process.env; + +describe('runner group cache store selection', () => { + beforeEach(() => { + vi.clearAllMocks(); + process.env = { ...cleanEnv }; + delete process.env.RUNNER_CONFIG_STORAGE_PROVIDER; + resetRunnerGroupCacheStore(); + }); + + it.each([undefined, '', ' '])('uses aws_ssm for default selector input %j', (provider) => { + setProvider(provider); + const store = stubStore(); + + expect(getRunnerGroupCacheStore()).toBe(store); + expect(createAwsSsmRunnerGroupCacheStoreMock).toHaveBeenCalledOnce(); + }); + + it.each(['aws_ssm', ' AWS_SSM '])('uses aws_ssm for explicit selector input %j', (provider) => { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = provider; + const store = stubStore(); + + expect(getRunnerGroupCacheStore()).toBe(store); + expect(createAwsSsmRunnerGroupCacheStoreMock).toHaveBeenCalledOnce(); + }); + + it('rejects an unsupported provider on first use', () => { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'not-registered'; + + expect(() => getRunnerGroupCacheStore()).toThrow("Unsupported runner config storage provider 'not-registered'"); + expect(createAwsSsmRunnerGroupCacheStoreMock).not.toHaveBeenCalled(); + }); + + it('selects lazily and caches the created store', () => { + const store = stubStore(); + + expect(createAwsSsmRunnerGroupCacheStoreMock).not.toHaveBeenCalled(); + const first = getRunnerGroupCacheStore(); + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'not-registered'; + const second = getRunnerGroupCacheStore(); + + expect(first).toBe(store); + expect(second).toBe(store); + expect(createAwsSsmRunnerGroupCacheStoreMock).toHaveBeenCalledOnce(); + }); + + it('selects again after the test reset', () => { + const firstStore = stubStore(); + expect(getRunnerGroupCacheStore()).toBe(firstStore); + + const secondStore = { get: vi.fn(), create: vi.fn() } satisfies RunnerGroupCacheStore; + createAwsSsmRunnerGroupCacheStoreMock.mockReturnValue(secondStore); + resetRunnerGroupCacheStore(); + + expect(getRunnerGroupCacheStore()).toBe(secondStore); + expect(createAwsSsmRunnerGroupCacheStoreMock).toHaveBeenCalledTimes(2); + }); +}); + +function setProvider(provider: string | undefined): void { + if (provider === undefined) { + delete process.env.RUNNER_CONFIG_STORAGE_PROVIDER; + } else { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = provider; + } +} + +function stubStore(): RunnerGroupCacheStore { + const store = { get: vi.fn(), create: vi.fn() } satisfies RunnerGroupCacheStore; + createAwsSsmRunnerGroupCacheStoreMock.mockReturnValue(store); + return store; +} diff --git a/lambdas/libs/storage-providers/runner-group-cache.ts b/lambdas/libs/storage-providers/runner-group-cache.ts new file mode 100644 index 0000000000..a28b00f48e --- /dev/null +++ b/lambdas/libs/storage-providers/runner-group-cache.ts @@ -0,0 +1,23 @@ +import { createAwsSsmRunnerGroupCacheStore } from './aws/ssm/runner-group-cache-store'; +import type { RunnerGroupCacheStore } from './core'; +import type {} from './environment'; +import { resolveRunnerConfigStorageProvider, type RunnerConfigStorageProvider } from './provider'; + +type RunnerGroupCacheStoreFactory = () => RunnerGroupCacheStore; + +const providerFactories = { + aws_ssm: createAwsSsmRunnerGroupCacheStore, +} as const satisfies Record; + +let runnerGroupCacheStore: RunnerGroupCacheStore | undefined; + +export function getRunnerGroupCacheStore(): RunnerGroupCacheStore { + runnerGroupCacheStore ??= + providerFactories[resolveRunnerConfigStorageProvider(process.env.RUNNER_CONFIG_STORAGE_PROVIDER)](); + return runnerGroupCacheStore; +} + +// Test-only reset for cases that need to exercise first-use environment selection. +export function resetRunnerGroupCacheStore(): void { + runnerGroupCacheStore = undefined; +} diff --git a/lambdas/libs/storage-providers/vitest.config.ts b/lambdas/libs/storage-providers/vitest.config.ts index a5812ad13e..af85b8946d 100644 --- a/lambdas/libs/storage-providers/vitest.config.ts +++ b/lambdas/libs/storage-providers/vitest.config.ts @@ -7,7 +7,7 @@ export default mergeConfig(defaultConfig, { test: { setupFiles: [resolve(__dirname, '../../aws-vitest-setup.ts')], coverage: { - include: ['index.ts', 'runner-config.ts', 'core/**/*.ts', 'aws/**/*.ts'], + include: ['index.ts', 'provider.ts', 'runner-config.ts', 'runner-group-cache.ts', 'core/**/*.ts', 'aws/**/*.ts'], exclude: ['**/*.test.ts', '**/*.d.ts'], }, }, diff --git a/lambdas/yarn.lock b/lambdas/yarn.lock index 84a1a7069c..836224f666 100644 --- a/lambdas/yarn.lock +++ b/lambdas/yarn.lock @@ -214,7 +214,11 @@ __metadata: version: 0.0.0-use.local resolution: "@aws-github-runner/storage-providers@workspace:libs/storage-providers" dependencies: + "@aws-github-runner/aws-powertools-util": "npm:*" "@aws-github-runner/aws-ssm-util": "npm:*" + "@aws-sdk/client-ssm": "npm:^3.1009.0" + aws-sdk-client-mock: "npm:^4.1.0" + aws-sdk-client-mock-jest: "npm:^4.1.0" languageName: unknown linkType: soft From 82a4165cf6ee99322b55dfb2fe28c30cfdd9cb8f Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Tue, 18 Aug 2026 23:43:43 +0200 Subject: [PATCH 3/6] refactor(storage): move local housekeeper harness --- .../aws/ssm/local-runner-config-housekeeper.ts} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename lambdas/{functions/control-plane/src/local-ssm-housekeeper.ts => libs/storage-providers/aws/ssm/local-runner-config-housekeeper.ts} (71%) diff --git a/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts b/lambdas/libs/storage-providers/aws/ssm/local-runner-config-housekeeper.ts similarity index 71% rename from lambdas/functions/control-plane/src/local-ssm-housekeeper.ts rename to lambdas/libs/storage-providers/aws/ssm/local-runner-config-housekeeper.ts index 08b062a193..79518a8157 100644 --- a/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts +++ b/lambdas/libs/storage-providers/aws/ssm/local-runner-config-housekeeper.ts @@ -1,4 +1,4 @@ -import { getRunnerConfigStore } from '@aws-github-runner/storage-providers'; +import { createAwsSsmRunnerConfigStore } from './runner-config-store'; export function run(): void { process.env.SSM_CLEANUP_CONFIG = JSON.stringify({ @@ -7,7 +7,7 @@ export function run(): void { tokenPath: '/ghr/my-env/runners/tokens', }); - getRunnerConfigStore() + createAwsSsmRunnerConfigStore() .houseKeeper() .then() .catch((e) => { From 06d882aaf6f9c8f126ce1792c37b48d89e4c5b39 Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Tue, 18 Aug 2026 23:48:43 +0200 Subject: [PATCH 4/6] refactor(storage): preserve SSM cleanup names --- .../aws/ssm/runner-config-housekeeper.ts | 6 +++--- .../libs/storage-providers/aws/ssm/runner-config-store.ts | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts index 2c0c22359c..30bc1d20ca 100644 --- a/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts @@ -1,13 +1,13 @@ import { DeleteParameterCommand, GetParametersByPathCommand, SSMClient } from '@aws-sdk/client-ssm'; import { getTracedAWSV3Client, logger } from '@aws-github-runner/aws-powertools-util'; -export interface SsmRunnerConfigCleanupOptions { +export interface SSMCleanupOptions { dryRun: boolean; minimumDaysOld: number; tokenPath: string; } -function validateOptions(options: SsmRunnerConfigCleanupOptions): void { +function validateOptions(options: SSMCleanupOptions): void { const errorMessages: string[] = []; if (!options.minimumDaysOld || options.minimumDaysOld < 1) { errorMessages.push(`minimumDaysOld must be greater then 0, value is set to "${options.minimumDaysOld}"`); @@ -20,7 +20,7 @@ function validateOptions(options: SsmRunnerConfigCleanupOptions): void { } } -export async function cleanSsmRunnerConfigs(options: SsmRunnerConfigCleanupOptions): Promise { +export async function cleanSSMTokens(options: SSMCleanupOptions): Promise { logger.info(`Cleaning tokens / JIT config older then ${options.minimumDaysOld} days, dryRun: ${options.dryRun}`); logger.debug('Cleaning with options', { options }); validateOptions(options); diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts index dec2241086..1959a6192a 100644 --- a/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-store.ts @@ -3,19 +3,19 @@ import { putParameter } from '@aws-github-runner/aws-ssm-util'; import type { RunnerConfigMetadata, RunnerConfigRecord, RunnerConfigStore } from '../../core'; import type {} from './environment'; import { loadSsmParameterStoreTagsFromEnvironment } from './parameter-store-tags'; -import { cleanSsmRunnerConfigs, type SsmRunnerConfigCleanupOptions } from './runner-config-housekeeper'; +import { cleanSSMTokens, type SSMCleanupOptions } from './runner-config-housekeeper'; interface AwsSsmRunnerConfigStoreConfig { tokenPath?: string; parameterStoreTags: { Key: string; Value: string }[]; - cleanupOptions?: SsmRunnerConfigCleanupOptions; + cleanupOptions?: SSMCleanupOptions; } export function createAwsSsmRunnerConfigStore(): RunnerConfigStore { const tokenPath = process.env.SSM_TOKEN_PATH; const cleanupOptions = process.env.SSM_CLEANUP_CONFIG !== undefined - ? (JSON.parse(process.env.SSM_CLEANUP_CONFIG) as SsmRunnerConfigCleanupOptions) + ? (JSON.parse(process.env.SSM_CLEANUP_CONFIG) as SSMCleanupOptions) : undefined; const hasWriterConfig = tokenPath !== undefined && tokenPath.trim() !== ''; @@ -53,6 +53,6 @@ class AwsSsmRunnerConfigStore implements RunnerConfigStore { throw new Error('Environment variable SSM_CLEANUP_CONFIG is not set'); } - await cleanSsmRunnerConfigs(this.config.cleanupOptions); + await cleanSSMTokens(this.config.cleanupOptions); } } From 66a0527f0410ded307780802c5d59b70866f43ad Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Wed, 19 Aug 2026 00:09:49 +0200 Subject: [PATCH 5/6] test(storage): decouple scale-up tests from SSM --- .../src/scale-runners/scale-up.test.ts | 508 +++++++----------- lambdas/libs/aws-ssm-util/src/index.test.ts | 21 + 2 files changed, 217 insertions(+), 312 deletions(-) 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 9ae7e6a5da..83cb77cd5e 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 @@ -1,9 +1,12 @@ -import { PutParameterCommand, SSMClient } from '@aws-sdk/client-ssm'; -import { mockClient } from 'aws-sdk-client-mock'; -import 'aws-sdk-client-mock-jest/vitest'; -// Using vi.mocked instead of jest-mock +import { + getRunnerConfigStore, + getRunnerGroupCacheStore, + type RunnerConfigStore, + type RunnerGroupCacheStore, +} from '@aws-github-runner/storage-providers'; +import type { Octokit } from '@octokit/rest'; import nock from 'nock'; -import { performance } from 'perf_hooks'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import { controlPlaneProviderRegistry } from '../control-plane-providers'; import * as ghAuth from '../github/auth'; @@ -16,10 +19,6 @@ import type { CreateScaleUpRunnersInput, ScaleUpComputeProvider, } from './types'; -import { getParameter } from '@aws-github-runner/aws-ssm-util'; -import { resetRunnerConfigStore, resetRunnerGroupCacheStore } from '@aws-github-runner/storage-providers'; -import { beforeEach, describe, expect, it, vi } from 'vitest'; -import type { Octokit } from '@octokit/rest'; const mockOctokit = { paginate: vi.fn(), @@ -54,8 +53,21 @@ const createRunner = vi.fn<(input: TestRunnerCreationInput) => Promise Promise>(); const mockCreateRunner = vi.mocked(createRunner); const mockListRunners = vi.mocked(listRunners); -const mockSSMClient = mockClient(SSMClient); -const mockSSMgetParameter = vi.mocked(getParameter); +const mockGetRunnerConfigStore = vi.mocked(getRunnerConfigStore); +const mockGetRunnerGroupCacheStore = vi.mocked(getRunnerGroupCacheStore); +const mockRunnerConfigCreate = vi.fn(); +const mockRunnerConfigHouseKeeper = vi.fn(); +const mockRunnerGroupCacheGet = vi.fn(); +const mockRunnerGroupCacheCreate = vi.fn(); +const mockRunnerConfigStore: RunnerConfigStore = { + maxWritesPerSecond: 40, + create: mockRunnerConfigCreate, + houseKeeper: mockRunnerConfigHouseKeeper, +}; +const mockRunnerGroupCacheStore: RunnerGroupCacheStore = { + get: mockRunnerGroupCacheGet, + create: mockRunnerGroupCacheCreate, +}; const mockPublishRetryMessage = vi.mocked(publishRetryMessage); const testProviderState = { provider: 'test' }; const mockComputeProvider: ScaleUpComputeProvider = { @@ -87,16 +99,10 @@ vi.mock('../github/auth', async () => ({ getStoredInstallationId: vi.fn().mockResolvedValue(undefined), })); -vi.mock('@aws-github-runner/aws-ssm-util', async () => { - const actual = (await vi.importActual( - '@aws-github-runner/aws-ssm-util', - )) as typeof import('@aws-github-runner/aws-ssm-util'); - - return { - ...actual, - getParameter: vi.fn(), - }; -}); +vi.mock('@aws-github-runner/storage-providers', () => ({ + getRunnerConfigStore: vi.fn(), + getRunnerGroupCacheStore: vi.fn(), +})); vi.mock('./job-retry', () => ({ publishRetryMessage: vi.fn(), @@ -148,8 +154,6 @@ function setDefaults() { process.env.GITHUB_APP_CLIENT_SECRET = 'TEST_CLIENT_SECRET'; process.env.RUNNERS_MAXIMUM_COUNT = '3'; process.env.ENVIRONMENT = EXPECTED_RUNNER_PARAMS.environment; - process.env.SSM_CONFIG_PATH = '/github-action-runners/default/runners/config'; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; } async function createTestProviderRunners(input: CreateScaleUpRunnersInput): Promise { @@ -190,10 +194,12 @@ beforeEach(() => { vi.resetModules(); vi.clearAllMocks(); setDefaults(); - resetRunnerConfigStore(); - resetRunnerGroupCacheStore(); - - defaultSSMGetParameterMockImpl(); + mockGetRunnerConfigStore.mockReturnValue(mockRunnerConfigStore); + mockGetRunnerGroupCacheStore.mockReturnValue(mockRunnerGroupCacheStore); + mockRunnerConfigCreate.mockResolvedValue(); + mockRunnerConfigHouseKeeper.mockResolvedValue(); + mockRunnerGroupCacheGet.mockResolvedValue(1); + mockRunnerGroupCacheCreate.mockResolvedValue(); defaultOctokitMockImpl(); mockedResolveCapability.mockReturnValue(() => mockComputeProvider); @@ -273,12 +279,9 @@ describe('scaleUp with GHES', () => { process.env.ENABLE_EPHEMERAL_RUNNERS = 'true'; process.env.RUNNER_NAME_PREFIX = 'unit-test-'; process.env.RUNNER_GROUP_NAME = 'Default'; - process.env.SSM_CONFIG_PATH = '/github-action-runners/default/runners/config'; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; process.env.RUNNER_LABELS = 'label1,label2'; expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; - mockSSMClient.reset(); }); it('does not create a token when maximum runners has been reached', async () => { @@ -333,9 +336,7 @@ describe('scaleUp with GHES', () => { it('returns a retryable failure if runner group lookup fails for ephemeral runners', async () => { process.env.RUNNER_GROUP_NAME = 'test-runner-group'; - mockSSMgetParameter.mockImplementation(async () => { - throw new Error('ParameterNotFound'); - }); + mockRunnerGroupCacheGet.mockRejectedValue(new Error('Cache entry not found')); await expect(scaleUpModule.scaleUp(TEST_DATA)).resolves.toEqual(['foobar']); @@ -350,24 +351,27 @@ describe('scaleUp with GHES', () => { expect(createRunner).not.toHaveBeenCalled(); }); - it('create SSM parameter for runner group id if it does not exist', async () => { - mockSSMgetParameter.mockImplementation(async () => { - throw new Error('ParameterNotFound'); - }); + it('caches the runner group id if it does not exist', async () => { + mockRunnerGroupCacheGet.mockResolvedValue(undefined); + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockRunnerGroupCacheGet).toHaveBeenCalledWith('Default'); expect(mockOctokit.paginate).toHaveBeenCalledTimes(1); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 2); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: `${process.env.SSM_CONFIG_PATH}/runner-group/${process.env.RUNNER_GROUP_NAME}`, - Value: '1', - Type: 'String', + expect(mockRunnerGroupCacheCreate).toHaveBeenCalledWith({ + runnerGroupName: 'Default', + runnerGroupId: 1, }); + expect(mockRunnerConfigCreate).toHaveBeenCalledTimes(1); }); - it('Does not create SSM parameter for runner group id if it exists', async () => { + it('reuses a cached runner group id', async () => { await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockRunnerGroupCacheGet).toHaveBeenCalledWith('Default'); expect(mockOctokit.paginate).toHaveBeenCalledTimes(0); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 1); + expect(mockRunnerGroupCacheCreate).not.toHaveBeenCalled(); + expect(mockRunnerConfigCreate).toHaveBeenCalledTimes(1); }); it('create start runner config for ephemeral runners ', async () => { @@ -380,17 +384,10 @@ describe('scaleUp with GHES', () => { runner_group_id: 1, labels: ['label1', 'label2'], }); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-12345', - Value: 'TEST_JIT_CONFIG_ORG', - Type: 'SecureString', - Tags: [ - { - Key: 'RunnerId', - Value: 'i-12345', - }, - ], - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-12345', value: 'TEST_JIT_CONFIG_ORG' }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); it('create start runner config for non-ephemeral runners ', async () => { @@ -399,19 +396,15 @@ describe('scaleUp with GHES', () => { await scaleUpModule.scaleUp(TEST_DATA); expect(mockOctokit.actions.generateRunnerJitconfigForOrg).not.toBeCalled(); expect(mockOctokit.actions.createRegistrationTokenForOrg).toBeCalled(); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-12345', - Value: - '--url https://github.enterprise.something/Codertocat --token 1234abcd ' + - '--labels label1,label2 --runnergroup Default', - Type: 'SecureString', - Tags: [ - { - Key: 'RunnerId', - Value: 'i-12345', - }, - ], - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { + runnerId: 'i-12345', + value: + '--url https://github.enterprise.something/Codertocat --token 1234abcd ' + + '--labels label1,label2 --runnergroup Default', + }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); it('quotes runner labels with semicolon separators in non-ephemeral runner config', async () => { @@ -426,19 +419,15 @@ describe('scaleUp with GHES', () => { }, ]); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-12345', - Value: - '--url https://github.enterprise.something/Codertocat --token 1234abcd ' + - "--labels 'label1,label2,ghr-provider-capability:intel;amd' --runnergroup Default", - Type: 'SecureString', - Tags: [ - { - Key: 'RunnerId', - Value: 'i-12345', - }, - ], - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { + runnerId: 'i-12345', + value: + '--url https://github.enterprise.something/Codertocat --token 1234abcd ' + + "--labels 'label1,label2,ghr-provider-capability:intel;amd' --runnergroup Default", + }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); it('should create JIT config for all remaining instances even when GitHub API fails for one instance', async () => { @@ -498,23 +487,18 @@ describe('scaleUp with GHES', () => { labels: ['label1', 'label2'], }); - expect(mockSSMClient).toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-1', - Value: 'TEST_JIT_CONFIG_unit-test-i-instance-1', - Type: 'SecureString', - Tags: [{ Key: 'RunnerId', Value: 'i-instance-1' }], - }); - - expect(mockSSMClient).toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-3', - Value: 'TEST_JIT_CONFIG_unit-test-i-instance-3', - Type: 'SecureString', - Tags: [{ Key: 'RunnerId', Value: 'i-instance-3' }], - }); - - expect(mockSSMClient).not.toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-2', - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-instance-1', value: 'TEST_JIT_CONFIG_unit-test-i-instance-1' }, + { metadata: [{ key: 'RunnerId', value: 'i-instance-1' }] }, + ); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-instance-3', value: 'TEST_JIT_CONFIG_unit-test-i-instance-3' }, + { metadata: [{ key: 'RunnerId', value: 'i-instance-3' }] }, + ); + expect(mockRunnerConfigCreate).not.toHaveBeenCalledWith( + expect.objectContaining({ runnerId: 'i-instance-2' }), + expect.anything(), + ); }); it('should handle retryable errors with error handling logic', async () => { @@ -550,16 +534,14 @@ describe('scaleUp with GHES', () => { await scaleUpModule.scaleUp(TEST_DATA); - expect(mockSSMClient).toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-2', - Value: 'TEST_JIT_CONFIG_unit-test-i-instance-2', - Type: 'SecureString', - Tags: [{ Key: 'RunnerId', Value: 'i-instance-2' }], - }); - - expect(mockSSMClient).not.toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-1', - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-instance-2', value: 'TEST_JIT_CONFIG_unit-test-i-instance-2' }, + { metadata: [{ key: 'RunnerId', value: 'i-instance-2' }] }, + ); + expect(mockRunnerConfigCreate).not.toHaveBeenCalledWith( + expect.objectContaining({ runnerId: 'i-instance-1' }), + expect.anything(), + ); }); it('should handle non-retryable 4xx errors gracefully', async () => { @@ -596,79 +578,62 @@ describe('scaleUp with GHES', () => { await scaleUpModule.scaleUp(TEST_DATA); - expect(mockSSMClient).toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-2', - Value: 'TEST_JIT_CONFIG_unit-test-i-instance-2', - Type: 'SecureString', - Tags: [{ Key: 'RunnerId', Value: 'i-instance-2' }], - }); - - expect(mockSSMClient).not.toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-1', - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-instance-2', value: 'TEST_JIT_CONFIG_unit-test-i-instance-2' }, + { metadata: [{ key: 'RunnerId', value: 'i-instance-2' }] }, + ); + expect(mockRunnerConfigCreate).not.toHaveBeenCalledWith( + expect.objectContaining({ runnerId: 'i-instance-1' }), + expect.anything(), + ); }); it.each(RUNNER_TYPES)( - 'calls create start runner config of 40' + ' instances (ssm rate limit condition) to test time delay ', + 'paces 40 runner-config writes at the store throughput limit for %s runners', async (type: RunnerLifecycle) => { process.env.ENABLE_EPHEMERAL_RUNNERS = type === 'ephemeral' ? 'true' : 'false'; process.env.RUNNERS_MAXIMUM_COUNT = '40'; + const instances = Array.from({ length: 40 }, (_, index) => `i-${index + 1}`); mockCreateRunner.mockImplementation(async () => { return createRunnerResult(instances); }); mockListRunners.mockImplementation(async () => { return []; }); - const startTime = performance.now(); - const instances = [ - 'i-1234', - 'i-5678', - 'i-5567', - 'i-5569', - 'i-5561', - 'i-5560', - 'i-5566', - 'i-5536', - 'i-5526', - 'i-5516', - 'i-122', - 'i-123', - 'i-124', - 'i-125', - 'i-126', - 'i-127', - 'i-128', - 'i-129', - 'i-130', - 'i-131', - 'i-132', - 'i-133', - 'i-134', - 'i-135', - 'i-136', - 'i-137', - 'i-138', - 'i-139', - 'i-140', - 'i-141', - 'i-142', - 'i-143', - 'i-144', - 'i-145', - 'i-146', - 'i-147', - 'i-148', - 'i-149', - 'i-150', - 'i-151', - ]; - await scaleUpModule.scaleUp(TEST_DATA); - const endTime = performance.now(); - expect(endTime - startTime).toBeGreaterThan(1000); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 40); + + const setTimeoutSpy = vi.spyOn(globalThis, 'setTimeout').mockImplementation((callback) => { + callback(); + return 0 as unknown as NodeJS.Timeout; + }); + try { + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockRunnerConfigCreate).toHaveBeenCalledTimes(40); + expect(setTimeoutSpy).toHaveBeenCalledTimes(40); + expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 25); + } finally { + setTimeoutSpy.mockRestore(); + } }, - 10000, ); + + it('does not pace 39 runner-config writes below the store throughput limit', async () => { + process.env.ENABLE_EPHEMERAL_RUNNERS = 'false'; + process.env.RUNNERS_MAXIMUM_COUNT = '39'; + const instances = Array.from({ length: 39 }, (_, index) => `i-${index + 1}`); + mockCreateRunner.mockResolvedValue(createRunnerResult(instances)); + mockListRunners.mockResolvedValue([]); + const setTimeoutSpy = vi.spyOn(globalThis, 'setTimeout'); + + try { + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockRunnerConfigCreate).toHaveBeenCalledTimes(39); + expect(setTimeoutSpy).not.toHaveBeenCalled(); + } finally { + setTimeoutSpy.mockRestore(); + } + }); }); describe('dynamic label groups', () => { @@ -679,7 +644,6 @@ describe('scaleUp with GHES', () => { process.env.RUNNER_LABELS = 'base-label'; process.env.RUNNER_NAME_PREFIX = 'unit-test'; expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; - mockSSMClient.reset(); mockResolveLabelsForRunners.mockImplementation(async (labels) => ({ runnerLabels: labels.filter((label) => label.startsWith('ghr-')), @@ -1155,8 +1119,6 @@ describe('scaleUp with public GH', () => { describe('on repo level', () => { beforeEach(() => { - mockSSMClient.reset(); - process.env.ENABLE_ORGANIZATION_RUNNERS = 'false'; process.env.RUNNER_NAME_PREFIX = 'unit-test'; expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; @@ -1200,45 +1162,32 @@ describe('scaleUp with public GH', () => { it('creates a ephemeral runner with JIT config.', async () => { process.env.ENABLE_EPHEMERAL_RUNNERS = 'true'; process.env.ENABLE_JOB_QUEUED_CHECK = 'false'; - delete process.env.SSM_CONFIG_PATH; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; await scaleUpModule.scaleUp(TEST_DATA); expect(mockOctokit.actions.getJobForWorkflowRun).not.toBeCalled(); expect(createRunner).toBeCalledWith(expectedRunnerParams); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-12345', - Value: 'TEST_JIT_CONFIG_REPO', - Type: 'SecureString', - Tags: [ - { - Key: 'RunnerId', - Value: 'i-12345', - }, - ], - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-12345', value: 'TEST_JIT_CONFIG_REPO' }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); + expect(mockGetRunnerGroupCacheStore).not.toHaveBeenCalled(); }); it('creates a ephemeral runner with registration token.', async () => { process.env.ENABLE_EPHEMERAL_RUNNERS = 'true'; process.env.ENABLE_JIT_CONFIG = 'false'; process.env.ENABLE_JOB_QUEUED_CHECK = 'false'; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; await scaleUpModule.scaleUp(TEST_DATA); expect(mockOctokit.actions.getJobForWorkflowRun).not.toBeCalled(); expect(createRunner).toBeCalledWith(expectedRunnerParams); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-12345', - Value: '--url https://github.com/Codertocat/hello-world --token 1234abcd --ephemeral', - Type: 'SecureString', - Tags: [ - { - Key: 'RunnerId', - Value: 'i-12345', - }, - ], - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { + runnerId: 'i-12345', + value: '--url https://github.com/Codertocat/hello-world --token 1234abcd --ephemeral', + }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); it('JIT config is ignored for non-ephemeral runners.', async () => { @@ -1246,23 +1195,18 @@ describe('scaleUp with public GH', () => { process.env.ENABLE_JIT_CONFIG = 'true'; process.env.ENABLE_JOB_QUEUED_CHECK = 'false'; process.env.RUNNER_LABELS = 'jit'; - delete process.env.SSM_CONFIG_PATH; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; await scaleUpModule.scaleUp(TEST_DATA); expect(mockOctokit.actions.getJobForWorkflowRun).not.toBeCalled(); expect(createRunner).toBeCalledWith(expectedRunnerParams); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-12345', - Value: '--url https://github.com/Codertocat/hello-world --token 1234abcd --labels jit', - Type: 'SecureString', - Tags: [ - { - Key: 'RunnerId', - Value: 'i-12345', - }, - ], - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { + runnerId: 'i-12345', + value: '--url https://github.com/Codertocat/hello-world --token 1234abcd --labels jit', + }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); + expect(mockGetRunnerGroupCacheStore).not.toHaveBeenCalled(); }); it('creates a ephemeral runner after checking job is queued.', async () => { @@ -1541,12 +1485,9 @@ describe('scaleUp with Github Data Residency', () => { process.env.ENABLE_EPHEMERAL_RUNNERS = 'true'; process.env.RUNNER_NAME_PREFIX = 'unit-test-'; process.env.RUNNER_GROUP_NAME = 'Default'; - process.env.SSM_CONFIG_PATH = '/github-action-runners/default/runners/config'; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; process.env.RUNNER_LABELS = 'label1,label2'; expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; - mockSSMClient.reset(); }); it('does not create a token when maximum runners has been reached', async () => { @@ -1589,24 +1530,27 @@ describe('scaleUp with Github Data Residency', () => { expect(createRunner).not.toHaveBeenCalled(); }); - it('create SSM parameter for runner group id if it does not exist', async () => { - mockSSMgetParameter.mockImplementation(async () => { - throw new Error('ParameterNotFound'); - }); + it('caches the runner group id if it does not exist', async () => { + mockRunnerGroupCacheGet.mockResolvedValue(undefined); + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockRunnerGroupCacheGet).toHaveBeenCalledWith('Default'); expect(mockOctokit.paginate).toHaveBeenCalledTimes(1); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 2); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: `${process.env.SSM_CONFIG_PATH}/runner-group/${process.env.RUNNER_GROUP_NAME}`, - Value: '1', - Type: 'String', + expect(mockRunnerGroupCacheCreate).toHaveBeenCalledWith({ + runnerGroupName: 'Default', + runnerGroupId: 1, }); + expect(mockRunnerConfigCreate).toHaveBeenCalledTimes(1); }); - it('Does not create SSM parameter for runner group id if it exists', async () => { + it('reuses a cached runner group id', async () => { await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockRunnerGroupCacheGet).toHaveBeenCalledWith('Default'); expect(mockOctokit.paginate).toHaveBeenCalledTimes(0); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 1); + expect(mockRunnerGroupCacheCreate).not.toHaveBeenCalled(); + expect(mockRunnerConfigCreate).toHaveBeenCalledTimes(1); }); it('create start runner config for ephemeral runners ', async () => { @@ -1619,17 +1563,10 @@ describe('scaleUp with Github Data Residency', () => { runner_group_id: 1, labels: ['label1', 'label2'], }); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-12345', - Value: 'TEST_JIT_CONFIG_ORG', - Type: 'SecureString', - Tags: [ - { - Key: 'RunnerId', - Value: 'i-12345', - }, - ], - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-12345', value: 'TEST_JIT_CONFIG_ORG' }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); it('create start runner config for non-ephemeral runners ', async () => { @@ -1638,80 +1575,43 @@ describe('scaleUp with Github Data Residency', () => { await scaleUpModule.scaleUp(TEST_DATA); expect(mockOctokit.actions.generateRunnerJitconfigForOrg).not.toBeCalled(); expect(mockOctokit.actions.createRegistrationTokenForOrg).toBeCalled(); - expect(mockSSMClient).toHaveReceivedNthSpecificCommandWith(1, PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-12345', - Value: - '--url https://companyname.ghe.com/Codertocat --token 1234abcd ' + - '--labels label1,label2 --runnergroup Default', - Type: 'SecureString', - Tags: [ - { - Key: 'RunnerId', - Value: 'i-12345', - }, - ], - }); + expect(mockRunnerConfigCreate).toHaveBeenCalledWith( + { + runnerId: 'i-12345', + value: + '--url https://companyname.ghe.com/Codertocat --token 1234abcd ' + + '--labels label1,label2 --runnergroup Default', + }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); it.each(RUNNER_TYPES)( - 'calls create start runner config of 40' + ' instances (ssm rate limit condition) to test time delay ', + 'paces 40 runner-config writes at the store throughput limit for %s runners', async (type: RunnerLifecycle) => { process.env.ENABLE_EPHEMERAL_RUNNERS = type === 'ephemeral' ? 'true' : 'false'; process.env.RUNNERS_MAXIMUM_COUNT = '40'; + const instances = Array.from({ length: 40 }, (_, index) => `i-${index + 1}`); mockCreateRunner.mockImplementation(async () => { return createRunnerResult(instances); }); mockListRunners.mockImplementation(async () => { return []; }); - const startTime = performance.now(); - const instances = [ - 'i-1234', - 'i-5678', - 'i-5567', - 'i-5569', - 'i-5561', - 'i-5560', - 'i-5566', - 'i-5536', - 'i-5526', - 'i-5516', - 'i-122', - 'i-123', - 'i-124', - 'i-125', - 'i-126', - 'i-127', - 'i-128', - 'i-129', - 'i-130', - 'i-131', - 'i-132', - 'i-133', - 'i-134', - 'i-135', - 'i-136', - 'i-137', - 'i-138', - 'i-139', - 'i-140', - 'i-141', - 'i-142', - 'i-143', - 'i-144', - 'i-145', - 'i-146', - 'i-147', - 'i-148', - 'i-149', - 'i-150', - 'i-151', - ]; - await scaleUpModule.scaleUp(TEST_DATA); - const endTime = performance.now(); - expect(endTime - startTime).toBeGreaterThan(1000); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 40); + + const setTimeoutSpy = vi.spyOn(globalThis, 'setTimeout').mockImplementation((callback) => { + callback(); + return 0 as unknown as NodeJS.Timeout; + }); + try { + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockRunnerConfigCreate).toHaveBeenCalledTimes(40); + expect(setTimeoutSpy).toHaveBeenCalledTimes(40); + expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 25); + } finally { + setTimeoutSpy.mockRestore(); + } }, - 10000, ); }); describe('on repo level', () => { @@ -1996,7 +1896,6 @@ describe('Retry mechanism tests', () => { process.env.ENABLE_JOB_QUEUED_CHECK = 'true'; process.env.RUNNERS_MAXIMUM_COUNT = '10'; expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; - mockSSMClient.reset(); }); const createTestMessages = ( @@ -2182,11 +2081,8 @@ describe('Multi-app round-robin', () => { process.env.RUNNERS_MAXIMUM_COUNT = '10'; process.env.RUNNER_NAME_PREFIX = 'unit-test-'; process.env.RUNNER_GROUP_NAME = 'Default'; - process.env.SSM_CONFIG_PATH = '/github-action-runners/default/runners/config'; - process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config'; process.env.RUNNER_LABELS = 'label1,label2'; expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; - mockSSMClient.reset(); }); it('passes the same appIndex to createGithubInstallationAuth when multi-app is active', async () => { @@ -2270,7 +2166,7 @@ describe('Multi-app round-robin', () => { }); it('stored installationId takes precedence over webhook payload for additional app', async () => { - // Additional app (index 1) with a pre-configured installation id stored in SSM + // Additional app (index 1) with a pre-configured installation id mockedGetAppCount.mockResolvedValue(2); mockedGetStoredInstallationId.mockResolvedValue(77); mockedAppAuth.mockResolvedValue({ @@ -2340,15 +2236,3 @@ function defaultOctokitMockImpl() { mockOctokit.apps.getOrgInstallation.mockImplementation(() => mockInstallationIdReturnValueOrgs); mockOctokit.apps.getRepoInstallation.mockImplementation(() => mockInstallationIdReturnValueRepos); } - -function defaultSSMGetParameterMockImpl() { - mockSSMgetParameter.mockImplementation(async (name: string) => { - if (name === `${process.env.SSM_CONFIG_PATH}/runner-group/${process.env.RUNNER_GROUP_NAME}`) { - return '1'; - } else if (name === `${process.env.PARAMETER_GITHUB_APP_ID_NAME}`) { - return `${process.env.GITHUB_APP_ID}`; - } else { - throw new Error(`ParameterNotFound: ${name}`); - } - }); -} diff --git a/lambdas/libs/aws-ssm-util/src/index.test.ts b/lambdas/libs/aws-ssm-util/src/index.test.ts index 8a1d8d3864..f027b16293 100644 --- a/lambdas/libs/aws-ssm-util/src/index.test.ts +++ b/lambdas/libs/aws-ssm-util/src/index.test.ts @@ -127,6 +127,27 @@ describe('Test getParameter and putParameter', () => { }); }); + it('passes tags to the PutParameter command', async () => { + const parameterValue = 'test'; + const parameterName = 'testParam'; + const tags = [{ Key: 'InstanceId', Value: 'i-123' }]; + const output: PutParameterCommandOutput = { + $metadata: { + httpStatusCode: 200, + }, + }; + mockSSMClient.on(PutParameterCommand).resolves(output); + + await putParameter(parameterName, parameterValue, true, { tags }); + + expect(mockSSMClient).toHaveReceivedCommandWith(PutParameterCommand, { + Name: parameterName, + Value: parameterValue, + Type: 'SecureString', + Tags: tags, + }); + }); + it('Gets invalid parameters and returns string', async () => { // Arrange const parameterName = 'invalid'; From 47b53cb79d64955ab6671f2feac9a66fa1d0cbd3 Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Wed, 19 Aug 2026 10:15:33 +0200 Subject: [PATCH 6/6] refactor(storage): extract GitHub App credentials --- lambdas/functions/control-plane/package.json | 1 - .../control-plane/src/github/auth.test.ts | 228 +++++------------- .../control-plane/src/github/auth.ts | 54 +---- .../src/github/rate-limit.test.ts | 181 ++++++-------- .../control-plane/src/github/rate-limit.ts | 16 +- .../control-plane/src/lambda.test.ts | 1 - .../functions/control-plane/src/modules.d.ts | 2 - .../src/scale-runners/scale-up.test.ts | 1 - .../aws/ssm/environment.d.ts | 3 + .../ssm/github-app-credentials-store.test.ts | 154 ++++++++++++ .../aws/ssm/github-app-credentials-store.ts | 61 +++++ lambdas/libs/storage-providers/core/index.ts | 10 + .../github-app-credentials.test.ts | 83 +++++++ .../github-app-credentials.ts | 23 ++ lambdas/libs/storage-providers/index.ts | 3 + .../libs/storage-providers/vitest.config.ts | 10 +- lambdas/yarn.lock | 1 - 17 files changed, 486 insertions(+), 346 deletions(-) create mode 100644 lambdas/libs/storage-providers/aws/ssm/github-app-credentials-store.test.ts create mode 100644 lambdas/libs/storage-providers/aws/ssm/github-app-credentials-store.ts create mode 100644 lambdas/libs/storage-providers/github-app-credentials.test.ts create mode 100644 lambdas/libs/storage-providers/github-app-credentials.ts diff --git a/lambdas/functions/control-plane/package.json b/lambdas/functions/control-plane/package.json index 0f443fc849..ee3aedd210 100644 --- a/lambdas/functions/control-plane/package.json +++ b/lambdas/functions/control-plane/package.json @@ -31,7 +31,6 @@ }, "dependencies": { "@aws-github-runner/aws-powertools-util": "*", - "@aws-github-runner/aws-ssm-util": "*", "@aws-github-runner/compute-providers": "*", "@aws-github-runner/storage-providers": "*", "@aws-lambda-powertools/parameters": "^2.31.0", diff --git a/lambdas/functions/control-plane/src/github/auth.test.ts b/lambdas/functions/control-plane/src/github/auth.test.ts index dd2cf3b8c2..f87053819e 100644 --- a/lambdas/functions/control-plane/src/github/auth.test.ts +++ b/lambdas/functions/control-plane/src/github/auth.test.ts @@ -2,13 +2,19 @@ import { createAppAuth } from '@octokit/auth-app'; import { StrategyOptions } from '@octokit/auth-app/dist-types/types'; import { request } from '@octokit/request'; import { RequestInterface, RequestParameters } from '@octokit/types'; -import { getParameters } from '@aws-github-runner/aws-ssm-util'; +import { + getGitHubAppCredentialsStore, + type GitHubAppCredential, + type GitHubAppCredentialsStore, +} from '@aws-github-runner/storage-providers'; import { generateKeyPairSync } from 'node:crypto'; import * as nock from 'nock'; import { createGithubAppAuth, createOctokitClient, + getAppCount, + getAppId, getStoredInstallationId, onRateLimit, onSecondaryRateLimit, @@ -25,24 +31,27 @@ type MockProxy = T & { // eslint-disable-next-line @typescript-eslint/no-explicit-any const mock = (implementation?: any): MockProxy => vi.fn(implementation) as any; -vi.mock('@aws-github-runner/aws-ssm-util'); +vi.mock('@aws-github-runner/storage-providers', () => ({ + getGitHubAppCredentialsStore: vi.fn(), +})); vi.mock('@octokit/auth-app'); const cleanEnv = process.env; -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 GITHUB_APP_ID = 1; -const mockedGetParameters = vi.mocked(getParameters); +const mockedGetGitHubAppCredentialsStore = vi.mocked(getGitHubAppCredentialsStore); +const mockCredentialsGet = vi.fn(); +const credentialsStore = { + get: mockCredentialsGet, +} satisfies GitHubAppCredentialsStore; beforeEach(() => { vi.resetModules(); vi.clearAllMocks(); + mockCredentialsGet.mockReset(); 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; + mockedGetGitHubAppCredentialsStore.mockReturnValue(credentialsStore); nock.disableNetConnect(); }); @@ -80,38 +89,18 @@ describe('Test createGithubAppAuth', () => { const authType = 'app'; const token = '123456'; const decryptedValue = 'decryptedValue'; - const b64 = Buffer.from(decryptedValue, 'binary').toString('base64'); - - beforeEach(() => { - process.env.ENVIRONMENT = ENVIRONMENT; - }); - it('Throws early when PARAMETER_GITHUB_APP_ID_NAME is not set', async () => { - delete process.env.PARAMETER_GITHUB_APP_ID_NAME; + it('Propagates errors from the credential store', async () => { + const error = new Error('Unable to load GitHub App credentials'); + mockCredentialsGet.mockRejectedValueOnce(error); - await expect(createGithubAppAuth(installationId)).rejects.toThrow( - 'Environment variable PARAMETER_GITHUB_APP_ID_NAME is not set', - ); - expect(mockedGetParameters).not.toHaveBeenCalled(); - }); - - 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( - 'Environment variable PARAMETER_GITHUB_APP_KEY_BASE64_NAME is not set', - ); - expect(mockedGetParameters).not.toHaveBeenCalled(); + await expect(createGithubAppAuth(installationId)).rejects.toBe(error); + expect(mockCredentialsGet).toHaveBeenCalledOnce(); }); it('Creates auth object with createJwt callback including jti claim', async () => { // Arrange - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - ]), - ); + mockCredentialsGet.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: decryptedValue }]); const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token }); @@ -124,7 +113,7 @@ describe('Test createGithubAppAuth', () => { // Assert expect(mockedCreatAppAuth).toBeCalledTimes(1); const callArgs = mockedCreatAppAuth.mock.calls[0][0] as Record; - expect(callArgs.appId).toBe(parseInt(GITHUB_APP_ID)); + expect(callArgs.appId).toBe(GITHUB_APP_ID); expect(callArgs.createJwt).toBeTypeOf('function'); expect(callArgs).not.toHaveProperty('privateKey'); expect(callArgs.installationId).toBe(installationId); @@ -137,14 +126,7 @@ describe('Test createGithubAppAuth', () => { privateKeyEncoding: { type: 'pkcs8', format: 'pem' }, publicKeyEncoding: { type: 'spki', format: 'pem' }, }); - const b64Key = Buffer.from(privateKey as string).toString('base64'); - - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64Key], - ]), - ); + mockCredentialsGet.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: privateKey as string }]); let capturedCreateJwt: (appId: string | number, timeDifference?: number) => Promise<{ jwt: string }>; mockedCreatAppAuth.mockImplementation((opts: StrategyOptions) => { @@ -173,41 +155,9 @@ describe('Test createGithubAppAuth', () => { expect(payload).toHaveProperty('iss'); }); - it('Creates auth object with line breaks in SSH key.', async () => { - // Arrange - const b64PrivateKeyWithLineBreaks = Buffer.from(decryptedValue + '\n' + decryptedValue, 'binary').toString( - 'base64', - ); - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64PrivateKeyWithLineBreaks], - ]), - ); - - const mockedAuth = vi.fn(); - mockedAuth.mockResolvedValue({ token }); - const mockWithHook = Object.assign(mockedAuth, { hook: vi.fn() }); - mockedCreatAppAuth.mockReturnValue(mockWithHook); - - // Act - const result = await createGithubAppAuth(installationId); - - // Assert - expect(getParameters).toBeCalledWith([PARAMETER_GITHUB_APP_ID_NAME, PARAMETER_GITHUB_APP_KEY_BASE64_NAME]); - expect(mockedCreatAppAuth).toBeCalledTimes(1); - expect(mockedAuth).toBeCalledWith({ type: authType }); - expect(result.token).toBe(token); - }); - it('Creates auth object for public GitHub', async () => { // Arrange - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - ]), - ); + mockCredentialsGet.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: decryptedValue }]); const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token }); @@ -218,11 +168,9 @@ describe('Test createGithubAppAuth', () => { const result = await createGithubAppAuth(installationId); // Assert - expect(getParameters).toBeCalledWith([PARAMETER_GITHUB_APP_ID_NAME, PARAMETER_GITHUB_APP_KEY_BASE64_NAME]); - expect(mockedCreatAppAuth).toBeCalledTimes(1); const callArgs = mockedCreatAppAuth.mock.calls[0][0] as Record; - expect(callArgs.appId).toBe(parseInt(GITHUB_APP_ID)); + expect(callArgs.appId).toBe(GITHUB_APP_ID); expect(callArgs.createJwt).toBeTypeOf('function'); expect(callArgs.installationId).toBe(installationId); expect(mockedAuth).toBeCalledWith({ type: authType }); @@ -238,12 +186,7 @@ describe('Test createGithubAppAuth', () => { () => mockedRequestInterface as RequestInterface, ); - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - ]), - ); + mockCredentialsGet.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: decryptedValue }]); const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token }); // eslint-disable-next-line @typescript-eslint/no-unused-vars @@ -255,11 +198,9 @@ describe('Test createGithubAppAuth', () => { const result = await createGithubAppAuth(installationId, githubServerUrl); // Assert - expect(getParameters).toBeCalledWith([PARAMETER_GITHUB_APP_ID_NAME, PARAMETER_GITHUB_APP_KEY_BASE64_NAME]); - expect(mockedCreatAppAuth).toBeCalledTimes(1); const callArgs = mockedCreatAppAuth.mock.calls[0][0] as Record; - expect(callArgs.appId).toBe(parseInt(GITHUB_APP_ID)); + expect(callArgs.appId).toBe(GITHUB_APP_ID); expect(callArgs.createJwt).toBeTypeOf('function'); expect(callArgs.installationId).toBe(installationId); expect(callArgs.request).toBeDefined(); @@ -278,12 +219,7 @@ describe('Test createGithubAppAuth', () => { const installationId = undefined; - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - ]), - ); + mockCredentialsGet.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: decryptedValue }]); const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token }); const mockWithHook = Object.assign(mockedAuth, { hook: vi.fn() }); @@ -293,11 +229,9 @@ describe('Test createGithubAppAuth', () => { const result = await createGithubAppAuth(installationId, githubServerUrl); // Assert - expect(getParameters).toBeCalledWith([PARAMETER_GITHUB_APP_ID_NAME, PARAMETER_GITHUB_APP_KEY_BASE64_NAME]); - expect(mockedCreatAppAuth).toBeCalledTimes(1); const callArgs = mockedCreatAppAuth.mock.calls[0][0] as Record; - expect(callArgs.appId).toBe(parseInt(GITHUB_APP_ID)); + expect(callArgs.appId).toBe(GITHUB_APP_ID); expect(callArgs.createJwt).toBeTypeOf('function'); expect(callArgs).not.toHaveProperty('installationId'); expect(callArgs.request).toBeDefined(); @@ -330,98 +264,48 @@ describe('Test throttling retry caps', () => { }); }); -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); - }); - +describe('Test GitHub App credential accessors', () => { 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; - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - [installationIdParam, '12345'], - ]), - ); + mockCredentialsGet.mockResolvedValueOnce([ + { appId: GITHUB_APP_ID, privateKey: 'private-key', installationId: 12345 }, + ]); const result = await getStoredInstallationId(0); expect(result).toBe(12345); }); - it('returns undefined when installation ID param is empty', async () => { - process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME = ''; - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, 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; - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - ]), - ); + it('returns undefined when the credential has no installation ID', async () => { + mockCredentialsGet.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: 'private-key' }]); 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 = ''; - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - ]), - ); + mockCredentialsGet.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: 'private-key' }]); 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}`; - - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [app1IdParam, '1'], - [app1KeyParam, b64], - [app2IdParam, '2'], - [app2KeyParam, b64], - [app2InstallParam, '67890'], - ]), - ); + it('loads multi-app credentials once and exposes values by index', async () => { + const credentials: GitHubAppCredential[] = [ + { appId: 1, privateKey: 'private-key-1' }, + { appId: 2, privateKey: 'private-key-2', installationId: 67890 }, + ]; + mockCredentialsGet.mockResolvedValueOnce(credentials); + + await expect(getAppCount()).resolves.toBe(2); + await expect(getAppId()).resolves.toBe('1'); + await expect(getAppId(1)).resolves.toBe('2'); + await expect(getStoredInstallationId(0)).resolves.toBeUndefined(); + await expect(getStoredInstallationId(1)).resolves.toBe(67890); + expect(mockCredentialsGet).toHaveBeenCalledOnce(); + }); - // Primary app (index 0) has no stored installation ID - const result0 = await getStoredInstallationId(0); - expect(result0).toBeUndefined(); + it('throws a clear error for an out-of-bounds app ID index', async () => { + mockCredentialsGet.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: 'private-key' }]); - // Additional app (index 1) has stored installation ID - const result1 = await getStoredInstallationId(1); - expect(result1).toBe(67890); + await expect(getAppId(99)).rejects.toThrow('GitHub App credential at index 99 not found'); }); }); diff --git a/lambdas/functions/control-plane/src/github/auth.ts b/lambdas/functions/control-plane/src/github/auth.ts index f64ac00b30..e4ade0b38a 100644 --- a/lambdas/functions/control-plane/src/github/auth.ts +++ b/lambdas/functions/control-plane/src/github/auth.ts @@ -22,7 +22,7 @@ import { Octokit } from '@octokit/rest'; import { retry } from '@octokit/plugin-retry'; import { throttling } from '@octokit/plugin-throttling'; import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; -import { getParameters } from '@aws-github-runner/aws-ssm-util'; +import { getGitHubAppCredentialsStore, type GitHubAppCredential } from '@aws-github-runner/storage-providers'; import { EndpointDefaults } from '@octokit/types'; const logger = createChildLogger('gh-auth'); @@ -69,52 +69,10 @@ export function onSecondaryRateLimit( return retryCount < MAX_SECONDARY_RATE_LIMIT_RETRIES; } -interface GitHubAppCredential { - appId: number; - privateKey: string; - installationId?: number; -} - let appCredentialsPromise: Promise | null = null; async function loadAppCredentials(): Promise { - if (!process.env.PARAMETER_GITHUB_APP_ID_NAME) { - throw new Error('Environment variable PARAMETER_GITHUB_APP_ID_NAME is not set'); - } - if (!process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME) { - throw new Error('Environment variable PARAMETER_GITHUB_APP_KEY_BASE64_NAME is not set'); - } - 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`); - } - // Batch fetch all SSM parameters in a single call to reduce API calls - const allParamNames = [...idParams, ...keyParams, ...installationIdParams.filter((p) => p.length > 0)]; - const params = await getParameters(allParamNames); - - const credentials: GitHubAppCredential[] = []; - for (let i = 0; i < idParams.length; i++) { - const appIdValue = params.get(idParams[i]); - if (!appIdValue) { - throw new Error(`Parameter ${idParams[i]} not found`); - } - const appId = parseInt(appIdValue, 10); - const privateKeyBase64 = params.get(keyParams[i]); - if (!privateKeyBase64) { - throw new Error(`Parameter ${keyParams[i]} not found`); - } - // 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 - const privateKey = Buffer.from(privateKeyBase64, 'base64').toString().replace(/\\n/g, '\n'); - const installationIdParam = installationIdParams[i]; - const installationIdValue = - installationIdParam && installationIdParam.length > 0 ? params.get(installationIdParam) : undefined; - const installationId = installationIdValue ? parseInt(installationIdValue, 10) : undefined; - credentials.push({ appId, privateKey, installationId }); - } + const credentials = await getGitHubAppCredentialsStore().get(); logger.info(`Loaded ${credentials.length} GitHub App credential(s)`); return credentials; } @@ -137,6 +95,14 @@ export async function getStoredInstallationId(appIndex: number): Promise { + const credential = (await getAppCredentials())[appIndex]; + if (!credential) { + throw new Error(`GitHub App credential at index ${appIndex} not found`); + } + return credential.appId.toString(); +} + export async function createOctokitClient(token: string, ghesApiUrl = ''): Promise { const CustomOctokit = Octokit.plugin(retry, throttling); const ocktokitOptions: OctokitOptions = { diff --git a/lambdas/functions/control-plane/src/github/rate-limit.test.ts b/lambdas/functions/control-plane/src/github/rate-limit.test.ts index d9d18c5921..93e6d24ba0 100644 --- a/lambdas/functions/control-plane/src/github/rate-limit.test.ts +++ b/lambdas/functions/control-plane/src/github/rate-limit.test.ts @@ -1,50 +1,37 @@ -import { ResponseHeaders } from '@octokit/types'; -import { createSingleMetric } from '@aws-github-runner/aws-powertools-util'; import { MetricUnit } from '@aws-lambda-powertools/metrics'; +import { createSingleMetric } from '@aws-github-runner/aws-powertools-util'; +import type { ResponseHeaders } from '@octokit/types'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { getAppId } from './auth'; import { metricGitHubAppRateLimit } from './rate-limit'; -import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; -import { getParameter } from '@aws-github-runner/aws-ssm-util'; - -process.env.PARAMETER_GITHUB_APP_ID_NAME = 'test'; -vi.mock('@aws-github-runner/aws-ssm-util', async () => { - // Return only what we need without spreading actual - return { - getParameter: vi.fn((name: string) => { - if (name === process.env.PARAMETER_GITHUB_APP_ID_NAME) { - return '1234'; - } else { - return ''; - } - }), - }; -}); -vi.mock('@aws-github-runner/aws-powertools-util', async () => { - // Provide only what's needed without spreading actual - return { - // Mock the logger - logger: { - debug: vi.fn(), - info: vi.fn(), - warn: vi.fn(), - error: vi.fn(), - }, - // eslint-disable-next-line @typescript-eslint/no-unused-vars - createSingleMetric: vi.fn((name: string, unit: string, value: number, dimensions?: Record) => { - return { - addMetadata: vi.fn(), - }; - }), - }; +vi.mock('./auth', () => ({ + getAppId: vi.fn(), +})); + +vi.mock('@aws-github-runner/aws-powertools-util', () => ({ + logger: { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, + createSingleMetric: vi.fn(() => ({ addMetadata: vi.fn() })), +})); + +const cleanEnv = process.env; +const mockedGetAppId = vi.mocked(getAppId); + +beforeEach(() => { + vi.clearAllMocks(); + mockedGetAppId.mockReset(); + mockedGetAppId.mockResolvedValue('1234'); + process.env = { ...cleanEnv }; }); describe('metricGitHubAppRateLimit', () => { - beforeEach(() => { - vi.clearAllMocks(); - }); - - it('should update rate limit metric', async () => { - // set process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT to true + it('updates the rate limit metric', async () => { process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'true'; const headers: ResponseHeaders = { 'x-ratelimit-remaining': '10', @@ -53,13 +40,13 @@ describe('metricGitHubAppRateLimit', () => { await metricGitHubAppRateLimit(headers); + expect(mockedGetAppId).toHaveBeenCalledWith(undefined); expect(createSingleMetric).toHaveBeenCalledWith('GitHubAppRateLimitRemaining', MetricUnit.Count, 10, { AppId: '1234', }); }); - it('should not update rate limit metric', async () => { - // set process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT to false + it('does not update the rate limit metric when disabled', async () => { process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'false'; const headers: ResponseHeaders = { 'x-ratelimit-remaining': '10', @@ -68,107 +55,85 @@ describe('metricGitHubAppRateLimit', () => { await metricGitHubAppRateLimit(headers); + expect(mockedGetAppId).not.toHaveBeenCalled(); expect(createSingleMetric).not.toHaveBeenCalled(); }); - it('should not update rate limit metric if headers are undefined', async () => { - // set process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT to true + it('does not update the rate limit metric if headers are undefined', async () => { process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'true'; await metricGitHubAppRateLimit(undefined as unknown as ResponseHeaders); + expect(mockedGetAppId).not.toHaveBeenCalled(); expect(createSingleMetric).not.toHaveBeenCalled(); }); - it('should cache GitHub App ID and only call getParameter once', async () => { - // Reset modules to clear the appIdPromises Map cache - vi.resetModules(); - const { metricGitHubAppRateLimit: freshMetricFunction } = await import('./rate-limit'); - + it('does not update the metric when the app ID lookup fails', async () => { process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'true'; + mockedGetAppId.mockRejectedValueOnce(new Error('credential store unavailable')); const headers: ResponseHeaders = { 'x-ratelimit-remaining': '10', 'x-ratelimit-limit': '60', }; - const mockGetParameter = vi.mocked(getParameter); - mockGetParameter.mockClear(); + await expect(metricGitHubAppRateLimit(headers)).resolves.not.toThrow(); - await freshMetricFunction(headers); - await freshMetricFunction(headers); - await freshMetricFunction(headers); - - // getParameter should only be called once due to caching (index 0 cached after first call) - expect(mockGetParameter).toHaveBeenCalledTimes(1); - // split(':')[0] of 'test' is still 'test' - expect(mockGetParameter).toHaveBeenCalledWith(process.env.PARAMETER_GITHUB_APP_ID_NAME); + expect(createSingleMetric).not.toHaveBeenCalled(); }); }); describe('metricGitHubAppRateLimit multi-app', () => { - let freshMetricFunction: typeof metricGitHubAppRateLimit; - let mockGetParam: ReturnType; - - beforeEach(async () => { - // Reset modules to get a clean appIdPromises Map for each test - vi.resetModules(); - - process.env.PARAMETER_GITHUB_APP_ID_NAME = 'app0:app1'; + beforeEach(() => { process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'true'; - - mockGetParam = vi.fn((name: string) => { - if (name === 'app0') return Promise.resolve('1234'); - if (name === 'app1') return Promise.resolve('5678'); - return Promise.resolve(''); + mockedGetAppId.mockImplementation(async (appIndex = 0) => { + if (appIndex === 0) return '1234'; + if (appIndex === 1) return '5678'; + throw new Error(`GitHub App credential at index ${appIndex} not found`); }); - - vi.doMock('@aws-github-runner/aws-ssm-util', () => ({ getParameter: mockGetParam })); - vi.doMock('@aws-github-runner/aws-powertools-util', () => ({ - logger: { debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn() }, - createSingleMetric: vi.fn(() => ({ addMetadata: vi.fn() })), - })); - - const mod = await import('./rate-limit'); - freshMetricFunction = mod.metricGitHubAppRateLimit; - }); - - afterEach(() => { - vi.resetModules(); - process.env.PARAMETER_GITHUB_APP_ID_NAME = 'test'; }); - it('should label metric with correct appId for index 0 (primary app)', async () => { - const { createSingleMetric: mockMetric } = await import('@aws-github-runner/aws-powertools-util'); + it('labels the metric with the primary app ID', async () => { const headers: ResponseHeaders = { 'x-ratelimit-remaining': '50', 'x-ratelimit-limit': '5000' }; - await freshMetricFunction(headers, 0); - expect(mockMetric).toHaveBeenCalledWith('GitHubAppRateLimitRemaining', MetricUnit.Count, 50, { AppId: '1234' }); + + await metricGitHubAppRateLimit(headers, 0); + + expect(mockedGetAppId).toHaveBeenCalledWith(0); + expect(createSingleMetric).toHaveBeenCalledWith('GitHubAppRateLimitRemaining', MetricUnit.Count, 50, { + AppId: '1234', + }); }); - it('should label metric with correct appId for index 1 (additional app)', async () => { - const { createSingleMetric: mockMetric } = await import('@aws-github-runner/aws-powertools-util'); + it('labels the metric with an additional app ID', async () => { const headers: ResponseHeaders = { 'x-ratelimit-remaining': '100', 'x-ratelimit-limit': '5000' }; - await freshMetricFunction(headers, 1); - expect(mockMetric).toHaveBeenCalledWith('GitHubAppRateLimitRemaining', MetricUnit.Count, 100, { AppId: '5678' }); + + await metricGitHubAppRateLimit(headers, 1); + + expect(mockedGetAppId).toHaveBeenCalledWith(1); + expect(createSingleMetric).toHaveBeenCalledWith('GitHubAppRateLimitRemaining', MetricUnit.Count, 100, { + AppId: '5678', + }); }); - it('should default to index 0 when no appIndex is passed', async () => { - const { createSingleMetric: mockMetric } = await import('@aws-github-runner/aws-powertools-util'); + it('defaults to the primary app when no app index is passed', async () => { const headers: ResponseHeaders = { 'x-ratelimit-remaining': '75', 'x-ratelimit-limit': '5000' }; - await freshMetricFunction(headers); - expect(mockMetric).toHaveBeenCalledWith('GitHubAppRateLimitRemaining', MetricUnit.Count, 75, { AppId: '1234' }); + + await metricGitHubAppRateLimit(headers); + + expect(mockedGetAppId).toHaveBeenCalledWith(undefined); + expect(createSingleMetric).toHaveBeenCalledWith('GitHubAppRateLimitRemaining', MetricUnit.Count, 75, { + AppId: '1234', + }); }); - it('should cache per index and call getParameter separately for each index', async () => { + it('forwards each app index to the shared credential accessor', async () => { const headers: ResponseHeaders = { 'x-ratelimit-remaining': '10', 'x-ratelimit-limit': '5000' }; - // Two calls with index 1, then one with index 0 - await freshMetricFunction(headers, 1); - await freshMetricFunction(headers, 1); - await freshMetricFunction(headers, 0); + await metricGitHubAppRateLimit(headers, 1); + await metricGitHubAppRateLimit(headers, 1); + await metricGitHubAppRateLimit(headers, 0); - // getParameter should be called exactly once per distinct index - expect(mockGetParam).toHaveBeenCalledTimes(2); - expect(mockGetParam).toHaveBeenCalledWith('app1'); - expect(mockGetParam).toHaveBeenCalledWith('app0'); + expect(mockedGetAppId).toHaveBeenNthCalledWith(1, 1); + expect(mockedGetAppId).toHaveBeenNthCalledWith(2, 1); + expect(mockedGetAppId).toHaveBeenNthCalledWith(3, 0); }); }); diff --git a/lambdas/functions/control-plane/src/github/rate-limit.ts b/lambdas/functions/control-plane/src/github/rate-limit.ts index df2372a255..b5559a5d82 100644 --- a/lambdas/functions/control-plane/src/github/rate-limit.ts +++ b/lambdas/functions/control-plane/src/github/rate-limit.ts @@ -2,22 +2,8 @@ import { ResponseHeaders } from '@octokit/types'; import { createSingleMetric, logger } from '@aws-github-runner/aws-powertools-util'; import { MetricUnit } from '@aws-lambda-powertools/metrics'; import yn from 'yn'; -import { getParameter } from '@aws-github-runner/aws-ssm-util'; -// Cache the app ID per app index to avoid repeated SSM calls across Lambda invocations. -// In multi-app mode PARAMETER_GITHUB_APP_ID_NAME is a ':'-joined list of SSM param names, -// one per app in app-index order; index 0 is the primary app. -const appIdPromises = new Map>(); - -async function getAppId(appIndex = 0): Promise { - let cached = appIdPromises.get(appIndex); - if (!cached) { - const paramName = process.env.PARAMETER_GITHUB_APP_ID_NAME.split(':')[appIndex]; - cached = getParameter(paramName); - appIdPromises.set(appIndex, cached); - } - return cached; -} +import { getAppId } from './auth'; export async function metricGitHubAppRateLimit(headers: ResponseHeaders, appIndex?: number): Promise { try { diff --git a/lambdas/functions/control-plane/src/lambda.test.ts b/lambdas/functions/control-plane/src/lambda.test.ts index f93b4eac49..4c61f2c585 100644 --- a/lambdas/functions/control-plane/src/lambda.test.ts +++ b/lambdas/functions/control-plane/src/lambda.test.ts @@ -66,7 +66,6 @@ vi.mock('./scale-runners/scale-down'); vi.mock('./scale-runners/scale-up'); vi.mock('./scale-runners/job-retry'); vi.mock('@aws-github-runner/aws-powertools-util'); -vi.mock('@aws-github-runner/aws-ssm-util'); vi.mock('@aws-github-runner/storage-providers', () => ({ getRunnerConfigStore: vi.fn(), })); diff --git a/lambdas/functions/control-plane/src/modules.d.ts b/lambdas/functions/control-plane/src/modules.d.ts index d32f8431e0..af537afcba 100644 --- a/lambdas/functions/control-plane/src/modules.d.ts +++ b/lambdas/functions/control-plane/src/modules.d.ts @@ -13,8 +13,6 @@ declare namespace NodeJS { MINIMUM_RUNNING_TIME_IN_MINUTES: string; PARAMETER_GITHUB_APP_CLIENT_ID_NAME: string; PARAMETER_GITHUB_APP_CLIENT_SECRET_NAME: string; - PARAMETER_GITHUB_APP_ID_NAME: string; - PARAMETER_GITHUB_APP_KEY_BASE64_NAME: string; RUNNER_OWNER: string; COMPUTE_PROVIDER_TYPE?: string; SCALE_DOWN_CONFIG: string; 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 83cb77cd5e..16d609d3e7 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 @@ -147,7 +147,6 @@ let expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; function setDefaults() { process.env = { ...cleanEnv }; - process.env.PARAMETER_GITHUB_APP_ID_NAME = 'github-app-id'; process.env.GITHUB_APP_KEY_BASE64 = 'TEST_CERTIFICATE_DATA'; process.env.GITHUB_APP_ID = '1337'; process.env.GITHUB_APP_CLIENT_ID = 'TEST_CLIENT_ID'; diff --git a/lambdas/libs/storage-providers/aws/ssm/environment.d.ts b/lambdas/libs/storage-providers/aws/ssm/environment.d.ts index b3fb63cdbe..ba0e7afda0 100644 --- a/lambdas/libs/storage-providers/aws/ssm/environment.d.ts +++ b/lambdas/libs/storage-providers/aws/ssm/environment.d.ts @@ -3,6 +3,9 @@ export {}; declare global { namespace NodeJS { interface ProcessEnv { + PARAMETER_GITHUB_APP_ID_NAME?: string; + PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME?: string; + PARAMETER_GITHUB_APP_KEY_BASE64_NAME?: string; SSM_CONFIG_PATH?: string; SSM_CLEANUP_CONFIG?: string; SSM_PARAMETER_STORE_TAGS?: string; diff --git a/lambdas/libs/storage-providers/aws/ssm/github-app-credentials-store.test.ts b/lambdas/libs/storage-providers/aws/ssm/github-app-credentials-store.test.ts new file mode 100644 index 0000000000..4ca2e93914 --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/github-app-credentials-store.test.ts @@ -0,0 +1,154 @@ +import { getParameters } from '@aws-github-runner/aws-ssm-util'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { createAwsSsmGitHubAppCredentialsStore } from './github-app-credentials-store'; + +vi.mock('@aws-github-runner/aws-ssm-util', () => ({ + getParameters: vi.fn(), +})); + +const getParametersMock = vi.mocked(getParameters); +const cleanEnv = process.env; +const primaryIdParameter = '/actions-runner/test/github_app_id'; +const primaryKeyParameter = '/actions-runner/test/github_app_key_base64'; + +describe('aws_ssm GitHub App credentials store', () => { + beforeEach(() => { + vi.clearAllMocks(); + process.env = { ...cleanEnv }; + process.env.PARAMETER_GITHUB_APP_ID_NAME = primaryIdParameter; + process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME = primaryKeyParameter; + delete process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME; + }); + + it('batch reads and maps the primary GitHub App credential', async () => { + const privateKey = 'fake-private-key'; + getParametersMock.mockResolvedValue( + new Map([ + [primaryIdParameter, '123'], + [primaryKeyParameter, Buffer.from(privateKey).toString('base64')], + ]), + ); + const store = createAwsSsmGitHubAppCredentialsStore(); + + await expect(store.get()).resolves.toEqual([{ appId: 123, privateKey, installationId: undefined }]); + expect(getParametersMock).toHaveBeenCalledOnce(); + expect(getParametersMock).toHaveBeenCalledWith([primaryIdParameter, primaryKeyParameter]); + }); + + it('preserves multi-app order and optional installation-id slots', async () => { + const additionalIdParameter = '/actions-runner/test/additional_github_app_0_id'; + const additionalKeyParameter = '/actions-runner/test/additional_github_app_0_key_base64'; + const additionalInstallationIdParameter = '/actions-runner/test/additional_github_app_0_installation_id'; + process.env.PARAMETER_GITHUB_APP_ID_NAME = `${primaryIdParameter}:${additionalIdParameter}`; + process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME = `${primaryKeyParameter}:${additionalKeyParameter}`; + process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME = `:${additionalInstallationIdParameter}`; + getParametersMock.mockResolvedValue( + new Map([ + [primaryIdParameter, '123'], + [primaryKeyParameter, Buffer.from('primary-key').toString('base64')], + [additionalIdParameter, '456'], + [additionalKeyParameter, Buffer.from('additional-key').toString('base64')], + [additionalInstallationIdParameter, '789'], + ]), + ); + const store = createAwsSsmGitHubAppCredentialsStore(); + + await expect(store.get()).resolves.toEqual([ + { appId: 123, privateKey: 'primary-key', installationId: undefined }, + { appId: 456, privateKey: 'additional-key', installationId: 789 }, + ]); + expect(getParametersMock).toHaveBeenCalledWith([ + primaryIdParameter, + additionalIdParameter, + primaryKeyParameter, + additionalKeyParameter, + additionalInstallationIdParameter, + ]); + }); + + it('decodes literal newline escapes in a base64 private key', async () => { + getParametersMock.mockResolvedValue( + new Map([ + [primaryIdParameter, '123'], + [primaryKeyParameter, Buffer.from('first-line\\nsecond-line').toString('base64')], + ]), + ); + const store = createAwsSsmGitHubAppCredentialsStore(); + + await expect(store.get()).resolves.toEqual([ + { appId: 123, privateKey: 'first-line\nsecond-line', installationId: undefined }, + ]); + }); + + it('preserves parseInt behavior for stored numeric values', async () => { + const installationIdParameter = '/actions-runner/test/github_app_installation_id'; + process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME = installationIdParameter; + getParametersMock.mockResolvedValue( + new Map([ + [primaryIdParameter, '123app'], + [primaryKeyParameter, Buffer.from('fake-private-key').toString('base64')], + [installationIdParameter, '789installation'], + ]), + ); + const store = createAwsSsmGitHubAppCredentialsStore(); + + await expect(store.get()).resolves.toEqual([{ appId: 123, privateKey: 'fake-private-key', installationId: 789 }]); + }); + + it.each([ + ['PARAMETER_GITHUB_APP_ID_NAME', undefined], + ['PARAMETER_GITHUB_APP_ID_NAME', ''], + ['PARAMETER_GITHUB_APP_KEY_BASE64_NAME', undefined], + ['PARAMETER_GITHUB_APP_KEY_BASE64_NAME', ''], + ] as const)('rejects missing environment value %s=%j before reading', async (name, value) => { + setEnvironmentValue(name, value); + const store = createAwsSsmGitHubAppCredentialsStore(); + + await expect(store.get()).rejects.toThrow(`Environment variable ${name} is not set`); + expect(getParametersMock).not.toHaveBeenCalled(); + }); + + it('rejects mismatched GitHub App id and key parameter counts before reading', async () => { + process.env.PARAMETER_GITHUB_APP_ID_NAME = `${primaryIdParameter}:/additional/id`; + const store = createAwsSsmGitHubAppCredentialsStore(); + + await expect(store.get()).rejects.toThrow('GitHub App parameter count mismatch: 2 IDs vs 1 keys'); + expect(getParametersMock).not.toHaveBeenCalled(); + }); + + it('rejects a missing GitHub App id parameter', async () => { + getParametersMock.mockResolvedValue( + new Map([[primaryKeyParameter, Buffer.from('fake-private-key').toString('base64')]]), + ); + const store = createAwsSsmGitHubAppCredentialsStore(); + + await expect(store.get()).rejects.toThrow(`Parameter ${primaryIdParameter} not found`); + }); + + it('rejects a missing GitHub App private-key parameter', async () => { + getParametersMock.mockResolvedValue(new Map([[primaryIdParameter, '123']])); + const store = createAwsSsmGitHubAppCredentialsStore(); + + await expect(store.get()).rejects.toThrow(`Parameter ${primaryKeyParameter} not found`); + }); + + it('propagates parameter-store read errors', async () => { + const error = new Error('access denied'); + getParametersMock.mockRejectedValue(error); + const store = createAwsSsmGitHubAppCredentialsStore(); + + await expect(store.get()).rejects.toBe(error); + }); +}); + +function setEnvironmentValue( + name: 'PARAMETER_GITHUB_APP_ID_NAME' | 'PARAMETER_GITHUB_APP_KEY_BASE64_NAME', + value: string | undefined, +): void { + if (value === undefined) { + delete process.env[name]; + } else { + process.env[name] = value; + } +} diff --git a/lambdas/libs/storage-providers/aws/ssm/github-app-credentials-store.ts b/lambdas/libs/storage-providers/aws/ssm/github-app-credentials-store.ts new file mode 100644 index 0000000000..5e5ca2e501 --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/github-app-credentials-store.ts @@ -0,0 +1,61 @@ +import { getParameters } from '@aws-github-runner/aws-ssm-util'; + +import type { GitHubAppCredential, GitHubAppCredentialsStore } from '../../core'; +import type {} from './environment'; + +export function createAwsSsmGitHubAppCredentialsStore(): GitHubAppCredentialsStore { + return new AwsSsmGitHubAppCredentialsStore(); +} + +class AwsSsmGitHubAppCredentialsStore implements GitHubAppCredentialsStore { + async get(): Promise { + if (!process.env.PARAMETER_GITHUB_APP_ID_NAME) { + throw new Error('Environment variable PARAMETER_GITHUB_APP_ID_NAME is not set'); + } + if (!process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME) { + throw new Error('Environment variable PARAMETER_GITHUB_APP_KEY_BASE64_NAME is not set'); + } + + const idParameters = process.env.PARAMETER_GITHUB_APP_ID_NAME.split(':').filter(Boolean); + const keyParameters = process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME.split(':').filter(Boolean); + const installationIdParameters = (process.env.PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME || '').split(':'); + if (idParameters.length !== keyParameters.length) { + throw new Error( + `GitHub App parameter count mismatch: ${idParameters.length} IDs vs ${keyParameters.length} keys`, + ); + } + + const parameterNames = [ + ...idParameters, + ...keyParameters, + ...installationIdParameters.filter((parameter) => parameter.length > 0), + ]; + const parameters = await getParameters(parameterNames); + + const credentials: GitHubAppCredential[] = []; + for (let index = 0; index < idParameters.length; index++) { + const appIdValue = parameters.get(idParameters[index]); + if (!appIdValue) { + throw new Error(`Parameter ${idParameters[index]} not found`); + } + + const privateKeyBase64 = parameters.get(keyParameters[index]); + if (!privateKeyBase64) { + throw new Error(`Parameter ${keyParameters[index]} not found`); + } + + const installationIdParameter = installationIdParameters[index]; + const installationIdValue = installationIdParameter ? parameters.get(installationIdParameter) : undefined; + + credentials.push({ + appId: parseInt(appIdValue, 10), + // Match the GitHub Terraform provider's handling of keys stored as a + // single-line base64 value containing literal newline escapes. + privateKey: Buffer.from(privateKeyBase64, 'base64').toString().replace(/\\n/g, '\n'), + installationId: installationIdValue ? parseInt(installationIdValue, 10) : undefined, + }); + } + + return credentials; + } +} diff --git a/lambdas/libs/storage-providers/core/index.ts b/lambdas/libs/storage-providers/core/index.ts index 6fba06f035..a044489348 100644 --- a/lambdas/libs/storage-providers/core/index.ts +++ b/lambdas/libs/storage-providers/core/index.ts @@ -1,3 +1,13 @@ +export interface GitHubAppCredential { + appId: number; + privateKey: string; + installationId?: number; +} + +export interface GitHubAppCredentialsStore { + get(): Promise; +} + export interface RunnerConfigMetadata { key: string; value: string; diff --git a/lambdas/libs/storage-providers/github-app-credentials.test.ts b/lambdas/libs/storage-providers/github-app-credentials.test.ts new file mode 100644 index 0000000000..fe1870e387 --- /dev/null +++ b/lambdas/libs/storage-providers/github-app-credentials.test.ts @@ -0,0 +1,83 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { createAwsSsmGitHubAppCredentialsStore } from './aws/ssm/github-app-credentials-store'; +import type { GitHubAppCredentialsStore } from './core'; +import { getGitHubAppCredentialsStore, resetGitHubAppCredentialsStore } from './github-app-credentials'; + +vi.mock('./aws/ssm/github-app-credentials-store', () => ({ + createAwsSsmGitHubAppCredentialsStore: vi.fn(), +})); + +const createAwsSsmGitHubAppCredentialsStoreMock = vi.mocked(createAwsSsmGitHubAppCredentialsStore); +const cleanEnv = process.env; + +describe('GitHub App credentials store selection', () => { + beforeEach(() => { + vi.clearAllMocks(); + process.env = { ...cleanEnv }; + delete process.env.RUNNER_CONFIG_STORAGE_PROVIDER; + resetGitHubAppCredentialsStore(); + }); + + it.each([undefined, '', ' '])('uses aws_ssm for default selector input %j', (provider) => { + setProvider(provider); + const store = stubStore(); + + expect(getGitHubAppCredentialsStore()).toBe(store); + expect(createAwsSsmGitHubAppCredentialsStoreMock).toHaveBeenCalledOnce(); + }); + + it.each(['aws_ssm', ' AWS_SSM '])('uses aws_ssm for explicit selector input %j', (provider) => { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = provider; + const store = stubStore(); + + expect(getGitHubAppCredentialsStore()).toBe(store); + expect(createAwsSsmGitHubAppCredentialsStoreMock).toHaveBeenCalledOnce(); + }); + + it('rejects an unsupported provider on first use', () => { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'not-registered'; + + expect(() => getGitHubAppCredentialsStore()).toThrow("Unsupported runner config storage provider 'not-registered'"); + expect(createAwsSsmGitHubAppCredentialsStoreMock).not.toHaveBeenCalled(); + }); + + it('selects lazily and caches the created store', () => { + const store = stubStore(); + + expect(createAwsSsmGitHubAppCredentialsStoreMock).not.toHaveBeenCalled(); + const first = getGitHubAppCredentialsStore(); + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = 'not-registered'; + const second = getGitHubAppCredentialsStore(); + + expect(first).toBe(store); + expect(second).toBe(store); + expect(createAwsSsmGitHubAppCredentialsStoreMock).toHaveBeenCalledOnce(); + }); + + it('selects again after the test reset', () => { + const firstStore = stubStore(); + expect(getGitHubAppCredentialsStore()).toBe(firstStore); + + const secondStore = { get: vi.fn() } satisfies GitHubAppCredentialsStore; + createAwsSsmGitHubAppCredentialsStoreMock.mockReturnValue(secondStore); + resetGitHubAppCredentialsStore(); + + expect(getGitHubAppCredentialsStore()).toBe(secondStore); + expect(createAwsSsmGitHubAppCredentialsStoreMock).toHaveBeenCalledTimes(2); + }); +}); + +function setProvider(provider: string | undefined): void { + if (provider === undefined) { + delete process.env.RUNNER_CONFIG_STORAGE_PROVIDER; + } else { + process.env.RUNNER_CONFIG_STORAGE_PROVIDER = provider; + } +} + +function stubStore(): GitHubAppCredentialsStore { + const store = { get: vi.fn() } satisfies GitHubAppCredentialsStore; + createAwsSsmGitHubAppCredentialsStoreMock.mockReturnValue(store); + return store; +} diff --git a/lambdas/libs/storage-providers/github-app-credentials.ts b/lambdas/libs/storage-providers/github-app-credentials.ts new file mode 100644 index 0000000000..683ab6bb3e --- /dev/null +++ b/lambdas/libs/storage-providers/github-app-credentials.ts @@ -0,0 +1,23 @@ +import { createAwsSsmGitHubAppCredentialsStore } from './aws/ssm/github-app-credentials-store'; +import type { GitHubAppCredentialsStore } from './core'; +import type {} from './environment'; +import { resolveRunnerConfigStorageProvider, type RunnerConfigStorageProvider } from './provider'; + +type GitHubAppCredentialsStoreFactory = () => GitHubAppCredentialsStore; + +const providerFactories = { + aws_ssm: createAwsSsmGitHubAppCredentialsStore, +} as const satisfies Record; + +let githubAppCredentialsStore: GitHubAppCredentialsStore | undefined; + +export function getGitHubAppCredentialsStore(): GitHubAppCredentialsStore { + githubAppCredentialsStore ??= + providerFactories[resolveRunnerConfigStorageProvider(process.env.RUNNER_CONFIG_STORAGE_PROVIDER)](); + return githubAppCredentialsStore; +} + +// Test-only reset for cases that need to exercise first-use environment selection. +export function resetGitHubAppCredentialsStore(): void { + githubAppCredentialsStore = undefined; +} diff --git a/lambdas/libs/storage-providers/index.ts b/lambdas/libs/storage-providers/index.ts index 8001457df5..05a1d5e285 100644 --- a/lambdas/libs/storage-providers/index.ts +++ b/lambdas/libs/storage-providers/index.ts @@ -1,9 +1,12 @@ export type { + GitHubAppCredential, + GitHubAppCredentialsStore, RunnerConfigMetadata, RunnerConfigRecord, RunnerConfigStore, RunnerGroupCacheRecord, RunnerGroupCacheStore, } from './core'; +export { getGitHubAppCredentialsStore, resetGitHubAppCredentialsStore } from './github-app-credentials'; export { getRunnerConfigStore, resetRunnerConfigStore } from './runner-config'; export { getRunnerGroupCacheStore, resetRunnerGroupCacheStore } from './runner-group-cache'; diff --git a/lambdas/libs/storage-providers/vitest.config.ts b/lambdas/libs/storage-providers/vitest.config.ts index af85b8946d..d43a3721eb 100644 --- a/lambdas/libs/storage-providers/vitest.config.ts +++ b/lambdas/libs/storage-providers/vitest.config.ts @@ -7,7 +7,15 @@ export default mergeConfig(defaultConfig, { test: { setupFiles: [resolve(__dirname, '../../aws-vitest-setup.ts')], coverage: { - include: ['index.ts', 'provider.ts', 'runner-config.ts', 'runner-group-cache.ts', 'core/**/*.ts', 'aws/**/*.ts'], + include: [ + 'index.ts', + 'provider.ts', + 'github-app-credentials.ts', + 'runner-config.ts', + 'runner-group-cache.ts', + 'core/**/*.ts', + 'aws/**/*.ts', + ], exclude: ['**/*.test.ts', '**/*.d.ts'], }, }, diff --git a/lambdas/yarn.lock b/lambdas/yarn.lock index 836224f666..b4502d20fc 100644 --- a/lambdas/yarn.lock +++ b/lambdas/yarn.lock @@ -162,7 +162,6 @@ __metadata: resolution: "@aws-github-runner/control-plane@workspace:functions/control-plane" dependencies: "@aws-github-runner/aws-powertools-util": "npm:*" - "@aws-github-runner/aws-ssm-util": "npm:*" "@aws-github-runner/compute-providers": "npm:*" "@aws-github-runner/storage-providers": "npm:*" "@aws-lambda-powertools/parameters": "npm:^2.31.0"