From 8db70dcefdd2f07d390da7105a3e18946c026b80 Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Fri, 18 Sep 2026 23:10:24 +0200 Subject: [PATCH 1/2] refactor(storage): route lambda access through storage providers --- lambdas/functions/control-plane/package.json | 1 - .../control-plane/src/github/auth.test.ts | 251 +++--------- .../control-plane/src/lambda.test.ts | 9 +- .../src/local-ssm-housekeeper.ts | 15 - .../functions/control-plane/src/modules.d.ts | 20 - .../src/scale-runners/scale-up.test.ts | 362 ++++++++---------- .../src/scale-runners/ssm-housekeeper.ts | 6 - .../termination-watcher/package.json | 2 +- .../src/deregister.test.ts | 27 +- .../termination-watcher/src/deregister.ts | 28 +- .../aws/ec2/src/environment.d.ts | 1 + .../aws/ssm/environment.d.ts | 1 + .../ssm/runner-config-housekeeper.test.ts} | 2 +- lambdas/yarn.lock | 3 +- 14 files changed, 248 insertions(+), 480 deletions(-) delete mode 100644 lambdas/functions/control-plane/src/local-ssm-housekeeper.ts delete mode 100644 lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts rename lambdas/{functions/control-plane/src/scale-runners/ssm-housekeeper.test.ts => libs/storage-providers/aws/ssm/runner-config-housekeeper.test.ts} (98%) 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 3010e18abb..35ab43952b 100644 --- a/lambdas/functions/control-plane/src/github/auth.test.ts +++ b/lambdas/functions/control-plane/src/github/auth.test.ts @@ -2,7 +2,7 @@ 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 { getParameter, getParameters } from '@aws-github-runner/aws-ssm-util'; +import { createCommonStorage, type GitHubAppCredentialsStore } from '@aws-github-runner/storage-providers'; import { generateKeyPairSync } from 'node:crypto'; import * as nock from 'nock'; @@ -27,25 +27,25 @@ 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', () => ({ + createCommonStorage: 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 mockedGetParameters = vi.mocked(getParameters); -const mockedGetParameter = vi.mocked(getParameter); +const GITHUB_APP_ID = 1; +const mockedCreateCommonStorage = vi.mocked(createCommonStorage); +const mockedGetCredentials = vi.fn(); +const credentialsStore = { get: mockedGetCredentials } satisfies GitHubAppCredentialsStore; +const defaultCredentials = [{ appId: GITHUB_APP_ID, privateKey: 'private-key' }]; beforeEach(() => { vi.resetModules(); vi.clearAllMocks(); resetAppCredentialsCache(); process.env = { ...cleanEnv }; - process.env.PARAMETER_GITHUB_APP_ID_NAME = PARAMETER_GITHUB_APP_ID_NAME; - process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME = PARAMETER_GITHUB_APP_KEY_BASE64_NAME; + mockedGetCredentials.mockResolvedValue(defaultCredentials); + mockedCreateCommonStorage.mockReturnValue({ githubAppCredentials: credentialsStore }); nock.disableNetConnect(); }); @@ -83,38 +83,10 @@ 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; - - 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(); - }); 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], - ]), - ); + mockedGetCredentials.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: decryptedValue }]); const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token }); @@ -127,7 +99,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); @@ -140,14 +112,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], - ]), - ); + mockedGetCredentials.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: privateKey as string }]); let capturedCreateJwt: (appId: string | number, timeDifference?: number) => Promise<{ jwt: string }>; mockedCreatAppAuth.mockImplementation((opts: StrategyOptions) => { @@ -178,15 +143,8 @@ describe('Test createGithubAppAuth', () => { 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 privateKeyWithLineBreaks = decryptedValue + '\n' + decryptedValue; + mockedGetCredentials.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: privateKeyWithLineBreaks }]); const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token }); @@ -197,7 +155,6 @@ 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); expect(mockedAuth).toBeCalledWith({ type: authType }); expect(result.token).toBe(token); @@ -205,12 +162,7 @@ describe('Test createGithubAppAuth', () => { 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], - ]), - ); + mockedGetCredentials.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: decryptedValue }]); const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token }); @@ -221,11 +173,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 }); @@ -241,12 +191,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], - ]), - ); + mockedGetCredentials.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: decryptedValue }]); const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token }); // eslint-disable-next-line @typescript-eslint/no-unused-vars @@ -258,11 +203,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(); @@ -281,12 +224,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], - ]), - ); + mockedGetCredentials.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: decryptedValue }]); const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token }); const mockWithHook = Object.assign(mockedAuth, { hook: vi.fn() }); @@ -296,11 +234,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(); @@ -334,132 +270,45 @@ describe('Test throttling retry caps', () => { }); describe('Test getStoredInstallationId', () => { - const decryptedValue = 'decryptedValue'; - const b64 = Buffer.from(decryptedValue, 'binary').toString('base64'); + it('returns stored installation ID for an additional app', async () => { + mockedGetCredentials.mockResolvedValueOnce([ + { appId: GITHUB_APP_ID, privateKey: 'private-key' }, + { appId: 2, privateKey: 'additional-private-key', installationId: 12345 }, + ]); - beforeEach(() => { - const mockedAuth = vi.fn(); - mockedAuth.mockResolvedValue({ token: 'token' }); - const mockWithHook = Object.assign(mockedAuth, { hook: vi.fn() }); - vi.mocked(createAppAuth).mockReturnValue(mockWithHook); + await expect(getStoredInstallationId(1)).resolves.toBe(12345); }); - it('returns stored installation ID when configured for an additional app', async () => { - const appIdParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_id`; - const appKeyParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_key_base64`; - const installationIdParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_installation_id`; - process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME = `/actions-runner/${ENVIRONMENT}/additional_github_apps_manifest`; - mockedGetParameter.mockResolvedValueOnce( - JSON.stringify([ - { idParamName: appIdParam, keyParamName: appKeyParam, installationIdParamName: installationIdParam }, - ]), - ); - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - [appIdParam, '2'], - [appKeyParam, b64], - [installationIdParam, '12345'], - ]), - ); - - const result = await getStoredInstallationId(1); - expect(result).toBe(12345); + it('returns undefined when a credential has no stored installation ID', async () => { + await expect(getStoredInstallationId(0)).resolves.toBeUndefined(); }); - it('returns undefined when the manifest env var is empty', async () => { - process.env.PARAMETER_GITHUB_APPS_MANIFEST_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 for an out-of-bounds app index', async () => { + await expect(getStoredInstallationId(99)).resolves.toBeUndefined(); }); - it('returns undefined when the manifest env var is not set', async () => { - delete process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME; - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - ]), - ); + it('loads installation IDs for multiple credentials in order', async () => { + mockedGetCredentials.mockResolvedValueOnce([ + { appId: GITHUB_APP_ID, privateKey: 'private-key' }, + { appId: 2, privateKey: 'additional-private-key', installationId: 67890 }, + ]); - const result = await getStoredInstallationId(0); - expect(result).toBeUndefined(); - }); - - it('returns undefined for out-of-bounds appIndex', async () => { - delete process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME; - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - ]), - ); - - const result = await getStoredInstallationId(99); - expect(result).toBeUndefined(); - }); - - it('loads installation IDs for multi-app setup from the manifest', async () => { - const app2IdParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_id`; - 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_APPS_MANIFEST_NAME = `/actions-runner/${ENVIRONMENT}/additional_github_apps_manifest`; - mockedGetParameter.mockResolvedValueOnce( - JSON.stringify([ - { idParamName: app2IdParam, keyParamName: app2KeyParam, installationIdParamName: app2InstallParam }, - ]), - ); - mockedGetParameters.mockResolvedValueOnce( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, '1'], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - [app2IdParam, '2'], - [app2KeyParam, b64], - [app2InstallParam, '67890'], - ]), - ); - - // Primary app (index 0) has no stored installation ID - const result0 = await getStoredInstallationId(0); - expect(result0).toBeUndefined(); - - // Additional app (index 1) has stored installation ID - const result1 = await getStoredInstallationId(1); - expect(result1).toBe(67890); + await expect(getStoredInstallationId(0)).resolves.toBeUndefined(); + await expect(getStoredInstallationId(1)).resolves.toBe(67890); }); }); describe('Test rate-limit aware app selection', () => { - const decryptedValue = 'decryptedValue'; - const b64 = Buffer.from(decryptedValue, 'binary').toString('base64'); - const app2IdParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_id`; - const app2KeyParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_key_base64`; - beforeEach(() => { const mockedAuth = vi.fn(); mockedAuth.mockResolvedValue({ token: 'token' }); const mockWithHook = Object.assign(mockedAuth, { hook: vi.fn() }); vi.mocked(createAppAuth).mockReturnValue(mockWithHook); - process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME = `/actions-runner/${ENVIRONMENT}/additional_github_apps_manifest`; - mockedGetParameter.mockResolvedValue(JSON.stringify([{ idParamName: app2IdParam, keyParamName: app2KeyParam }])); - mockedGetParameters.mockResolvedValue( - new Map([ - [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], - [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], - [app2IdParam, '2'], - [app2KeyParam, b64], - ]), - ); + mockedGetCredentials.mockResolvedValue([ + { appId: GITHUB_APP_ID, privateKey: 'private-key' }, + { appId: 2, privateKey: 'additional-private-key' }, + ]); // Pin the random start offset to 0 so selection is deterministic. vi.spyOn(Math, 'random').mockReturnValue(0); @@ -473,8 +322,8 @@ describe('Test rate-limit aware app selection', () => { expect(result.appIndex).toBe(1); }); - it('selects from the supplied credentials store without reading SSM', async () => { - const credentialsStore = { + it('selects from the supplied credentials store without reading the default store', async () => { + const suppliedCredentialsStore = { get: vi.fn().mockResolvedValue([ { appId: 10, privateKey: 'first-key' }, { appId: 20, privateKey: 'second-key' }, @@ -483,17 +332,15 @@ describe('Test rate-limit aware app selection', () => { reportAppRateLimit(0, 100); reportAppRateLimit(1, 5000); - const result = await createGithubAppAuth(undefined, '', undefined, credentialsStore); + const result = await createGithubAppAuth(undefined, '', undefined, suppliedCredentialsStore); expect(result.appIndex).toBe(1); expect(createAppAuth).toHaveBeenCalledWith(expect.objectContaining({ appId: 20, createJwt: expect.any(Function) })); - expect(mockedGetParameter).not.toHaveBeenCalled(); - expect(mockedGetParameters).not.toHaveBeenCalled(); + expect(mockedGetCredentials).not.toHaveBeenCalled(); }); it('assumes full budget for apps without observed state', async () => { reportAppRateLimit(0, 100); - // App 1 has no observed state and is assumed full. const result = await createGithubAppAuth(undefined); expect(result.appIndex).toBe(1); @@ -519,14 +366,14 @@ describe('Test rate-limit aware app selection', () => { }); it('short-circuits to the primary app in single-app deployments', async () => { - delete process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME; + mockedGetCredentials.mockResolvedValueOnce([{ appId: GITHUB_APP_ID, privateKey: 'private-key' }]); reportAppRateLimit(0, 0); const result = await createGithubAppAuth(undefined); expect(result.appIndex).toBe(0); }); - it('respects an explicitly provided appIndex', async () => { + it('respects an explicitly provided app index', async () => { reportAppRateLimit(0, 5000); reportAppRateLimit(1, 100); diff --git a/lambdas/functions/control-plane/src/lambda.test.ts b/lambdas/functions/control-plane/src/lambda.test.ts index b4540a9934..07567dcb87 100644 --- a/lambdas/functions/control-plane/src/lambda.test.ts +++ b/lambdas/functions/control-plane/src/lambda.test.ts @@ -2,7 +2,7 @@ import { captureLambdaHandler, logger } from '@aws-github-runner/aws-powertools- import { createRunnerConfigHousekeeper } from '@aws-github-runner/storage-providers'; import { Context, SQSEvent, SQSRecord } from 'aws-lambda'; -import { addMiddleware, adjustPool, scaleDownHandler, scaleUpHandler, ssmHousekeeper, jobRetryCheck } from './lambda'; +import { addMiddleware, adjustPool, scaleDownHandler, scaleUpHandler, runnerConfigHousekeeper, jobRetryCheck } from './lambda'; import { adjust } from './pool/pool'; import { scaleDown } from './scale-runners/scale-down'; import { scaleUp } from './scale-runners/scale-up'; @@ -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', () => ({ createRunnerConfigHousekeeper: vi.fn(), })); @@ -300,18 +299,18 @@ describe('Test middleware', () => { }); }); -describe('Test ssm housekeeper lambda wrapper.', () => { +describe('Test runnerConfigHousekeeper lambda wrapper.', () => { it('Invoke without errors.', async () => { const houseKeeper = vi.fn().mockResolvedValue(); mockedCreateRunnerConfigHousekeeper.mockReturnValue({ houseKeeper }); - await expect(ssmHousekeeper({}, context)).resolves.not.toThrow(); + await expect(runnerConfigHousekeeper({}, context)).resolves.not.toThrow(); expect(houseKeeper).toHaveBeenCalledOnce(); }); it('Errors not throws.', async () => { mockedCreateRunnerConfigHousekeeper.mockReturnValue({ houseKeeper: vi.fn().mockRejectedValue(new Error()) }); - await expect(ssmHousekeeper({}, context)).resolves.not.toThrow(); + await expect(runnerConfigHousekeeper({}, context)).resolves.not.toThrow(); }); }); diff --git a/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts b/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts deleted file mode 100644 index 81c4cbafd5..0000000000 --- a/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts +++ /dev/null @@ -1,15 +0,0 @@ -import { cleanSSMTokens } from '@aws-github-runner/storage-providers/aws/ssm/runner-config-housekeeper'; - -export function run(): void { - cleanSSMTokens({ - dryRun: true, - minimumDaysOld: 3, - tokenPath: '/ghr/my-env/runners/tokens', - }) - .then() - .catch((e) => { - console.log(e); - }); -} - -run(); diff --git a/lambdas/functions/control-plane/src/modules.d.ts b/lambdas/functions/control-plane/src/modules.d.ts index bc04aff0cb..08183fa3ea 100644 --- a/lambdas/functions/control-plane/src/modules.d.ts +++ b/lambdas/functions/control-plane/src/modules.d.ts @@ -2,35 +2,15 @@ declare namespace NodeJS { export interface ProcessEnv { AWS_REGION: string; ENABLE_METRIC_GITHUB_APP_RATE_LIMIT: string; - ENABLE_ON_DEMAND_FAILOVER_FOR_ERRORS: string; - SCALE_ERRORS: string; ENVIRONMENT: string; GHES_URL: string; JOB_RETRY_CONFIG: string; - LAUNCH_TEMPLATE_NAME: string; LOG_LEVEL: 'silly' | 'trace' | 'debug' | 'info' | 'warn' | 'error' | 'fatal'; LOG_TYPE: 'json' | 'pretty' | 'hidden'; MINIMUM_RUNNING_TIME_IN_MINUTES: string; SCALE_DOWN_IDLE_CONFIRMATION_SECONDS?: 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; - PARAMETER_GITHUB_APPS_MANIFEST_NAME?: string; 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'; - INSTANCE_MAX_SPOT_PRICE: string | undefined; - INSTANCE_ALLOCATION_STRATEGY: - | 'lowest-price' - | 'price-capacity-optimized' - | 'diversified' - | 'capacity-optimized' - | 'capacity-optimized-prioritized' - | 'prioritized'; } } 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 664bac60fb..a8ef79c3a9 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,7 +1,3 @@ -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 nock from 'nock'; import { performance } from 'perf_hooks'; @@ -18,10 +14,29 @@ import type { } from './types'; import { InvalidRunnerLabelsError } from '@aws-github-runner/compute-providers/core'; import { defaultComputeProvider } from '@aws-github-runner/compute-providers/provider-types'; -import { getParameter } from '@aws-github-runner/aws-ssm-util'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { Octokit } from '@octokit/rest'; +const { mockStorage, mockedCreateStorageProviders } = vi.hoisted(() => { + const storage = { + runnerConfig: { + maxWritesPerSecond: 40, + create: vi.fn(), + }, + runnerGroupCache: { + get: vi.fn(), + create: vi.fn(), + }, + consumer: { consume: vi.fn() }, + githubAppCredentials: { get: vi.fn() }, + }; + + return { + mockStorage: storage, + mockedCreateStorageProviders: vi.fn(() => storage), + }; +}); + const mockOctokit = { paginate: vi.fn(), checks: { get: vi.fn() }, @@ -55,8 +70,9 @@ 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 mockRunnerConfigCreate = mockStorage.runnerConfig.create; +const mockRunnerGroupCacheGet = mockStorage.runnerGroupCache.get; +const mockRunnerGroupCacheCreate = mockStorage.runnerGroupCache.create; const mockPublishRetryMessage = vi.mocked(publishRetryMessage); const testProviderState = { provider: 'test' }; const mockComputeProvider = { @@ -87,14 +103,14 @@ 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'); +vi.mock('@aws-github-runner/storage-providers', async () => { + const actual = await vi.importActual( + '@aws-github-runner/storage-providers', + ); return { ...actual, - getParameter: vi.fn(), + createStorageProviders: mockedCreateStorageProviders, }; }); @@ -149,8 +165,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_TOKEN_PATH = '/github-action-runners/default/runners/config'; - process.env.SSM_CONFIG_PATH = '/github-action-runners/default/runners/config'; } async function createTestProviderRunners(input: CreateScaleUpRunnersInput): Promise { @@ -172,6 +186,8 @@ async function createTestProviderRunners(input: CreateScaleUpRunnersInput [{ key: 'RunnerId', value: runnerId }], }, ); @@ -192,7 +208,10 @@ beforeEach(() => { vi.clearAllMocks(); setDefaults(); - defaultSSMGetParameterMockImpl(); + mockedCreateStorageProviders.mockReturnValue(mockStorage); + mockRunnerConfigCreate.mockResolvedValue(undefined); + mockRunnerGroupCacheGet.mockResolvedValue(1); + mockRunnerGroupCacheCreate.mockResolvedValue(undefined); defaultOctokitMockImpl(); mockedResolveCapability.mockReturnValue(() => mockComputeProvider); @@ -272,12 +291,10 @@ 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(); + resetStorageMocks(); }); it('does not create a token when maximum runners has been reached', async () => { @@ -332,11 +349,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 () => { - const error = new Error('ParameterNotFound'); - error.name = 'ParameterNotFound'; - throw error; - }); + mockRunnerGroupCacheGet.mockResolvedValue(undefined); await expect(scaleUpModule.scaleUp(TEST_DATA)).resolves.toEqual(['foobar']); @@ -351,26 +364,20 @@ 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 () => { - const error = new Error('ParameterNotFound'); - error.name = 'ParameterNotFound'; - throw error; - }); + it('creates the runner group cache record when it does not exist', async () => { + mockRunnerGroupCacheGet.mockResolvedValue(undefined); await scaleUpModule.scaleUp(TEST_DATA); 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: process.env.RUNNER_GROUP_NAME, + runnerGroupId: 1, }); }); - it('Does not create SSM parameter for runner group id if it exists', async () => { + it('does not create the runner group cache record if it exists', async () => { await scaleUpModule.scaleUp(TEST_DATA); expect(mockOctokit.paginate).toHaveBeenCalledTimes(0); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 1); + expect(mockRunnerGroupCacheCreate).not.toHaveBeenCalled(); }); it('create start runner config for ephemeral runners ', async () => { @@ -383,17 +390,11 @@ 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).toHaveBeenNthCalledWith( + 1, + { runnerId: 'i-12345', value: 'TEST_JIT_CONFIG_ORG' }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); it('create start runner config for non-ephemeral runners ', async () => { @@ -402,19 +403,16 @@ 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).toHaveBeenNthCalledWith( + 1, + { + 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 () => { @@ -429,19 +427,16 @@ 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).toHaveBeenNthCalledWith( + 1, + { + 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 () => { @@ -501,23 +496,20 @@ 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(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-instance-1', value: 'TEST_JIT_CONFIG_unit-test-i-instance-1' }, + { metadata: [{ 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(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-instance-3', value: 'TEST_JIT_CONFIG_unit-test-i-instance-3' }, + { metadata: [{ key: 'RunnerId', value: 'i-instance-3' }] }, + ); - expect(mockSSMClient).not.toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-2', - }); + expect(mockRunnerConfigCreate).not.toHaveBeenCalledWith( + expect.objectContaining({ runnerId: 'i-instance-2' }), + expect.anything(), + ); }); it('should handle retryable errors with error handling logic', async () => { @@ -553,16 +545,15 @@ 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(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-instance-2', value: 'TEST_JIT_CONFIG_unit-test-i-instance-2' }, + { metadata: [{ key: 'RunnerId', value: 'i-instance-2' }] }, + ); - expect(mockSSMClient).not.toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-1', - }); + expect(mockRunnerConfigCreate).not.toHaveBeenCalledWith( + expect.objectContaining({ runnerId: 'i-instance-1' }), + expect.anything(), + ); }); it('should handle non-retryable 4xx errors gracefully', async () => { @@ -599,20 +590,19 @@ 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(mockRunnerConfigCreate).toHaveBeenCalledWith( + { runnerId: 'i-instance-2', value: 'TEST_JIT_CONFIG_unit-test-i-instance-2' }, + { metadata: [{ key: 'RunnerId', value: 'i-instance-2' }] }, + ); - expect(mockSSMClient).not.toHaveReceivedCommandWith(PutParameterCommand, { - Name: '/github-action-runners/default/runners/config/i-instance-1', - }); + 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 ', + 'calls create start runner config of 40 instances (storage write rate limit) to test time delay', async (type: RunnerLifecycle) => { process.env.ENABLE_EPHEMERAL_RUNNERS = type === 'ephemeral' ? 'true' : 'false'; process.env.RUNNERS_MAXIMUM_COUNT = '40'; @@ -668,7 +658,7 @@ describe('scaleUp with GHES', () => { await scaleUpModule.scaleUp(TEST_DATA); const endTime = performance.now(); expect(endTime - startTime).toBeGreaterThan(1000); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 40); + expect(mockRunnerConfigCreate).toHaveBeenCalledTimes(40); }, 10000, ); @@ -682,7 +672,7 @@ describe('scaleUp with GHES', () => { process.env.RUNNER_LABELS = 'base-label'; process.env.RUNNER_NAME_PREFIX = 'unit-test'; expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; - mockSSMClient.reset(); + resetStorageMocks(); mockResolveLabelsForRunners.mockImplementation(async (labels) => ({ runnerLabels: labels.filter((label) => label.startsWith('ghr-')), @@ -1203,7 +1193,7 @@ describe('scaleUp with public GH', () => { describe('on repo level', () => { beforeEach(() => { - mockSSMClient.reset(); + resetStorageMocks(); process.env.ENABLE_ORGANIZATION_RUNNERS = 'false'; process.env.RUNNER_NAME_PREFIX = 'unit-test'; @@ -1248,44 +1238,33 @@ 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'; - 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).toHaveBeenNthCalledWith( + 1, + { runnerId: 'i-12345', value: 'TEST_JIT_CONFIG_REPO' }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); 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).toHaveBeenNthCalledWith( + 1, + { + 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 () => { @@ -1293,22 +1272,18 @@ describe('scaleUp with public GH', () => { process.env.ENABLE_JIT_CONFIG = 'true'; process.env.ENABLE_JOB_QUEUED_CHECK = 'false'; process.env.RUNNER_LABELS = 'jit'; - 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).toHaveBeenNthCalledWith( + 1, + { + runnerId: 'i-12345', + value: '--url https://github.com/Codertocat/hello-world --token 1234abcd --labels jit', + }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); it('creates a ephemeral runner after checking job is queued.', async () => { @@ -1587,12 +1562,10 @@ 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(); + resetStorageMocks(); }); it('does not create a token when maximum runners has been reached', async () => { @@ -1635,26 +1608,20 @@ 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 () => { - const error = new Error('ParameterNotFound'); - error.name = 'ParameterNotFound'; - throw error; - }); + it('creates the runner group cache record when it does not exist', async () => { + mockRunnerGroupCacheGet.mockResolvedValue(undefined); await scaleUpModule.scaleUp(TEST_DATA); 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: process.env.RUNNER_GROUP_NAME, + runnerGroupId: 1, }); }); - it('Does not create SSM parameter for runner group id if it exists', async () => { + it('does not create the runner group cache record if it exists', async () => { await scaleUpModule.scaleUp(TEST_DATA); expect(mockOctokit.paginate).toHaveBeenCalledTimes(0); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 1); + expect(mockRunnerGroupCacheCreate).not.toHaveBeenCalled(); }); it('create start runner config for ephemeral runners ', async () => { @@ -1667,17 +1634,11 @@ 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).toHaveBeenNthCalledWith( + 1, + { runnerId: 'i-12345', value: 'TEST_JIT_CONFIG_ORG' }, + { metadata: [{ key: 'RunnerId', value: 'i-12345' }] }, + ); }); it('create start runner config for non-ephemeral runners ', async () => { @@ -1686,22 +1647,19 @@ 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).toHaveBeenNthCalledWith( + 1, + { + 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 ', + 'calls create start runner config of 40 instances (storage write rate limit) to test time delay', async (type: RunnerLifecycle) => { process.env.ENABLE_EPHEMERAL_RUNNERS = type === 'ephemeral' ? 'true' : 'false'; process.env.RUNNERS_MAXIMUM_COUNT = '40'; @@ -1757,7 +1715,7 @@ describe('scaleUp with Github Data Residency', () => { await scaleUpModule.scaleUp(TEST_DATA); const endTime = performance.now(); expect(endTime - startTime).toBeGreaterThan(1000); - expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 40); + expect(mockRunnerConfigCreate).toHaveBeenCalledTimes(40); }, 10000, ); @@ -2044,7 +2002,7 @@ describe('Retry mechanism tests', () => { process.env.ENABLE_JOB_QUEUED_CHECK = 'true'; process.env.RUNNERS_MAXIMUM_COUNT = '10'; expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; - mockSSMClient.reset(); + resetStorageMocks(); }); const createTestMessages = ( @@ -2232,11 +2190,9 @@ 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(); + resetStorageMocks(); }); it('passes the same appIndex to createGithubInstallationAuth when multi-app is active', async () => { @@ -2322,7 +2278,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 from storage mockedGetAppCount.mockResolvedValue(2); mockedGetStoredInstallationId.mockResolvedValue(77); mockedAppAuth.mockResolvedValue({ @@ -2393,17 +2349,11 @@ function defaultOctokitMockImpl() { mockOctokit.apps.getRepoInstallation.mockImplementation(() => mockInstallationIdReturnValueRepos); } -function defaultSSMGetParameterMockImpl() { - mockSSMgetParameter.mockImplementation(async (name: string) => { - const runnerGroupName = process.env.RUNNER_GROUP_NAME || 'Default'; - if (name === `${process.env.SSM_CONFIG_PATH}/runner-group/${runnerGroupName}`) { - return '1'; - } else if (name === `${process.env.PARAMETER_GITHUB_APP_ID_NAME}`) { - return `${process.env.GITHUB_APP_ID}`; - } else { - const error = new Error(`ParameterNotFound: ${name}`); - error.name = 'ParameterNotFound'; - throw error; - } - }); +function resetStorageMocks() { + mockRunnerConfigCreate.mockReset(); + mockRunnerConfigCreate.mockResolvedValue(undefined); + mockRunnerGroupCacheGet.mockReset(); + mockRunnerGroupCacheGet.mockResolvedValue(1); + mockRunnerGroupCacheCreate.mockReset(); + mockRunnerGroupCacheCreate.mockResolvedValue(undefined); } diff --git a/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts b/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts deleted file mode 100644 index 52464b720d..0000000000 --- a/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts +++ /dev/null @@ -1,6 +0,0 @@ -/** @deprecated Import the SSM runner-config housekeeper from storage-providers. */ -export { - cleanSSMTokens, - createAwsSsmRunnerConfigHousekeeper, - type SSMCleanupOptions, -} from '@aws-github-runner/storage-providers/aws/ssm/runner-config-housekeeper'; diff --git a/lambdas/functions/termination-watcher/package.json b/lambdas/functions/termination-watcher/package.json index 87622843a9..3c58b5c272 100644 --- a/lambdas/functions/termination-watcher/package.json +++ b/lambdas/functions/termination-watcher/package.json @@ -24,7 +24,7 @@ }, "dependencies": { "@aws-github-runner/aws-powertools-util": "*", - "@aws-github-runner/aws-ssm-util": "*", + "@aws-github-runner/storage-providers": "*", "@aws-sdk/client-ec2": "^3.1009.0", "@aws-sdk/client-sqs": "^3.1009.0", "@middy/core": "^6.4.5", diff --git a/lambdas/functions/termination-watcher/src/deregister.test.ts b/lambdas/functions/termination-watcher/src/deregister.test.ts index dfc2854252..a804e5ce74 100644 --- a/lambdas/functions/termination-watcher/src/deregister.test.ts +++ b/lambdas/functions/termination-watcher/src/deregister.test.ts @@ -1,14 +1,18 @@ import { Instance } from '@aws-sdk/client-ec2'; +import { createCommonStorage, type GitHubAppCredentialsStore } from '@aws-github-runner/storage-providers'; import { describe, it, expect, beforeEach, vi } from 'vitest'; -import { deregisterRunner, createThrottleOptions } from './deregister'; +import { deregisterRunner, createThrottleOptions, resetAppCredentialsCache } from './deregister'; import { Config } from './ConfigResolver'; import type { EndpointDefaults } from '@octokit/types'; -const mockGetParameter = vi.fn(); -vi.mock('@aws-github-runner/aws-ssm-util', () => ({ - getParameter: (...args: unknown[]) => mockGetParameter(...args), +vi.mock('@aws-github-runner/storage-providers', () => ({ + createCommonStorage: vi.fn(), })); +const mockedCreateCommonStorage = vi.mocked(createCommonStorage); +const mockGetCredentials = vi.fn(); +const credentialsStore = { get: mockGetCredentials } satisfies GitHubAppCredentialsStore; + const mockCreateAppAuth = vi.fn(); vi.mock('@octokit/auth-app', () => ({ createAppAuth: (...args: unknown[]) => mockCreateAppAuth(...args), @@ -89,12 +93,7 @@ const repoInstance: Instance = { }; function setupAuthMocks() { - const appPrivateKey = Buffer.from('fake-private-key').toString('base64'); - mockGetParameter.mockImplementation((name: string) => { - if (name === 'github-app-id') return Promise.resolve('12345'); - if (name === 'github-app-key') return Promise.resolve(appPrivateKey); - return Promise.reject(new Error(`Unknown parameter: ${name}`)); - }); + mockGetCredentials.mockResolvedValue([{ appId: 12345, privateKey: 'fake-private-key' }]); // App auth returns app token const mockAuth = vi.fn(); @@ -110,20 +109,20 @@ function setupAuthMocks() { describe('deregisterRunner', () => { beforeEach(() => { vi.clearAllMocks(); - process.env.PARAMETER_GITHUB_APP_ID_NAME = 'github-app-id'; - process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME = 'github-app-key'; + resetAppCredentialsCache(); + mockedCreateCommonStorage.mockReturnValue({ githubAppCredentials: credentialsStore }); setupAuthMocks(); }); it('should skip deregistration when disabled', async () => { await deregisterRunner(orgInstance, { ...baseConfig, enableRunnerDeregistration: false }); - expect(mockGetParameter).not.toHaveBeenCalled(); + expect(mockGetCredentials).not.toHaveBeenCalled(); }); it('should skip deregistration when instance ID is missing', async () => { const instance: Instance = { ...orgInstance, InstanceId: undefined }; await deregisterRunner(instance, baseConfig); - expect(mockGetParameter).not.toHaveBeenCalled(); + expect(mockGetCredentials).not.toHaveBeenCalled(); }); it('should skip deregistration when ghr:Owner tag is missing', async () => { diff --git a/lambdas/functions/termination-watcher/src/deregister.ts b/lambdas/functions/termination-watcher/src/deregister.ts index ea53ad5240..1e811d2c86 100644 --- a/lambdas/functions/termination-watcher/src/deregister.ts +++ b/lambdas/functions/termination-watcher/src/deregister.ts @@ -5,7 +5,7 @@ import { request } from '@octokit/request'; import { Instance } from '@aws-sdk/client-ec2'; import { SQSClient, SendMessageCommand } from '@aws-sdk/client-sqs'; import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; -import { getParameter } from '@aws-github-runner/aws-ssm-util'; +import { createCommonStorage, type GitHubAppCredential } from '@aws-github-runner/storage-providers'; import type { EndpointDefaults } from '@octokit/types'; import type { Config } from './ConfigResolver'; @@ -21,6 +21,8 @@ const sqsClient = new SQSClient({ region: process.env.AWS_REGION }); const logger = createChildLogger('deregister'); +let appCredentialsPromise: Promise | undefined; + export function createThrottleOptions() { return { onRateLimit: (_retryAfter: number, options: Required) => { @@ -34,12 +36,24 @@ export function createThrottleOptions() { }; } -async function getAppCredentials(): Promise<{ appId: number; privateKey: string }> { - const appId = parseInt(await getParameter(process.env.PARAMETER_GITHUB_APP_ID_NAME!)); - const privateKey = Buffer.from(await getParameter(process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME!), 'base64') - .toString() - .replace('/[\\n]/g', String.fromCharCode(10)); - return { appId, privateKey }; +async function loadAppCredentials(): Promise { + const credentials = await createCommonStorage().githubAppCredentials.get(); + const credential = credentials[0]; + if (!credential) { + throw new Error('No GitHub App credentials found'); + } + return credential; +} + +function getAppCredentials(): Promise { + if (!appCredentialsPromise) { + appCredentialsPromise = loadAppCredentials(); + } + return appCredentialsPromise; +} + +export function resetAppCredentialsCache(): void { + appCredentialsPromise = undefined; } function createOctokitInstance(token: string, ghesApiUrl: string): Octokit { diff --git a/lambdas/libs/compute-providers/aws/ec2/src/environment.d.ts b/lambdas/libs/compute-providers/aws/ec2/src/environment.d.ts index 71ee01ff2f..c2d76c32cd 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/environment.d.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/environment.d.ts @@ -17,6 +17,7 @@ declare global { | 'capacity-optimized-prioritized' | 'prioritized'; SCALE_ERRORS: string; + ENABLE_ON_DEMAND_FAILOVER_FOR_ERRORS: string; } } } diff --git a/lambdas/libs/storage-providers/aws/ssm/environment.d.ts b/lambdas/libs/storage-providers/aws/ssm/environment.d.ts index 2604ce057e..985d362bbc 100644 --- a/lambdas/libs/storage-providers/aws/ssm/environment.d.ts +++ b/lambdas/libs/storage-providers/aws/ssm/environment.d.ts @@ -10,6 +10,7 @@ declare global { PARAMETER_GITHUB_APP_KEY_BASE64_NAME?: string; PARAMETER_GITHUB_APP_INSTALLATION_ID_NAME?: string; PARAMETER_GITHUB_APPS_MANIFEST_NAME?: string; + SSM_CLEANUP_CONFIG: string; } } } diff --git a/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.test.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.test.ts similarity index 98% rename from lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.test.ts rename to lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.test.ts index a848526a50..1c16607333 100644 --- a/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.test.ts +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.test.ts @@ -1,7 +1,7 @@ import { DeleteParameterCommand, GetParametersByPathCommand, SSMClient } from '@aws-sdk/client-ssm'; import { mockClient } from 'aws-sdk-client-mock'; import 'aws-sdk-client-mock-jest/vitest'; -import { cleanSSMTokens } from './ssm-housekeeper'; +import { cleanSSMTokens } from './runner-config-housekeeper'; import { describe, it, expect, beforeEach } from 'vitest'; process.env.AWS_REGION = 'eu-east-1'; diff --git a/lambdas/yarn.lock b/lambdas/yarn.lock index a703ceaa7a..65e2f362ba 100644 --- a/lambdas/yarn.lock +++ b/lambdas/yarn.lock @@ -153,7 +153,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" @@ -218,7 +217,7 @@ __metadata: resolution: "@aws-github-runner/termination-watcher@workspace:functions/termination-watcher" dependencies: "@aws-github-runner/aws-powertools-util": "npm:*" - "@aws-github-runner/aws-ssm-util": "npm:*" + "@aws-github-runner/storage-providers": "npm:*" "@aws-sdk/client-ec2": "npm:^3.1009.0" "@aws-sdk/client-sqs": "npm:^3.1009.0" "@aws-sdk/types": "npm:^3.973.6" From 02f0018de4f04ac2714118640e132c54eeb4c77e Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Fri, 18 Sep 2026 23:15:44 +0200 Subject: [PATCH 2/2] style: fix formatting issues --- lambdas/functions/control-plane/src/lambda.test.ts | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/lambdas/functions/control-plane/src/lambda.test.ts b/lambdas/functions/control-plane/src/lambda.test.ts index 07567dcb87..884dc48001 100644 --- a/lambdas/functions/control-plane/src/lambda.test.ts +++ b/lambdas/functions/control-plane/src/lambda.test.ts @@ -2,7 +2,14 @@ import { captureLambdaHandler, logger } from '@aws-github-runner/aws-powertools- import { createRunnerConfigHousekeeper } from '@aws-github-runner/storage-providers'; import { Context, SQSEvent, SQSRecord } from 'aws-lambda'; -import { addMiddleware, adjustPool, scaleDownHandler, scaleUpHandler, runnerConfigHousekeeper, jobRetryCheck } from './lambda'; +import { + addMiddleware, + adjustPool, + scaleDownHandler, + scaleUpHandler, + runnerConfigHousekeeper, + jobRetryCheck, +} from './lambda'; import { adjust } from './pool/pool'; import { scaleDown } from './scale-runners/scale-down'; import { scaleUp } from './scale-runners/scale-up';