diff --git a/lambdas/libs/storage-providers/aws/dynamodb/runner-config-consumer.test.ts b/lambdas/libs/storage-providers/aws/dynamodb/runner-config-consumer.test.ts new file mode 100644 index 0000000000..f3601f7981 --- /dev/null +++ b/lambdas/libs/storage-providers/aws/dynamodb/runner-config-consumer.test.ts @@ -0,0 +1,230 @@ +import { ConditionalCheckFailedException, DeleteItemCommand, type DynamoDBClient } from '@aws-sdk/client-dynamodb'; +import { afterEach, describe, expect, it, vi } from 'vitest'; + +import { + AwsSdkDynamoDbRunnerConfigApi, + createAwsDynamoDbRunnerConfigConsumer, + type AwsDynamoDbRunnerConfigApi, +} from './runner-config-consumer'; + +const dynamoDbEnvironment = { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb', + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'runner-state', +} as const; + +function namedError(name: string, message = 'provider detail'): Error { + const error = new Error(message); + error.name = name; + return error; +} + +describe('AWS SDK DynamoDB runner config API', () => { + afterEach(() => { + vi.useRealTimers(); + }); + + it('atomically deletes an unexpired composite-key item and returns its previous value', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const send = vi.fn().mockResolvedValue({ Attributes: { value: { S: 'encoded-jit' } } }); + const api = new AwsSdkDynamoDbRunnerConfigApi({ send } as unknown as DynamoDBClient); + const signal = new AbortController().signal; + + await expect(api.deleteItem('runner-state', 'microvm-123', signal)).resolves.toBe('encoded-jit'); + + expect(send.mock.calls[0][0]).toBeInstanceOf(DeleteItemCommand); + expect(send.mock.calls[0][0].input).toEqual({ + TableName: 'runner-state', + Key: { + scope: { S: 'microvm-123' }, + id: { S: 'config' }, + }, + ConditionExpression: 'attribute_exists(#expires_at) AND #expires_at > :now', + ExpressionAttributeNames: { '#expires_at': 'expires_at' }, + ExpressionAttributeValues: { ':now': { N: '1767225600' } }, + ReturnValues: 'ALL_OLD', + }); + expect(send.mock.calls[0][1]).toEqual({ abortSignal: signal }); + }); + + it('treats missing or expired records as unavailable after the conditional delete', async () => { + const conditional = new ConditionalCheckFailedException({ + $metadata: {}, + message: 'record is absent or expired', + }); + const api = new AwsSdkDynamoDbRunnerConfigApi({ + send: vi.fn().mockRejectedValue(conditional), + } as unknown as DynamoDBClient); + + await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).resolves.toBeUndefined(); + }); + + it('supports a deserialized conditional error without exposing its message', async () => { + const api = new AwsSdkDynamoDbRunnerConfigApi({ + send: vi.fn().mockRejectedValue(namedError('ConditionalCheckFailedException', 'expired-secret-detail')), + } as unknown as DynamoDBClient); + + await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).resolves.toBeUndefined(); + }); + + it('propagates non-conditional provider errors', async () => { + const error = namedError('AccessDeniedException'); + const api = new AwsSdkDynamoDbRunnerConfigApi({ + send: vi.fn().mockRejectedValue(error), + } as unknown as DynamoDBClient); + + await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).rejects.toBe(error); + }); + + it('returns undefined when no item was deleted', async () => { + const api = new AwsSdkDynamoDbRunnerConfigApi({ + send: vi.fn().mockResolvedValue({}), + } as unknown as DynamoDBClient); + + await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).resolves.toBeUndefined(); + }); + + it.each([{ Attributes: {} }, { Attributes: { value: { N: '1' } } }, { Attributes: { value: { S: '' } } }])( + 'rejects a deleted item without a string value %#', + async (response) => { + const api = new AwsSdkDynamoDbRunnerConfigApi({ + send: vi.fn().mockResolvedValue(response), + } as unknown as DynamoDBClient); + + await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).rejects.toThrow( + 'runner configuration record has an invalid value', + ); + }, + ); +}); + +describe('DynamoDB runner config consumer', () => { + afterEach(() => { + vi.useRealTimers(); + }); + + it('polls until an atomic delete returns the stored configuration', async () => { + const deleteItem = vi + .fn() + .mockResolvedValueOnce(undefined) + .mockResolvedValueOnce('encoded-jit'); + const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, { + api: { deleteItem }, + callTimeoutMs: 100, + configTimeoutMs: 500, + pollIntervalMs: 1, + }); + + await expect( + consumer.consume('microvm-123', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).resolves.toBe('encoded-jit'); + expect(deleteItem).toHaveBeenCalledTimes(2); + expect(deleteItem).toHaveBeenCalledWith('runner-state', 'microvm-123', expect.any(AbortSignal)); + }); + + it('retries transient provider failures', async () => { + const deleteItem = vi + .fn() + .mockRejectedValueOnce(namedError('ProvisionedThroughputExceededException')) + .mockResolvedValueOnce('encoded-jit'); + const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, { + api: { deleteItem }, + callTimeoutMs: 100, + configTimeoutMs: 500, + pollIntervalMs: 1, + }); + + await expect( + consumer.consume('microvm-123', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).resolves.toBe('encoded-jit'); + expect(deleteItem).toHaveBeenCalledTimes(2); + }); + + it('sanitizes non-retryable provider failures', async () => { + const api: AwsDynamoDbRunnerConfigApi = { + deleteItem: vi.fn().mockRejectedValue(namedError('AccessDeniedException', 'encoded-jit-secret')), + }; + const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, { + api, + callTimeoutMs: 100, + configTimeoutMs: 100, + pollIntervalMs: 1, + }); + + const pending = consumer.consume('microvm-123', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }); + await expect(pending).rejects.toThrow('failed to consume runner configuration from DynamoDB'); + await expect(pending).rejects.not.toThrow('encoded-jit-secret'); + }); + + it('validates the complete composite key before calling DynamoDB', async () => { + const api: AwsDynamoDbRunnerConfigApi = { deleteItem: vi.fn() }; + const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, { + api, + callTimeoutMs: 100, + configTimeoutMs: 100, + pollIntervalMs: 1, + }); + + await expect( + consumer.consume('invalid/scope', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toThrow('runnerId is invalid'); + expect(api.deleteItem).not.toHaveBeenCalled(); + }); + + it('times out while missing or expired items remain unavailable', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const api: AwsDynamoDbRunnerConfigApi = { deleteItem: vi.fn().mockResolvedValue(undefined) }; + const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, { + api, + callTimeoutMs: 10, + configTimeoutMs: 20, + pollIntervalMs: 5, + }); + const pending = consumer.consume('microvm-123', { + deadlineMs: Date.now() + 100, + signal: new AbortController().signal, + }); + const rejection = expect(pending).rejects.toThrow( + 'runner configuration did not become available before the deadline', + ); + + await vi.runAllTimersAsync(); + + await rejection; + expect(api.deleteItem).toHaveBeenCalled(); + }); + + it('stops a provider call immediately when the caller aborts', async () => { + const api: AwsDynamoDbRunnerConfigApi = { + deleteItem: vi.fn().mockReturnValue(new Promise(() => undefined)), + }; + const controller = new AbortController(); + const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, { + api, + callTimeoutMs: 10_000, + configTimeoutMs: 10_000, + pollIntervalMs: 1, + }); + const pending = consumer.consume('microvm-123', { + deadlineMs: Date.now() + 10_000, + signal: controller.signal, + }); + + controller.abort(); + + await expect(pending).rejects.toThrow('runner configuration consumption was cancelled'); + }); +}); diff --git a/lambdas/libs/storage-providers/aws/dynamodb/runner-config-consumer.ts b/lambdas/libs/storage-providers/aws/dynamodb/runner-config-consumer.ts new file mode 100644 index 0000000000..0dfca04c65 --- /dev/null +++ b/lambdas/libs/storage-providers/aws/dynamodb/runner-config-consumer.ts @@ -0,0 +1,147 @@ +import { + ConditionalCheckFailedException, + DeleteItemCommand, + DynamoDBClient, + type DeleteItemCommandOutput, +} from '@aws-sdk/client-dynamodb'; + +import type { + AwsDynamoDbRunnerConfigStorageEnvironment, + RunnerConfigConsumer, + RunnerConfigConsumeOptions, +} from '../../core'; +import { + delay, + isRetryableProviderError, + resolvePollingOptions, + throwIfCancelled, + validateConsumeOptions, + validateDynamoDbRunnerConfigKey, + validateDynamoDbTableName, + withCallDeadline, + type RunnerConfigPollingOptions, +} from '../../runner-config-consumer-common'; +import { EXPIRES_AT_ATTRIBUTE, ID_ATTRIBUTE, RUNNER_CONFIG_ID, SCOPE_ATTRIBUTE, VALUE_ATTRIBUTE } from './keys'; + +export interface AwsDynamoDbRunnerConfigApi { + deleteItem(tableName: string, scope: string, signal: AbortSignal): Promise; +} + +export class AwsSdkDynamoDbRunnerConfigApi implements AwsDynamoDbRunnerConfigApi { + private client?: DynamoDBClient; + + public constructor(client?: DynamoDBClient) { + this.client = client; + } + + private getClient(): DynamoDBClient { + // Do not use the Lambda tracing wrapper here: lifecycle hooks run inside + // the runner image and may be snapshotted before their first request. + this.client ??= new DynamoDBClient({ maxAttempts: 1 }); + return this.client; + } + + public async deleteItem(tableName: string, scope: string, signal: AbortSignal): Promise { + let response: DeleteItemCommandOutput; + try { + response = await this.getClient().send( + new DeleteItemCommand({ + TableName: tableName, + Key: { + [SCOPE_ATTRIBUTE]: { S: scope }, + [ID_ATTRIBUTE]: { S: RUNNER_CONFIG_ID }, + }, + ConditionExpression: 'attribute_exists(#expires_at) AND #expires_at > :now', + ExpressionAttributeNames: { '#expires_at': EXPIRES_AT_ATTRIBUTE }, + ExpressionAttributeValues: { ':now': { N: Math.floor(Date.now() / 1_000).toString() } }, + ReturnValues: 'ALL_OLD', + }), + { abortSignal: signal }, + ); + } catch (error) { + if ( + error instanceof ConditionalCheckFailedException || + (error !== null && + typeof error === 'object' && + 'name' in error && + error.name === 'ConditionalCheckFailedException') + ) { + return undefined; + } + throw error; + } + if (response.Attributes === undefined) { + return undefined; + } + const value = response.Attributes[VALUE_ATTRIBUTE]; + if (value?.S === undefined || value.S.length === 0) { + throw new Error('runner configuration record has an invalid value'); + } + return value.S; + } +} + +export interface AwsDynamoDbRunnerConfigConsumerOptions extends RunnerConfigPollingOptions { + api?: AwsDynamoDbRunnerConfigApi; +} + +export function createAwsDynamoDbRunnerConfigConsumer( + environment: AwsDynamoDbRunnerConfigStorageEnvironment, + options: AwsDynamoDbRunnerConfigConsumerOptions = {}, +): RunnerConfigConsumer { + return new AwsDynamoDbRunnerConfigConsumer( + environment.RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME, + options.api ?? new AwsSdkDynamoDbRunnerConfigApi(), + options, + ); +} + +class AwsDynamoDbRunnerConfigConsumer implements RunnerConfigConsumer { + private readonly callTimeoutMs: number; + private readonly configTimeoutMs: number; + private readonly pollIntervalMs: number; + + public constructor( + private readonly tableName: string, + private readonly api: AwsDynamoDbRunnerConfigApi, + options: AwsDynamoDbRunnerConfigConsumerOptions, + ) { + const polling = resolvePollingOptions(options); + this.callTimeoutMs = polling.callTimeoutMs; + this.configTimeoutMs = polling.configTimeoutMs; + this.pollIntervalMs = polling.pollIntervalMs; + } + + public async consume(runnerId: string, options: RunnerConfigConsumeOptions): Promise { + validateConsumeOptions(options); + const tableName = validateDynamoDbTableName(this.tableName); + validateDynamoDbRunnerConfigKey(runnerId, RUNNER_CONFIG_ID); + const pollDeadline = Math.min(Date.now() + this.configTimeoutMs, options.deadlineMs); + + while (Date.now() < pollDeadline) { + throwIfCancelled(options.signal); + try { + const runnerConfig = await withCallDeadline(options.signal, pollDeadline, this.callTimeoutMs, (callSignal) => + this.api.deleteItem(tableName, runnerId, callSignal), + ); + if (runnerConfig !== undefined) { + return runnerConfig; + } + } catch (error) { + if (options.signal.aborted) { + throw new Error('runner configuration consumption was cancelled'); + } + if (!isRetryableProviderError(error)) { + throw new Error('failed to consume runner configuration from DynamoDB'); + } + } + + const remaining = pollDeadline - Date.now(); + if (remaining > 0) { + await delay(Math.min(this.pollIntervalMs, remaining), options.signal); + } + } + + throw new Error('runner configuration did not become available before the deadline'); + } +} diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-consumer.test.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-consumer.test.ts new file mode 100644 index 0000000000..2a5d2217ea --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-consumer.test.ts @@ -0,0 +1,220 @@ +import { DeleteParameterCommand, GetParameterCommand, type SSMClient } from '@aws-sdk/client-ssm'; +import { afterEach, describe, expect, it, vi } from 'vitest'; + +import { + AwsSdkSsmRunnerConfigApi, + createAwsSsmRunnerConfigConsumer, + type AwsSsmRunnerConfigApi, +} from './runner-config-consumer'; + +function namedError(name: string, message = 'provider detail'): Error { + const error = new Error(message); + error.name = name; + return error; +} + +describe('AWS SDK SSM runner config API', () => { + it('decrypts the parameter and deletes it with the caller abort signal', async () => { + const send = vi + .fn() + .mockResolvedValueOnce({ Parameter: { Value: 'encoded-jit' } }) + .mockResolvedValueOnce({}); + const api = new AwsSdkSsmRunnerConfigApi({ send } as unknown as SSMClient); + const signal = new AbortController().signal; + + await expect(api.getParameter('/runner/tokens/runner-123', signal)).resolves.toBe('encoded-jit'); + await expect(api.deleteParameter('/runner/tokens/runner-123', signal)).resolves.toBeUndefined(); + + expect(send.mock.calls[0][0]).toBeInstanceOf(GetParameterCommand); + expect(send.mock.calls[0][0].input).toEqual({ + Name: '/runner/tokens/runner-123', + WithDecryption: true, + }); + expect(send.mock.calls[0][1]).toEqual({ abortSignal: signal }); + expect(send.mock.calls[1][0]).toBeInstanceOf(DeleteParameterCommand); + expect(send.mock.calls[1][0].input).toEqual({ Name: '/runner/tokens/runner-123' }); + expect(send.mock.calls[1][1]).toEqual({ abortSignal: signal }); + }); +}); + +describe('SSM runner config consumer', () => { + afterEach(() => { + vi.useRealTimers(); + }); + + it('polls a missing parameter, reads it, and deletes it before returning', async () => { + const getParameter = vi + .fn() + .mockRejectedValueOnce(namedError('ParameterNotFound')) + .mockResolvedValueOnce('encoded-jit'); + const deleteParameter = vi.fn().mockResolvedValue(undefined); + const consumer = createAwsSsmRunnerConfigConsumer( + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' }, + { + api: { getParameter, deleteParameter }, + callTimeoutMs: 100, + configTimeoutMs: 500, + pollIntervalMs: 1, + }, + ); + + await expect( + consumer.consume('runner-123', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).resolves.toBe('encoded-jit'); + expect(getParameter).toHaveBeenCalledTimes(2); + expect(deleteParameter).toHaveBeenCalledOnce(); + expect(deleteParameter).toHaveBeenCalledWith('/runner/tokens/runner-123', expect.any(AbortSignal)); + }); + + it('retries a transient delete failure without returning the value early', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const api: AwsSsmRunnerConfigApi = { + getParameter: vi.fn().mockResolvedValue('encoded-jit'), + deleteParameter: vi + .fn() + .mockRejectedValueOnce(namedError('ThrottlingException')) + .mockResolvedValueOnce(undefined), + }; + const consumer = createAwsSsmRunnerConfigConsumer( + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' }, + { api, callTimeoutMs: 100, configTimeoutMs: 2_000, deleteAttempts: 2, pollIntervalMs: 1 }, + ); + + const pending = consumer.consume('runner-123', { + deadlineMs: Date.now() + 3_000, + signal: new AbortController().signal, + }); + await vi.runAllTimersAsync(); + + await expect(pending).resolves.toBe('encoded-jit'); + expect(api.deleteParameter).toHaveBeenCalledTimes(2); + }); + + it('fails closed when another reader deletes the SSM parameter first', async () => { + const api: AwsSsmRunnerConfigApi = { + getParameter: vi.fn().mockResolvedValue('encoded-jit'), + deleteParameter: vi.fn().mockRejectedValue(namedError('ParameterNotFound')), + }; + const consumer = createAwsSsmRunnerConfigConsumer( + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' }, + { api, callTimeoutMs: 100, configTimeoutMs: 100, deleteAttempts: 3, pollIntervalMs: 1 }, + ); + + await expect( + consumer.consume('runner-123', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toThrow('runner configuration could not be deleted from SSM'); + expect(api.deleteParameter).toHaveBeenCalledOnce(); + }); + + it('sanitizes non-retryable provider failures', async () => { + const api: AwsSsmRunnerConfigApi = { + getParameter: vi.fn().mockRejectedValue(namedError('AccessDeniedException', 'encoded-jit-secret')), + deleteParameter: vi.fn(), + }; + const consumer = createAwsSsmRunnerConfigConsumer( + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' }, + { api, callTimeoutMs: 100, configTimeoutMs: 100, pollIntervalMs: 1 }, + ); + + const pending = consumer.consume('runner-123', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }); + await expect(pending).rejects.toThrow('failed to read runner configuration from SSM'); + await expect(pending).rejects.not.toThrow('encoded-jit-secret'); + expect(api.deleteParameter).not.toHaveBeenCalled(); + }); + + it('rejects an empty SSM parameter value without attempting deletion', async () => { + const api: AwsSsmRunnerConfigApi = { + getParameter: vi.fn().mockResolvedValue(''), + deleteParameter: vi.fn(), + }; + const consumer = createAwsSsmRunnerConfigConsumer( + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' }, + { api, callTimeoutMs: 100, configTimeoutMs: 100, pollIntervalMs: 1 }, + ); + + await expect( + consumer.consume('runner-123', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toThrow('failed to read runner configuration from SSM'); + expect(api.deleteParameter).not.toHaveBeenCalled(); + }); + + it('validates the full parameter name before calling SSM', async () => { + const api: AwsSsmRunnerConfigApi = { + getParameter: vi.fn(), + deleteParameter: vi.fn(), + }; + const consumer = createAwsSsmRunnerConfigConsumer( + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: `/${'x'.repeat(890)}` }, + { api, callTimeoutMs: 100, configTimeoutMs: 100, pollIntervalMs: 1 }, + ); + + await expect( + consumer.consume('runner-1234567890', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toThrow('aws_ssm runner configuration key is invalid'); + expect(api.getParameter).not.toHaveBeenCalled(); + }); + + it('stops a provider call immediately when the caller aborts', async () => { + const api: AwsSsmRunnerConfigApi = { + getParameter: vi.fn().mockReturnValue(new Promise(() => undefined)), + deleteParameter: vi.fn(), + }; + const controller = new AbortController(); + const consumer = createAwsSsmRunnerConfigConsumer( + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' }, + { api, callTimeoutMs: 10_000, configTimeoutMs: 10_000, pollIntervalMs: 1 }, + ); + const pending = consumer.consume('runner-123', { + deadlineMs: Date.now() + 10_000, + signal: controller.signal, + }); + + controller.abort(); + + await expect(pending).rejects.toThrow('runner configuration consumption was cancelled'); + }); + + it('reserves a bounded delete attempt when the value appears near the polling deadline', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const startedAt = Date.now(); + let deleteStartedAt: number | undefined; + const api: AwsSsmRunnerConfigApi = { + getParameter: vi.fn().mockResolvedValueOnce(undefined).mockResolvedValueOnce('encoded-jit'), + deleteParameter: vi.fn().mockImplementation(async () => { + deleteStartedAt = Date.now(); + }), + }; + const consumer = createAwsSsmRunnerConfigConsumer( + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' }, + { api, callTimeoutMs: 100, configTimeoutMs: 1_000, pollIntervalMs: 99 }, + ); + + const pending = consumer.consume('runner-123', { + deadlineMs: startedAt + 200, + signal: new AbortController().signal, + }); + await vi.runAllTimersAsync(); + + await expect(pending).resolves.toBe('encoded-jit'); + expect(api.getParameter).toHaveBeenCalledTimes(2); + expect(deleteStartedAt).toBe(startedAt + 99); + expect(deleteStartedAt).toBeLessThanOrEqual(startedAt + 100); + }); +}); diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-consumer.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-consumer.ts new file mode 100644 index 0000000000..7736f8e1ee --- /dev/null +++ b/lambdas/libs/storage-providers/aws/ssm/runner-config-consumer.ts @@ -0,0 +1,167 @@ +import { DeleteParameterCommand, GetParameterCommand, SSMClient } from '@aws-sdk/client-ssm'; + +import type { + AwsSsmRunnerConfigStorageEnvironment, + RunnerConfigConsumer, + RunnerConfigConsumeOptions, +} from '../../core'; +import { + composeSsmParameterName, + delay, + errorName, + isRetryableProviderError, + positiveIntegerOption, + resolvePollingOptions, + throwIfCancelled, + validateConsumeOptions, + withCallDeadline, + type RunnerConfigPollingOptions, +} from '../../runner-config-consumer-common'; + +const DEFAULT_DELETE_ATTEMPTS = 3; + +export interface AwsSsmRunnerConfigApi { + getParameter(name: string, signal: AbortSignal): Promise; + deleteParameter(name: string, signal: AbortSignal): Promise; +} + +export class AwsSdkSsmRunnerConfigApi implements AwsSsmRunnerConfigApi { + private client?: SSMClient; + + public constructor(client?: SSMClient) { + this.client = client; + } + + private getClient(): SSMClient { + // Lifecycle hooks can be snapshotted before their first request. Constructing + // the untraced client here avoids persisting connection state in that snapshot. + this.client ??= new SSMClient({ maxAttempts: 1 }); + return this.client; + } + + public async getParameter(name: string, signal: AbortSignal): Promise { + const response = await this.getClient().send(new GetParameterCommand({ Name: name, WithDecryption: true }), { + abortSignal: signal, + }); + return response.Parameter?.Value; + } + + public async deleteParameter(name: string, signal: AbortSignal): Promise { + await this.getClient().send(new DeleteParameterCommand({ Name: name }), { abortSignal: signal }); + } +} + +export interface AwsSsmRunnerConfigConsumerOptions extends RunnerConfigPollingOptions { + api?: AwsSsmRunnerConfigApi; + deleteAttempts?: number; +} + +export function createAwsSsmRunnerConfigConsumer( + environment: AwsSsmRunnerConfigStorageEnvironment, + options: AwsSsmRunnerConfigConsumerOptions = {}, +): RunnerConfigConsumer { + return new AwsSsmRunnerConfigConsumer( + environment.SSM_TOKEN_PATH, + options.api ?? new AwsSdkSsmRunnerConfigApi(), + options, + ); +} + +class AwsSsmRunnerConfigConsumer implements RunnerConfigConsumer { + private readonly callTimeoutMs: number; + private readonly configTimeoutMs: number; + private readonly deleteAttempts: number; + private readonly pollIntervalMs: number; + + public constructor( + private readonly tokenPath: string, + private readonly api: AwsSsmRunnerConfigApi, + options: AwsSsmRunnerConfigConsumerOptions, + ) { + const polling = resolvePollingOptions(options); + this.callTimeoutMs = polling.callTimeoutMs; + this.configTimeoutMs = polling.configTimeoutMs; + this.pollIntervalMs = polling.pollIntervalMs; + this.deleteAttempts = positiveIntegerOption('deleteAttempts', options.deleteAttempts, DEFAULT_DELETE_ATTEMPTS); + } + + public async consume(runnerId: string, options: RunnerConfigConsumeOptions): Promise { + validateConsumeOptions(options); + const parameterName = composeSsmParameterName(this.tokenPath, runnerId); + const startedAt = Date.now(); + const remainingMs = Math.max(0, options.deadlineMs - startedAt); + // Preserve enough of short hook budgets for at least one bounded delete + // attempt without reviving the old fixed reserve that could consume the + // entire polling window. + const deleteReserveMs = Math.min(this.callTimeoutMs, Math.max(1, Math.floor(remainingMs / 2))); + const pollDeadline = Math.min(startedAt + this.configTimeoutMs, options.deadlineMs - deleteReserveMs); + let runnerConfig: string | undefined; + + while (Date.now() < pollDeadline) { + throwIfCancelled(options.signal); + try { + runnerConfig = await this.read(parameterName, pollDeadline, options.signal); + if (runnerConfig !== undefined) { + if (runnerConfig.length === 0) { + throw new Error('runner configuration record has an invalid value'); + } + break; + } + } catch (error) { + if (options.signal.aborted) { + throw new Error('runner configuration consumption was cancelled'); + } + if (!isSsmNotFound(error) && !isRetryableProviderError(error)) { + throw new Error('failed to read runner configuration from SSM'); + } + } + + const remaining = pollDeadline - Date.now(); + if (remaining > 0) { + await delay(Math.min(this.pollIntervalMs, remaining), options.signal); + } + } + + if (runnerConfig === undefined) { + throw new Error('runner configuration did not become available before the deadline'); + } + + await this.delete(parameterName, options); + return runnerConfig; + } + + private async read(name: string, deadlineMs: number, signal: AbortSignal): Promise { + return withCallDeadline(signal, deadlineMs, this.callTimeoutMs, (callSignal) => + this.api.getParameter(name, callSignal), + ); + } + + private async delete(name: string, options: RunnerConfigConsumeOptions): Promise { + for (let attempt = 1; attempt <= this.deleteAttempts; attempt += 1) { + try { + await withCallDeadline(options.signal, options.deadlineMs, this.callTimeoutMs, (callSignal) => + this.api.deleteParameter(name, callSignal), + ); + return; + } catch (error) { + if (options.signal.aborted) { + throw new Error('runner configuration consumption was cancelled'); + } + if (!isRetryableProviderError(error) || attempt === this.deleteAttempts) { + break; + } + + const remaining = options.deadlineMs - Date.now(); + if (remaining <= 0) { + break; + } + await delay(Math.min(2 ** (attempt - 1) * 1_000, 5_000, remaining), options.signal); + } + } + throw new Error('runner configuration could not be deleted from SSM'); + } +} + +function isSsmNotFound(error: unknown): boolean { + return errorName(error) === 'ParameterNotFound'; +} diff --git a/lambdas/libs/storage-providers/core/index.ts b/lambdas/libs/storage-providers/core/index.ts index 57c3240c39..ea5559aef9 100644 --- a/lambdas/libs/storage-providers/core/index.ts +++ b/lambdas/libs/storage-providers/core/index.ts @@ -29,6 +29,33 @@ export interface RunnerConfigStore { houseKeeper(): Promise; } +export interface AwsSsmRunnerConfigStorageEnvironment { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm'; + SSM_TOKEN_PATH: string; +} + +export interface AwsDynamoDbRunnerConfigStorageEnvironment { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb'; + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: string; +} + +export type RunnerConfigStorageEnvironment = + | AwsSsmRunnerConfigStorageEnvironment + | AwsDynamoDbRunnerConfigStorageEnvironment; + +/** Exact environment-variable map accepted from `runHookPayload.context.storage`. */ +export type RunnerConfigStorageContext = RunnerConfigStorageEnvironment; + +export interface RunnerConfigConsumeOptions { + /** Absolute Unix time in milliseconds after which the operation must stop. */ + deadlineMs: number; + signal: AbortSignal; +} + +export interface RunnerConfigConsumer { + consume(runnerId: string, options: RunnerConfigConsumeOptions): Promise; +} + export interface RunnerGroupCacheRecord { runnerGroupName: string; runnerGroupId: number; diff --git a/lambdas/libs/storage-providers/index.ts b/lambdas/libs/storage-providers/index.ts index 5dde65a02e..48e9dfe122 100644 --- a/lambdas/libs/storage-providers/index.ts +++ b/lambdas/libs/storage-providers/index.ts @@ -3,6 +3,12 @@ export type { GitHubAppCredential, GitHubAppCredentialsStore, GitHubWebhookSecretStore, + AwsDynamoDbRunnerConfigStorageEnvironment, + AwsSsmRunnerConfigStorageEnvironment, + RunnerConfigConsumeOptions, + RunnerConfigConsumer, + RunnerConfigStorageContext, + RunnerConfigStorageEnvironment, RunnerConfigMetadata, RunnerConfigRecord, RunnerConfigStore, @@ -19,6 +25,15 @@ export type { } from './core'; export { getGitHubAppCredentialsStore, resetGitHubAppCredentialsStore } from './github-app-credentials'; export { getGitHubWebhookSecretStore, resetGitHubWebhookSecretStore } from './github-webhook-secret'; +export { + createRunnerConfigConsumerFromEnvironment, + exportRunnerConfigStorageEnvironment, + loadRunnerConfigConsumerConfigFromEnvironment, + loadRunnerConfigStorageContextFromEnvironment, + parseRunnerConfigStorageContext, + runnerConfigStorageEnvironment, + type RunnerConfigConsumerConfig, +} from './runner-config-consumer'; export { getRunnerConfigStore, resetRunnerConfigStore } from './runner-config'; export { getRunnerGroupCacheStore, resetRunnerGroupCacheStore } from './runner-group-cache'; export { getRunnerMatcherConfigStore, resetRunnerMatcherConfigStore } from './runner-matcher-config'; diff --git a/lambdas/libs/storage-providers/package.json b/lambdas/libs/storage-providers/package.json index 798cc52e32..5cca0dc512 100644 --- a/lambdas/libs/storage-providers/package.json +++ b/lambdas/libs/storage-providers/package.json @@ -3,7 +3,8 @@ "version": "1.0.0", "main": "index.ts", "exports": { - ".": "./index.ts" + ".": "./index.ts", + "./runner-config-consumer": "./runner-config-consumer.ts" }, "type": "module", "license": "MIT", diff --git a/lambdas/libs/storage-providers/runner-config-consumer-common.ts b/lambdas/libs/storage-providers/runner-config-consumer-common.ts new file mode 100644 index 0000000000..2eb61f380e --- /dev/null +++ b/lambdas/libs/storage-providers/runner-config-consumer-common.ts @@ -0,0 +1,243 @@ +import { Buffer } from 'node:buffer'; + +import type { RunnerConfigConsumeOptions } from './core'; + +export const DEFAULT_CALL_TIMEOUT_MS = 5_000; +export const DEFAULT_CONFIG_TIMEOUT_MS = 40_000; +export const DEFAULT_POLL_INTERVAL_MS = 2_000; + +const RUNNER_ID_PATTERN = /^[A-Za-z0-9_.-]{1,256}$/; +const SSM_PARAMETER_PATH_PATTERN = /^\/[A-Za-z0-9_.\-/]+$/; +const DYNAMODB_TABLE_NAME_PATTERN = /^[A-Za-z0-9_.-]{3,255}$/; +// AWS counts the partition/region/account ARN prefix toward its 1,011-character +// limit. Leave ample room for that deployment-specific prefix. +const MAX_SSM_PARAMETER_NAME_LENGTH = 900; +const MAX_DYNAMODB_PARTITION_KEY_BYTES = 2_048; +const MAX_DYNAMODB_SORT_KEY_BYTES = 1_024; + +const RETRYABLE_ERROR_NAMES = new Set([ + 'AbortError', + 'ConnectionError', + 'InternalServerException', + 'ProvisionedThroughputExceededException', + 'RequestLimitExceeded', + 'RequestTimeout', + 'ServiceUnavailable', + 'ThrottlingException', + 'TimeoutError', +]); + +export interface RunnerConfigPollingOptions { + callTimeoutMs?: number; + configTimeoutMs?: number; + pollIntervalMs?: number; +} + +export interface ResolvedRunnerConfigPollingOptions { + callTimeoutMs: number; + configTimeoutMs: number; + pollIntervalMs: number; +} + +class RunnerConfigCallDeadlineError extends Error { + public constructor() { + super('runner configuration provider call exceeded its deadline'); + this.name = 'RunnerConfigCallDeadlineError'; + } +} + +export function resolvePollingOptions(options: RunnerConfigPollingOptions): ResolvedRunnerConfigPollingOptions { + return { + callTimeoutMs: positiveIntegerOption('callTimeoutMs', options.callTimeoutMs, DEFAULT_CALL_TIMEOUT_MS), + configTimeoutMs: positiveIntegerOption('configTimeoutMs', options.configTimeoutMs, DEFAULT_CONFIG_TIMEOUT_MS), + pollIntervalMs: positiveIntegerOption('pollIntervalMs', options.pollIntervalMs, DEFAULT_POLL_INTERVAL_MS), + }; +} + +export function positiveIntegerOption(name: string, value: number | undefined, fallback: number): number { + if (value === undefined) { + return fallback; + } + if (!Number.isSafeInteger(value) || value <= 0) { + throw new Error(`${name} must be a positive integer`); + } + return value; +} + +export function validateRunnerId(runnerId: string): void { + if (!RUNNER_ID_PATTERN.test(runnerId)) { + throw new Error('runnerId is invalid'); + } +} + +export function canonicalSsmTokenPath(tokenPath: string): string { + if (tokenPath.includes('//')) { + throw new Error('aws_ssm tokenPath is invalid'); + } + const canonical = tokenPath.endsWith('/') ? tokenPath.slice(0, -1) : tokenPath; + const segments = canonical.split('/').slice(1); + if ( + canonical.length === 0 || + canonical.length > MAX_SSM_PARAMETER_NAME_LENGTH || + !SSM_PARAMETER_PATH_PATTERN.test(canonical) || + segments.length > 14 || + segments.some((segment) => segment === '' || segment === '.' || segment === '..') || + /^(aws|ssm)/i.test(segments[0] ?? '') + ) { + throw new Error('aws_ssm tokenPath is invalid'); + } + return canonical; +} + +export function composeSsmParameterName(tokenPath: string, runnerId: string): string { + validateRunnerId(runnerId); + const parameterName = `${canonicalSsmTokenPath(tokenPath)}/${runnerId}`; + const segments = parameterName.split('/').slice(1); + if (parameterName.length > MAX_SSM_PARAMETER_NAME_LENGTH || segments.length > 15) { + throw new Error('aws_ssm runner configuration key is invalid'); + } + return parameterName; +} + +export function validateDynamoDbTableName(tableName: string): string { + if (!DYNAMODB_TABLE_NAME_PATTERN.test(tableName)) { + throw new Error('aws_dynamodb tableName is invalid'); + } + return tableName; +} + +export function validateDynamoDbRunnerConfigKey(scope: string, id: string): void { + validateRunnerId(scope); + if ( + Buffer.byteLength(scope, 'utf8') > MAX_DYNAMODB_PARTITION_KEY_BYTES || + id.length === 0 || + Buffer.byteLength(id, 'utf8') > MAX_DYNAMODB_SORT_KEY_BYTES + ) { + throw new Error('aws_dynamodb runner configuration key is invalid'); + } +} + +export function validateConsumeOptions(options: RunnerConfigConsumeOptions): void { + if (!Number.isSafeInteger(options.deadlineMs) || options.deadlineMs <= 0) { + throw new Error('deadlineMs must be a positive integer'); + } + if ( + options.signal === null || + typeof options.signal !== 'object' || + typeof options.signal.aborted !== 'boolean' || + typeof options.signal.addEventListener !== 'function' || + typeof options.signal.removeEventListener !== 'function' + ) { + throw new Error('signal must be an AbortSignal'); + } +} + +export function errorName(error: unknown): string { + if (error !== null && typeof error === 'object' && 'name' in error && typeof error.name === 'string') { + return error.name; + } + return 'UnknownError'; +} + +function httpStatus(error: unknown): number | undefined { + if ( + error !== null && + typeof error === 'object' && + '$metadata' in error && + error.$metadata !== null && + typeof error.$metadata === 'object' && + 'httpStatusCode' in error.$metadata && + typeof error.$metadata.httpStatusCode === 'number' + ) { + return error.$metadata.httpStatusCode; + } + return undefined; +} + +export function isRetryableProviderError(error: unknown): boolean { + const status = httpStatus(error); + return ( + error instanceof RunnerConfigCallDeadlineError || + RETRYABLE_ERROR_NAMES.has(errorName(error)) || + (status !== undefined && status >= 500) + ); +} + +export function delay(ms: number, signal: AbortSignal): Promise { + if (signal.aborted) { + return Promise.reject(new Error('runner configuration consumption was cancelled')); + } + + return new Promise((resolve, reject) => { + let settled = false; + const cleanup = (): void => signal.removeEventListener('abort', cancel); + const finish = (): void => { + if (settled) { + return; + } + settled = true; + cleanup(); + resolve(); + }; + const timer = setTimeout(finish, ms); + const cancel = (): void => { + if (settled) { + return; + } + settled = true; + clearTimeout(timer); + cleanup(); + reject(new Error('runner configuration consumption was cancelled')); + }; + signal.addEventListener('abort', cancel, { once: true }); + }); +} + +export async function withCallDeadline( + parentSignal: AbortSignal, + deadlineMs: number, + callTimeoutMs: number, + operation: (signal: AbortSignal) => Promise, +): Promise { + if (parentSignal.aborted) { + throw new Error('runner configuration consumption was cancelled'); + } + + const remaining = deadlineMs - Date.now(); + if (remaining <= 0) { + throw new RunnerConfigCallDeadlineError(); + } + + const controller = new AbortController(); + let cancel!: () => void; + let timeout: ReturnType | undefined; + const deadline = new Promise((_resolve, reject) => { + cancel = (): void => { + reject(new Error('runner configuration consumption was cancelled')); + controller.abort(); + }; + parentSignal.addEventListener('abort', cancel, { once: true }); + timeout = setTimeout( + () => { + reject(new RunnerConfigCallDeadlineError()); + controller.abort(); + }, + Math.max(1, Math.min(remaining, callTimeoutMs)), + ); + }); + + try { + return await Promise.race([operation(controller.signal), deadline]); + } finally { + if (timeout !== undefined) { + clearTimeout(timeout); + } + parentSignal.removeEventListener('abort', cancel); + } +} + +export function throwIfCancelled(signal: AbortSignal): void { + if (signal.aborted) { + throw new Error('runner configuration consumption was cancelled'); + } +} diff --git a/lambdas/libs/storage-providers/runner-config-consumer-subpath.test.ts b/lambdas/libs/storage-providers/runner-config-consumer-subpath.test.ts new file mode 100644 index 0000000000..1ed795a5d1 --- /dev/null +++ b/lambdas/libs/storage-providers/runner-config-consumer-subpath.test.ts @@ -0,0 +1,31 @@ +import { + createRunnerConfigConsumerFromEnvironment, + exportRunnerConfigStorageEnvironment, + loadRunnerConfigStorageContextFromEnvironment, + parseRunnerConfigStorageContext, + type RunnerConfigConsumeOptions, + type RunnerConfigConsumer, + type RunnerConfigStorageContext, + type RunnerConfigStorageEnvironment, +} from '@aws-github-runner/storage-providers/runner-config-consumer'; +import { describe, expect, it } from 'vitest'; + +describe('runner config consumer package subpath', () => { + it('exposes the portable environment round-trip and consumer contract', () => { + const context: RunnerConfigStorageContext = { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', + SSM_TOKEN_PATH: '/runner/tokens', + }; + const options: RunnerConfigConsumeOptions = { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }; + const exported: RunnerConfigStorageEnvironment = exportRunnerConfigStorageEnvironment(context, {}); + const consumerFactory: () => RunnerConfigConsumer = createRunnerConfigConsumerFromEnvironment; + + expect(parseRunnerConfigStorageContext(context)).toEqual(context); + expect(loadRunnerConfigStorageContextFromEnvironment(exported)).toEqual(context); + expect(consumerFactory).toBeTypeOf('function'); + expect(options.signal.aborted).toBe(false); + }); +}); diff --git a/lambdas/libs/storage-providers/runner-config-consumer.test.ts b/lambdas/libs/storage-providers/runner-config-consumer.test.ts new file mode 100644 index 0000000000..f765b22780 --- /dev/null +++ b/lambdas/libs/storage-providers/runner-config-consumer.test.ts @@ -0,0 +1,238 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; + +import type { AwsDynamoDbRunnerConfigApi } from './aws/dynamodb/runner-config-consumer'; +import { AwsSdkSsmRunnerConfigApi, type AwsSsmRunnerConfigApi } from './aws/ssm/runner-config-consumer'; +import { + createRunnerConfigConsumerFromEnvironment, + exportRunnerConfigStorageEnvironment, + loadRunnerConfigConsumerConfigFromEnvironment, + loadRunnerConfigStorageContextFromEnvironment, + parseRunnerConfigStorageContext, + runnerConfigStorageEnvironment, +} from './runner-config-consumer'; + +const ssmContext = { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', + SSM_TOKEN_PATH: '/runner/tokens', +} as const; +const dynamoDbContext = { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb', + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'runner-state', +} as const; + +describe('runner config storage context', () => { + it('parses and freezes an exact SSM environment map while canonicalizing one trailing slash', () => { + const context = parseRunnerConfigStorageContext({ ...ssmContext, SSM_TOKEN_PATH: '/runner/tokens/' }); + + expect(context).toEqual(ssmContext); + expect(Object.isFrozen(context)).toBe(true); + expect(runnerConfigStorageEnvironment(context)).toEqual(ssmContext); + }); + + it('parses and freezes an exact DynamoDB environment map', () => { + const context = parseRunnerConfigStorageContext(dynamoDbContext); + + expect(context).toEqual(dynamoDbContext); + expect(Object.isFrozen(context)).toBe(true); + }); + + it.each([ + null, + [], + 'aws_ssm', + { RUNNER_CONFIG_STORAGE_PROVIDER: 'AWS_SSM', SSM_TOKEN_PATH: '/runner/tokens' }, + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm' }, + { ...ssmContext, unexpected: true }, + { ...ssmContext, AWS_ACCESS_KEY_ID: 'payload-must-not-export-credentials' }, + { ...ssmContext, RUNNER_CONFIG_TIMEOUT_SECONDS: '60' }, + { ...ssmContext, RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'runner-state' }, + { ...ssmContext, SSM_TOKEN_PATH: '/runner//tokens' }, + { ...ssmContext, SSM_TOKEN_PATH: '/runner/tokens//' }, + { ...ssmContext, SSM_TOKEN_PATH: '/awsParameters/tokens' }, + { ...ssmContext, SSM_TOKEN_PATH: '/ssm-private/tokens' }, + { ...ssmContext, SSM_TOKEN_PATH: '/runner/../tokens' }, + { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb', RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'ab' }, + { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb', + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'runner state', + }, + { ...dynamoDbContext, SSM_TOKEN_PATH: '/runner/tokens' }, + { provider: 'aws_dynamodb', tableName: 'runner-state' }, + ])('rejects a non-allowlisted or incomplete storage context %#', (value) => { + expect(() => parseRunnerConfigStorageContext(value)).toThrow(); + }); + + it('rejects symbol fields that would be hidden by JSON-style key enumeration', () => { + const context = { ...ssmContext }; + Object.defineProperty(context, Symbol('unexpected'), { value: true }); + + expect(() => parseRunnerConfigStorageContext(context)).toThrow('storage context is invalid'); + }); + + it.each([ + [ + { + ...ssmContext, + AWS_ACCESS_KEY_ID: 'producer-only', + RUNNER_CONFIG_TIMEOUT_SECONDS: '30', + UNRELATED: 'kept-out', + }, + ssmContext, + ], + [ + { + ...dynamoDbContext, + RUNNER_CONFIG_DYNAMODB_CONFIG_TABLE_NAME: 'durable-config', + RUNNER_CONFIG_DYNAMODB_ENTRY_ID: 'linux', + RUNNER_CONFIG_DYNAMODB_TTL_SECONDS: '3600', + RUNNER_CONFIG_STORAGE_VERSION: 'version-hash', + }, + dynamoDbContext, + ], + ])('selects only the chosen provider locator from a broader producer environment %#', (environment, expected) => { + expect(loadRunnerConfigStorageContextFromEnvironment(environment)).toEqual(expected); + }); + + it('defaults a legacy producer environment with only SSM_TOKEN_PATH to aws_ssm', () => { + expect(loadRunnerConfigStorageContextFromEnvironment({ SSM_TOKEN_PATH: '/runner/tokens' })).toEqual(ssmContext); + }); + + it('round-trips each producer environment through payload context and hook environment export', () => { + for (const producerEnvironment of [ssmContext, dynamoDbContext]) { + const payloadContext = loadRunnerConfigStorageContextFromEnvironment(producerEnvironment); + const hookEnvironment: Record = { + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'stale-table', + SSM_TOKEN_PATH: '/stale/path', + UNRELATED: 'preserved', + }; + + expect(exportRunnerConfigStorageEnvironment(payloadContext, hookEnvironment)).toEqual(payloadContext); + expect(loadRunnerConfigStorageContextFromEnvironment(hookEnvironment)).toEqual(payloadContext); + expect(hookEnvironment.UNRELATED).toBe('preserved'); + if (payloadContext.RUNNER_CONFIG_STORAGE_PROVIDER === 'aws_ssm') { + expect(hookEnvironment.RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME).toBeUndefined(); + } else { + expect(hookEnvironment.SSM_TOKEN_PATH).toBeUndefined(); + } + } + }); + + it('does not mutate a target when context validation fails', () => { + const target = { ...ssmContext } as Record; + + expect(() => + exportRunnerConfigStorageEnvironment({ ...dynamoDbContext, AWS_SECRET_ACCESS_KEY: 'forbidden' } as never, target), + ).toThrow(); + expect(target).toEqual(ssmContext); + }); +}); + +describe('runner config consumer environment factory', () => { + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it('creates an injected SSM consumer from exported environment', async () => { + const api: AwsSsmRunnerConfigApi = { + getParameter: vi.fn().mockResolvedValue('encoded-jit'), + deleteParameter: vi.fn().mockResolvedValue(undefined), + }; + const consumer = createRunnerConfigConsumerFromEnvironment(ssmContext, { + awsSsmApi: api, + callTimeoutMs: 100, + configTimeoutMs: 100, + pollIntervalMs: 1, + }); + + await expect( + consumer.consume('microvm-123', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).resolves.toBe('encoded-jit'); + expect(api.getParameter).toHaveBeenCalledWith('/runner/tokens/microvm-123', expect.any(AbortSignal)); + expect(api.deleteParameter).toHaveBeenCalledWith('/runner/tokens/microvm-123', expect.any(AbortSignal)); + }); + + it('creates an injected DynamoDB consumer using microvmId as the access scope', async () => { + const api: AwsDynamoDbRunnerConfigApi = { + deleteItem: vi.fn().mockResolvedValue('encoded-jit'), + }; + const consumer = createRunnerConfigConsumerFromEnvironment(dynamoDbContext, { + awsDynamoDbApi: api, + callTimeoutMs: 100, + configTimeoutMs: 100, + pollIntervalMs: 1, + }); + + await expect( + consumer.consume('microvm-123', { + deadlineMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).resolves.toBe('encoded-jit'); + expect(api.deleteItem).toHaveBeenCalledWith('runner-state', 'microvm-123', expect.any(AbortSignal)); + }); + + it('loads timing defaults and overrides from the supplied factory environment', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + vi.spyOn(AwsSdkSsmRunnerConfigApi.prototype, 'getParameter').mockResolvedValue(undefined); + const consumer = createRunnerConfigConsumerFromEnvironment( + { + ...ssmContext, + RUNNER_CONFIG_TIMEOUT_SECONDS: '1', + RUNNER_CONFIG_POLL_SECONDS: '1', + }, + undefined, + ); + const startedAt = Date.now(); + const pending = consumer.consume('microvm-123', { + deadlineMs: startedAt + 10_000, + signal: new AbortController().signal, + }); + const rejection = expect(pending).rejects.toThrow( + 'runner configuration did not become available before the deadline', + ); + + await vi.runAllTimersAsync(); + + await rejection; + expect(Date.now() - startedAt).toBe(1_000); + }); + + it('loads source-compatible timing defaults', () => { + expect(loadRunnerConfigConsumerConfigFromEnvironment({})).toEqual({ + callTimeoutMs: 5_000, + configTimeoutMs: 20_000, + deleteAttempts: 3, + pollIntervalMs: 2_000, + }); + }); + + it('loads bounded timing overrides from an environment', () => { + expect( + loadRunnerConfigConsumerConfigFromEnvironment({ + AWS_SDK_CALL_TIMEOUT_SECONDS: '7', + RUNNER_CONFIG_TIMEOUT_SECONDS: '31', + RUNNER_CONFIG_DELETE_ATTEMPTS: '4', + RUNNER_CONFIG_POLL_SECONDS: '3', + }), + ).toEqual({ + callTimeoutMs: 7_000, + configTimeoutMs: 31_000, + deleteAttempts: 4, + pollIntervalMs: 3_000, + }); + }); + + it.each([ + ['AWS_SDK_CALL_TIMEOUT_SECONDS', '0', 'callTimeoutMs', 5_000], + ['RUNNER_CONFIG_TIMEOUT_SECONDS', '61', 'configTimeoutMs', 20_000], + ['RUNNER_CONFIG_DELETE_ATTEMPTS', '11', 'deleteAttempts', 3], + ['RUNNER_CONFIG_POLL_SECONDS', 'not-a-number', 'pollIntervalMs', 2_000], + ])('falls back for invalid %s=%j', (name, value, property, expected) => { + expect(loadRunnerConfigConsumerConfigFromEnvironment({ [name]: value })).toHaveProperty(property, expected); + }); +}); diff --git a/lambdas/libs/storage-providers/runner-config-consumer.ts b/lambdas/libs/storage-providers/runner-config-consumer.ts new file mode 100644 index 0000000000..b820b1c092 --- /dev/null +++ b/lambdas/libs/storage-providers/runner-config-consumer.ts @@ -0,0 +1,189 @@ +import { + createAwsDynamoDbRunnerConfigConsumer, + type AwsDynamoDbRunnerConfigApi, +} from './aws/dynamodb/runner-config-consumer'; +import { createAwsSsmRunnerConfigConsumer, type AwsSsmRunnerConfigApi } from './aws/ssm/runner-config-consumer'; +import type { RunnerConfigConsumer, RunnerConfigStorageContext, RunnerConfigStorageEnvironment } from './core'; +import { + canonicalSsmTokenPath, + validateDynamoDbTableName, + type RunnerConfigPollingOptions, +} from './runner-config-consumer-common'; + +export type { + AwsDynamoDbRunnerConfigStorageEnvironment, + AwsSsmRunnerConfigStorageEnvironment, + RunnerConfigConsumeOptions, + RunnerConfigConsumer, + RunnerConfigStorageContext, + RunnerConfigStorageEnvironment, +} from './core'; +export type { AwsDynamoDbRunnerConfigApi } from './aws/dynamodb/runner-config-consumer'; +export type { AwsSsmRunnerConfigApi } from './aws/ssm/runner-config-consumer'; + +const STORAGE_PROVIDER_ENVIRONMENT_VARIABLE = 'RUNNER_CONFIG_STORAGE_PROVIDER'; +const SSM_TOKEN_PATH_ENVIRONMENT_VARIABLE = 'SSM_TOKEN_PATH'; +const DYNAMODB_TABLE_ENVIRONMENT_VARIABLE = 'RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME'; +const STORAGE_ENVIRONMENT_VARIABLES = [ + STORAGE_PROVIDER_ENVIRONMENT_VARIABLE, + SSM_TOKEN_PATH_ENVIRONMENT_VARIABLE, + DYNAMODB_TABLE_ENVIRONMENT_VARIABLE, +] as const; + +type Environment = Readonly>; +type MutableEnvironment = Record; + +export interface RunnerConfigConsumerConfig extends RunnerConfigPollingOptions { + awsDynamoDbApi?: AwsDynamoDbRunnerConfigApi; + awsSsmApi?: AwsSsmRunnerConfigApi; + deleteAttempts?: number; +} + +export function parseRunnerConfigStorageContext(value: unknown): RunnerConfigStorageContext { + if (!isPlainObject(value) || typeof value.RUNNER_CONFIG_STORAGE_PROVIDER !== 'string') { + throw new Error('runner configuration storage context is invalid'); + } + + if (value.RUNNER_CONFIG_STORAGE_PROVIDER === 'aws_ssm') { + if ( + !hasExactKeys(value, [STORAGE_PROVIDER_ENVIRONMENT_VARIABLE, SSM_TOKEN_PATH_ENVIRONMENT_VARIABLE]) || + typeof value.SSM_TOKEN_PATH !== 'string' + ) { + throw new Error('aws_ssm runner configuration storage context is invalid'); + } + return Object.freeze({ + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', + SSM_TOKEN_PATH: canonicalSsmTokenPath(value.SSM_TOKEN_PATH), + }); + } + + if (value.RUNNER_CONFIG_STORAGE_PROVIDER === 'aws_dynamodb') { + if ( + !hasExactKeys(value, [STORAGE_PROVIDER_ENVIRONMENT_VARIABLE, DYNAMODB_TABLE_ENVIRONMENT_VARIABLE]) || + typeof value.RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME !== 'string' + ) { + throw new Error('aws_dynamodb runner configuration storage context is invalid'); + } + return Object.freeze({ + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb', + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: validateDynamoDbTableName( + value.RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME, + ), + }); + } + + throw new Error('runner configuration storage provider is unsupported'); +} + +/** Selects only the consumer-safe storage variables from a broader producer environment. */ +export function loadRunnerConfigStorageContextFromEnvironment( + environment: Environment = process.env, +): RunnerConfigStorageContext { + const configuredProvider = environment[STORAGE_PROVIDER_ENVIRONMENT_VARIABLE]; + const provider = + configuredProvider === undefined || configuredProvider.trim() === '' ? 'aws_ssm' : configuredProvider; + + if (provider === 'aws_ssm') { + return parseRunnerConfigStorageContext({ + RUNNER_CONFIG_STORAGE_PROVIDER: provider, + SSM_TOKEN_PATH: environment[SSM_TOKEN_PATH_ENVIRONMENT_VARIABLE], + }); + } + if (provider === 'aws_dynamodb') { + return parseRunnerConfigStorageContext({ + RUNNER_CONFIG_STORAGE_PROVIDER: provider, + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: environment[DYNAMODB_TABLE_ENVIRONMENT_VARIABLE], + }); + } + throw new Error('runner configuration storage provider is unsupported'); +} + +/** Revalidates a payload context and returns its exact environment-variable map. */ +export function runnerConfigStorageEnvironment(context: RunnerConfigStorageContext): RunnerConfigStorageEnvironment { + return parseRunnerConfigStorageContext(context); +} + +/** + * Exports only allowlisted storage variables. Provider locators from a previous + * run are cleared so they cannot influence the selected consumer. + */ +export function exportRunnerConfigStorageEnvironment( + context: RunnerConfigStorageContext, + target: MutableEnvironment = process.env, +): RunnerConfigStorageEnvironment { + const environment = runnerConfigStorageEnvironment(context); + for (const name of STORAGE_ENVIRONMENT_VARIABLES) { + delete target[name]; + } + for (const [name, value] of Object.entries(environment)) { + target[name] = value; + } + return environment; +} + +export function createRunnerConfigConsumerFromEnvironment( + environment: Environment = process.env, + config?: RunnerConfigConsumerConfig, +): RunnerConfigConsumer { + const storage = loadRunnerConfigStorageContextFromEnvironment(environment); + const resolvedConfig = config ?? loadRunnerConfigConsumerConfigFromEnvironment(environment); + switch (storage.RUNNER_CONFIG_STORAGE_PROVIDER) { + case 'aws_ssm': + return createAwsSsmRunnerConfigConsumer(storage, { + api: resolvedConfig.awsSsmApi, + callTimeoutMs: resolvedConfig.callTimeoutMs, + configTimeoutMs: resolvedConfig.configTimeoutMs, + deleteAttempts: resolvedConfig.deleteAttempts, + pollIntervalMs: resolvedConfig.pollIntervalMs, + }); + case 'aws_dynamodb': + return createAwsDynamoDbRunnerConfigConsumer(storage, { + api: resolvedConfig.awsDynamoDbApi, + callTimeoutMs: resolvedConfig.callTimeoutMs, + configTimeoutMs: resolvedConfig.configTimeoutMs, + pollIntervalMs: resolvedConfig.pollIntervalMs, + }); + } +} + +export function loadRunnerConfigConsumerConfigFromEnvironment( + environment: Environment = process.env, +): RunnerConfigConsumerConfig { + return { + callTimeoutMs: secondsEnvironmentValue(environment, 'AWS_SDK_CALL_TIMEOUT_SECONDS', 5) * 1_000, + configTimeoutMs: secondsEnvironmentValue(environment, 'RUNNER_CONFIG_TIMEOUT_SECONDS', 20) * 1_000, + deleteAttempts: positiveIntegerEnvironmentValue(environment, 'RUNNER_CONFIG_DELETE_ATTEMPTS', 3, 10), + pollIntervalMs: secondsEnvironmentValue(environment, 'RUNNER_CONFIG_POLL_SECONDS', 2) * 1_000, + }; +} + +function secondsEnvironmentValue(environment: Environment, name: string, fallback: number): number { + return positiveIntegerEnvironmentValue(environment, name, fallback, 60); +} + +function positiveIntegerEnvironmentValue( + environment: Environment, + name: string, + fallback: number, + maximum: number, +): number { + const value = environment[name]; + if (value === undefined || !/^\d+$/.test(value)) { + return fallback; + } + const parsed = Number(value); + return Number.isSafeInteger(parsed) && parsed > 0 && parsed <= maximum ? parsed : fallback; +} + +function isPlainObject(value: unknown): value is Record { + if (value === null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const prototype = Object.getPrototypeOf(value) as unknown; + return prototype === Object.prototype || prototype === null; +} + +function hasExactKeys(value: object, expected: readonly string[]): boolean { + const keys = Reflect.ownKeys(value); + return keys.length === expected.length && keys.every((key) => typeof key === 'string' && expected.includes(key)); +} diff --git a/lambdas/libs/storage-providers/vitest.config.ts b/lambdas/libs/storage-providers/vitest.config.ts index 21626b5cc9..1b04035580 100644 --- a/lambdas/libs/storage-providers/vitest.config.ts +++ b/lambdas/libs/storage-providers/vitest.config.ts @@ -12,6 +12,8 @@ export default mergeConfig(defaultConfig, { 'provider.ts', 'github-app-credentials.ts', 'github-webhook-secret.ts', + 'runner-config-consumer.ts', + 'runner-config-consumer-common.ts', 'runner-config.ts', 'runner-group-cache.ts', 'runner-matcher-config.ts', diff --git a/lambdas/package.json b/lambdas/package.json index c6fa5d72c3..0a0b0b088a 100644 --- a/lambdas/package.json +++ b/lambdas/package.json @@ -3,7 +3,8 @@ "private": true, "workspaces": [ "functions/*", - "libs/*" + "libs/*", + "services/*" ], "scripts": { "build": "nx run-many --target=build --all", diff --git a/lambdas/services/microvm-lifecycle-hooks/README.md b/lambdas/services/microvm-lifecycle-hooks/README.md new file mode 100644 index 0000000000..6d75bc6d68 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/README.md @@ -0,0 +1,143 @@ +# Lambda MicroVM lifecycle hooks + +This service implements the lifecycle-hook HTTP server used to start one ephemeral GitHub Actions runner inside an AWS Lambda MicroVM. Storage-specific reads and one-time consumption are delegated to `@aws-github-runner/storage-providers`; this package owns only payload validation, lifecycle state, and the runner process boundary. + +## Build and run + +From `lambdas/`: + +```bash +yarn nx test @aws-github-runner/microvm-lifecycle-hooks +yarn workspace @aws-github-runner/microvm-lifecycle-hooks build +yarn workspace @aws-github-runner/microvm-lifecycle-hooks start +``` + +`build` uses NCC to create a self-contained `dist/`. It also writes `dist/package.json` with `type: module`, so the artifact runs after it is copied outside the Yarn workspace. Copy the **entire** directory; do not copy only `index.js`. + +To build before invoking Docker, run the workspace build above. In the existing MicroVM runner Dockerfile, which already installs s6-overlay and the GitHub runner's Node 24 runtime, copy the complete artifact and replace the old hook command with: + +```dockerfile +COPY lambdas/services/microvm-lifecycle-hooks/dist/ /opt/microvm-lifecycle-hooks/ +ENV RUNNER_ENTRYPOINT=/opt/microvm/entrypoint.sh +ENTRYPOINT ["/init"] +CMD ["/command/with-contenv", "/opt/actions-runner/externals/node24/bin/node", "/opt/microvm-lifecycle-hooks/index.js"] +``` + +Alternatively, build the service inside Docker with the repository root as the build context. Add this pinned builder stage: + +```dockerfile +ARG NODE_BUILDER_IMAGE=node:24-bookworm-slim@sha256:3638d9a6fe4030bd716be989438248074489337ba3275657f93595428be4fc03 +FROM ${NODE_BUILDER_IMAGE} AS lifecycle-build +WORKDIR /source +COPY lambdas/ ./lambdas/ +RUN corepack enable \ + && cd lambdas \ + && yarn install --immutable \ + && yarn workspace @aws-github-runner/microvm-lifecycle-hooks build +``` + +Use a clean checkout for that build context, or exclude local `node_modules/`, `coverage/`, and `dist/` directories with `.dockerignore`, so host-built dependencies are not copied into the Linux builder. + +Then copy the builder output into the existing final runner stage and use its supervisor and Node 24 runtime: + +```dockerfile +COPY --from=lifecycle-build \ + /source/lambdas/services/microvm-lifecycle-hooks/dist/ \ + /opt/microvm-lifecycle-hooks/ +ENV RUNNER_ENTRYPOINT=/opt/microvm/entrypoint.sh +ENTRYPOINT ["/init"] +CMD ["/command/with-contenv", "/opt/actions-runner/externals/node24/bin/node", "/opt/microvm-lifecycle-hooks/index.js"] +``` + +For an image without s6-overlay, start the artifact with `node /opt/microvm-lifecycle-hooks/index.js` under that image's process supervisor. The hook binds to `0.0.0.0:8080` by default. Restrict the port to the Lambda MicroVM lifecycle network; the protocol does not add a separate application authentication layer. + +## Run payloads + +AWS sends an outer JSON object whose `runHookPayload` is itself a JSON string. Version 1 remains strict and SSM-specific for backwards compatibility: + +```json +{ + "microvmId": "microvm-bdd2d536-3d87-35e4-8b40-18664608ebc1", + "runHookPayload": "{\"version\":1,\"runnerConfigSsmPath\":\"/github-action-runners/example/token\"}" +} +``` + +Version 1 is translated to the shared allowlisted SSM environment. Version 2 carries the exact environment-variable map under `context.storage`. SSM example: + +```json +{ + "version": 2, + "context": { + "storage": { + "RUNNER_CONFIG_STORAGE_PROVIDER": "aws_ssm", + "SSM_TOKEN_PATH": "/github-action-runners/example/token" + } + } +} +``` + +DynamoDB example: + +```json +{ + "version": 2, + "context": { + "storage": { + "RUNNER_CONFIG_STORAGE_PROVIDER": "aws_dynamodb", + "RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME": "github-runner-state" + } + } +} +``` + +Both versions reject missing, unknown, or provider-incompatible fields. Storage context accepts only the two keys for the selected provider; AWS credentials, timeout overrides, and arbitrary environment names are rejected. The validated storage map is exported once before the consumer is resolved. A retry may reuse the identical map, but it cannot change storage configuration after initialization. + +`microvmId` is an opaque path-safe `[A-Za-z0-9_.-]{1,256}` value. The resolved storage provider uses it to consume the one-time JIT configuration. Storage context variables are removed from the runner child environment. + +For DynamoDB, the producer must write the runner configuration with `accessScope` equal to `microvmId`; the consumer atomically removes the unexpired `{ scope: microvmId, id: "config" }` item. + +For a rolling upgrade, keep emitting version 1 SSM payloads until every deployed image contains this service. Old images do not understand version 2. DynamoDB requires the version 2 payload after the image rollout is complete. + +## Entrypoint contract + +On `/run`, the hook starts `${RUNNER_ENTRYPOINT:-/opt/microvm/entrypoint.sh} run` without a shell. It writes this versioned document to stdin: + +```json +{ + "jitConfig": "", + "microvmId": "", + "version": 1 +} +``` + +The entrypoint must write exactly `ready\n` to file descriptor 3 after the runner is ready. The JIT configuration, storage context, and AWS credential environment variables are not passed to the child process. `/terminate` sends `SIGTERM` to the detached process group and escalates to `SIGKILL` after the grace period. + +After the runner entrypoint exits on its own, the hook closes its HTTP server and exits with status `0` only when the runner exited cleanly. In the documented s6-overlay image layout above, that makes the foreground container command exit so s6 can stop the remaining image services and shut down the application container's PID 1. This path does not require `lambda:TerminateMicrovm` in the runner role. AWS documents only explicit termination and maximum duration as MicroVM termination triggers, so retain trusted control-plane cleanup and the maximum duration as failure backstops, and verify the container-exit behavior against a restored MicroVM before relying on it operationally. + +Useful environment variables are: + +| Variable | Default | Purpose | +| --------------------------------- | ---------------------------- | --------------------------------------------- | +| `HOOK_PORT` | `8080` | Lifecycle-hook HTTP port | +| `RUNNER_ENTRYPOINT` | `/opt/microvm/entrypoint.sh` | Image-specific runner supervisor | +| `RUN_HOOK_TIMEOUT_SECONDS` | `55` | Total `/run` budget, bounded to 40–55 seconds | +| `HOOK_HEADERS_TIMEOUT_SECONDS` | `5` | HTTP header receive timeout | +| `HOOK_REQUEST_TIMEOUT_SECONDS` | `10` | HTTP request receive timeout | +| `HOOK_KEEP_ALIVE_TIMEOUT_SECONDS` | `5` | Idle keep-alive timeout | +| `AWS_SDK_CALL_TIMEOUT_SECONDS` | `5` | Individual storage-provider call timeout | +| `RUNNER_CONFIG_TIMEOUT_SECONDS` | `20` | Total runner-configuration polling timeout | +| `RUNNER_CONFIG_POLL_SECONDS` | `2` | Delay between provider polling attempts | +| `RUNNER_CONFIG_DELETE_ATTEMPTS` | `3` | SSM one-time configuration delete attempts | + +The request body is capped at 20 KiB and HTTP headers at 16 KiB. Internal errors are returned generically and secret-bearing provider errors are never logged. + +## Runtime security + +Removing AWS credential and storage variables from the runner child prevents accidental environment inheritance; it is not an IAM boundary. A job can still obtain credentials made available to the runtime role, so scope that role to each lane and treat job code as untrusted. + +- For DynamoDB, grant only `dynamodb:DeleteItem` on the lane's runner-state table. Do not grant `Query` or `Scan`, keep the one-time records on a short TTL, and attach the MicroVM through a `NO_INGRESS` network connector. +- For SSM, grant only `ssm:GetParameter` and `ssm:DeleteParameter` on the lane's token path. Add `kms:Decrypt` only for the customer-managed key that encrypts those parameters. + +## TypeScript API + +The workspace service root is import-safe; importing it does not start the server. It exports the parser, lifecycle, process launcher, storage adapter, and server factories for composition and testing. `src/index.ts` is the executable-only NCC entrypoint. A producer can call `loadRunnerConfigStorageContextFromEnvironment` from `@aws-github-runner/storage-providers/runner-config-consumer` to copy only the selected provider and locator into `context.storage`. diff --git a/lambdas/services/microvm-lifecycle-hooks/package.json b/lambdas/services/microvm-lifecycle-hooks/package.json new file mode 100644 index 0000000000..4add4eb6ae --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/package.json @@ -0,0 +1,53 @@ +{ + "name": "@aws-github-runner/microvm-lifecycle-hooks", + "version": "1.0.0", + "private": true, + "description": "AWS Lambda MicroVM lifecycle hook server for ephemeral GitHub Actions runners", + "main": "src/public.ts", + "exports": { + ".": "./src/public.ts" + }, + "type": "module", + "license": "MIT", + "engines": { + "node": ">=24" + }, + "scripts": { + "start": "node dist/index.js", + "test": "NODE_ENV=test nx test", + "test:watch": "NODE_ENV=test nx test --watch", + "lint": "eslint src", + "build": "ncc build src/index.ts -o dist && cp runtime-package.json dist/package.json", + "format": "prettier --write \"**/*.{ts,json,md}\"", + "format-check": "prettier --check \"**/*.{ts,json,md}\"", + "all": "yarn build && yarn format && yarn lint && yarn test" + }, + "devDependencies": { + "@types/node": "^22.19.3", + "@vercel/ncc": "^0.38.4" + }, + "dependencies": { + "@aws-github-runner/storage-providers": "*" + }, + "nx": { + "targets": { + "build": { + "inputs": [ + "default", + "^default" + ], + "outputs": [ + "{projectRoot}/dist/**/*" + ] + } + }, + "includedScripts": [ + "build", + "format", + "format-check", + "lint", + "start", + "all" + ] + } +} diff --git a/lambdas/services/microvm-lifecycle-hooks/runtime-package.json b/lambdas/services/microvm-lifecycle-hooks/runtime-package.json new file mode 100644 index 0000000000..3dbc1ca591 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/runtime-package.json @@ -0,0 +1,3 @@ +{ + "type": "module" +} diff --git a/lambdas/services/microvm-lifecycle-hooks/src/contracts.ts b/lambdas/services/microvm-lifecycle-hooks/src/contracts.ts new file mode 100644 index 0000000000..9d49daeeae --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/contracts.ts @@ -0,0 +1,39 @@ +import type { RunnerConfigStorageContext } from '@aws-github-runner/storage-providers/runner-config-consumer'; + +export interface RunContext { + microvmId: string; + storage: RunnerConfigStorageContext; +} + +export interface ConsumeOptions { + deadlineMs: number; + signal: AbortSignal; +} + +export interface RunnerBootstrap { + jitConfig: string; +} + +/** Resolves and consumes a one-time runner configuration without exposing provider details. */ +export interface JitConfigSource { + consume(context: RunContext, options: ConsumeOptions): Promise; +} + +export interface ManagedProcess { + readonly ready: Promise; + readonly exit: Promise; + readonly exited: boolean; + stop(graceMs?: number): Promise; +} + +export interface RunnerLauncher { + launch(bootstrap: RunnerBootstrap, microvmId: string): ManagedProcess; +} + +export interface Logger { + info(message: string, ...values: unknown[]): void; + warn(message: string, ...values: unknown[]): void; + error(message: string, ...values: unknown[]): void; +} + +export const consoleLogger: Logger = console; diff --git a/lambdas/services/microvm-lifecycle-hooks/src/index.ts b/lambdas/services/microvm-lifecycle-hooks/src/index.ts new file mode 100644 index 0000000000..30cde16d7b --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/index.ts @@ -0,0 +1,7 @@ +import { consoleLogger } from './contracts'; +import { main } from './server'; + +void main().catch(() => { + consoleLogger.error('Lambda MicroVM lifecycle hook server failed to start'); + process.exitCode = 1; +}); diff --git a/lambdas/services/microvm-lifecycle-hooks/src/lifecycle.test.ts b/lambdas/services/microvm-lifecycle-hooks/src/lifecycle.test.ts new file mode 100644 index 0000000000..ea7bc02e73 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/lifecycle.test.ts @@ -0,0 +1,214 @@ +import type { JitConfigSource, Logger, ManagedProcess, RunContext, RunnerBootstrap, RunnerLauncher } from './contracts'; +import { RunnerLifecycle } from './lifecycle'; + +const quietLogger: Logger = { + error: () => undefined, + info: () => undefined, + warn: () => undefined, +}; + +const MICROVM_ID = 'microvm-bdd2d536-3d87-35e4-8b40-18664608ebc1'; + +afterEach(() => { + vi.restoreAllMocks(); +}); + +function runRequest(): string { + return JSON.stringify({ + microvmId: MICROVM_ID, + runHookPayload: JSON.stringify({ + runnerConfigSsmPath: '/runner/token', + version: 1, + }), + }); +} + +class DeferredProcess implements ManagedProcess { + public readonly ready = Promise.resolve(); + public readonly exit: Promise; + public exited = false; + private resolveExit!: (code: number | null) => void; + + public constructor() { + this.exit = new Promise((resolve) => { + this.resolveExit = resolve; + }); + } + + public finish(code: number | null): void { + this.exited = true; + this.resolveExit(code); + } + + public async stop(): Promise { + if (!this.exited) { + this.finish(null); + } + } +} + +describe('RunnerLifecycle', () => { + it('starts only once and waits for terminate cleanup after the runner exits', async () => { + const events: string[] = []; + const processHandle = new DeferredProcess(); + const source: JitConfigSource = { + async consume(context: RunContext): Promise { + events.push(`consume:${context.storage.RUNNER_CONFIG_STORAGE_PROVIDER}:${context.microvmId}`); + return { jitConfig: 'encoded-jit' }; + }, + }; + const launcher: RunnerLauncher = { + launch(bootstrap, id): ManagedProcess { + events.push(`launch:${id}:${bootstrap.jitConfig}`); + return processHandle; + }, + }; + const lifecycle = new RunnerLifecycle(source, launcher, quietLogger); + + await expect(lifecycle.start(runRequest())).resolves.toBe(true); + await expect(lifecycle.start(runRequest())).resolves.toBe(false); + expect(events).toEqual([`consume:aws_ssm:${MICROVM_ID}`, `launch:${MICROVM_ID}:encoded-jit`]); + + processHandle.finish(0); + await expect(lifecycle.completion).resolves.toBe(0); + await lifecycle.stop(); + expect(processHandle.exited).toBe(true); + }); + + it('does not report an externally requested stop as runner self-completion', async () => { + const processHandle = new DeferredProcess(); + const lifecycle = new RunnerLifecycle( + { consume: async () => ({ jitConfig: 'encoded-jit' }) }, + { launch: () => processHandle }, + quietLogger, + ); + + await lifecycle.start(runRequest()); + await lifecycle.stop(); + + await expect( + Promise.race([ + lifecycle.completion.then(() => 'completed'), + new Promise((resolve) => setImmediate(() => resolve('pending'))), + ]), + ).resolves.toBe('pending'); + }); + + it('reserves the runner startup budget before consuming configuration', async () => { + let consumeDeadline = 0; + const processHandle = new DeferredProcess(); + const lifecycle = new RunnerLifecycle( + { + async consume(_context, options): Promise { + consumeDeadline = options.deadlineMs; + return { jitConfig: 'encoded-jit' }; + }, + }, + { launch: () => processHandle }, + quietLogger, + ); + vi.spyOn(Date, 'now').mockReturnValue(1_000); + + await lifecycle.start(runRequest()); + + expect(consumeDeadline).toBe(21_000); + await lifecycle.stop(); + }); + + it('returns to idle if the configured entrypoint cannot launch', async () => { + const consume = vi.fn().mockResolvedValue({ jitConfig: 'encoded-jit' }); + const lifecycle = new RunnerLifecycle( + { consume }, + { + launch(): ManagedProcess { + throw new Error('spawn failed'); + }, + }, + quietLogger, + ); + + await expect(lifecycle.start(runRequest())).rejects.toThrow('spawn failed'); + await expect(lifecycle.start(runRequest())).rejects.toThrow('spawn failed'); + expect(consume).toHaveBeenCalledTimes(2); + }); + + it('aborts in-flight consumption when the run-hook deadline elapses', async () => { + let consumedSignal: AbortSignal | undefined; + let launched = false; + let releaseConsume = (): void => undefined; + const consumption = new Promise((resolve) => { + releaseConsume = resolve; + }); + const lifecycle = new RunnerLifecycle( + { + async consume(_context, options): Promise { + consumedSignal = options.signal; + await consumption; + return { jitConfig: 'encoded-jit' }; + }, + }, + { + launch(): ManagedProcess { + launched = true; + return new DeferredProcess(); + }, + }, + quietLogger, + ); + let calls = 0; + vi.spyOn(Date, 'now').mockImplementation(() => (calls++ === 0 ? 1_000 : 61_000)); + + await expect(lifecycle.start(runRequest())).rejects.toThrow('run-hook deadline elapsed'); + releaseConsume(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(consumedSignal?.aborted).toBe(true); + expect(launched).toBe(false); + }); + + it('waits for cleanup when terminate races with entrypoint readiness', async () => { + let finishCleanup = (): void => undefined; + let reportLaunched = (): void => undefined; + let stopCalled = false; + const cleanup = new Promise((resolve) => { + finishCleanup = resolve; + }); + const launched = new Promise((resolve) => { + reportLaunched = resolve; + }); + const processHandle: ManagedProcess = { + ready: new Promise(() => undefined), + exit: new Promise(() => undefined), + exited: false, + async stop(): Promise { + stopCalled = true; + await cleanup; + }, + }; + const lifecycle = new RunnerLifecycle( + { consume: async () => ({ jitConfig: 'encoded-jit' }) }, + { + launch(): ManagedProcess { + reportLaunched(); + return processHandle; + }, + }, + quietLogger, + ); + + const rejectedStart = expect(lifecycle.start(runRequest())).rejects.toThrow('runner start was cancelled'); + await launched; + let terminateSettled = false; + const terminate = lifecycle.stop().then(() => { + terminateSettled = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(stopCalled).toBe(true); + expect(terminateSettled).toBe(false); + + finishCleanup(); + await terminate; + await rejectedStart; + expect(terminateSettled).toBe(true); + }); +}); diff --git a/lambdas/services/microvm-lifecycle-hooks/src/lifecycle.ts b/lambdas/services/microvm-lifecycle-hooks/src/lifecycle.ts new file mode 100644 index 0000000000..4eae1aa091 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/lifecycle.ts @@ -0,0 +1,161 @@ +import type { JitConfigSource, Logger, ManagedProcess, RunnerBootstrap, RunnerLauncher } from './contracts'; +import { consoleLogger } from './contracts'; +import { parseRunRequest } from './payload'; +import { beforeDeadline, beforeDeadlineOrAbort } from './timing'; + +type LifecycleState = 'idle' | 'starting' | 'running' | 'stopping' | 'stopped'; + +function boundedNumber(value: string | undefined, fallback: number, minimum: number, maximum: number): number { + const parsed = Number(value); + return Number.isFinite(parsed) ? Math.max(minimum, Math.min(maximum, parsed)) : fallback; +} + +export class RunnerLifecycle { + private readonly runHookBudgetMs = boundedNumber(process.env.RUN_HOOK_TIMEOUT_SECONDS, 55, 40, 55) * 1_000; + // Reserve Lambda's 30-second service readiness window plus five seconds of local margin. + private readonly launchReserveMs = 35_000; + private state: LifecycleState = 'idle'; + private microvmId?: string; + private startAbort?: AbortController; + private startPromise?: Promise; + private runner?: ManagedProcess; + private resolveCompletion!: (exitCode: number | null) => void; + public readonly completion = new Promise((resolve) => { + this.resolveCompletion = resolve; + }); + + public constructor( + private readonly jitConfigSource: JitConfigSource, + private readonly launcher: RunnerLauncher, + private readonly logger: Logger = consoleLogger, + ) {} + + private currentState(): LifecycleState { + return this.state; + } + + public async start(body: string): Promise { + const context = parseRunRequest(body); + const deadlineMs = Date.now() + this.runHookBudgetMs; + + if (this.microvmId === context.microvmId && this.state === 'running') { + return false; + } + if (this.microvmId === context.microvmId && this.state === 'starting') { + if (this.startPromise === undefined) { + throw new Error('runner start state is inconsistent'); + } + await beforeDeadline(this.startPromise, deadlineMs); + if (this.currentState() === 'running') { + return false; + } + throw new Error('the preceding runner start did not succeed'); + } + if (this.state !== 'idle') { + throw new Error('another runner lifecycle is already active in this MicroVM'); + } + + const abort = new AbortController(); + this.state = 'starting'; + this.microvmId = context.microvmId; + this.startAbort = abort; + const startOperation = this.startRunner(context, deadlineMs, abort); + this.startPromise = startOperation; + const clearStartPromise = (): void => { + if (this.startPromise === startOperation) { + this.startPromise = undefined; + } + }; + void startOperation.then(clearStartPromise, clearStartPromise); + try { + await beforeDeadline(startOperation, deadlineMs); + return true; + } catch (error) { + // Cancel the underlying work so a timed-out hook cannot register a runner later. + abort.abort(); + throw error; + } + } + + private async startRunner( + context: ReturnType, + deadlineMs: number, + abort: AbortController, + ): Promise { + let bootstrap: RunnerBootstrap | undefined; + let processHandle: ManagedProcess | undefined; + try { + bootstrap = await this.jitConfigSource.consume(context, { + deadlineMs: deadlineMs - this.launchReserveMs, + signal: abort.signal, + }); + if (abort.signal.aborted) { + throw new Error('runner start was cancelled'); + } + + processHandle = this.launcher.launch(bootstrap, context.microvmId); + await beforeDeadlineOrAbort(processHandle.ready, deadlineMs, abort.signal); + if (abort.signal.aborted || this.state !== 'starting') { + throw new Error('runner start was cancelled'); + } + + this.runner = processHandle; + this.startAbort = undefined; + this.state = 'running'; + this.logger.info('GitHub Actions runner started for MicroVM %s', context.microvmId); + void this.monitorRunner(processHandle); + } catch (error) { + if (processHandle !== undefined) { + await processHandle.stop(); + } + if (this.state === 'stopping') { + this.state = 'stopped'; + } else { + this.state = 'idle'; + this.microvmId = undefined; + } + this.startAbort = undefined; + throw error; + } finally { + // JavaScript strings cannot be zeroized, but release the retained credential promptly. + if (bootstrap !== undefined) { + bootstrap.jitConfig = ''; + } + } + } + + private async monitorRunner(processHandle: ManagedProcess): Promise { + const exitCode = await processHandle.exit; + if (this.runner === processHandle) { + this.runner = undefined; + this.state = 'stopped'; + this.resolveCompletion(exitCode); + } + } + + public async stop(): Promise { + if (this.state === 'idle') { + this.state = 'stopped'; + } else if (this.state === 'starting' || this.state === 'running') { + this.state = 'stopping'; + } + this.startAbort?.abort(); + const starting = this.startPromise; + if (starting !== undefined) { + try { + await starting; + } catch { + // Cancellation is expected when terminate races with /run. + } + } + const running = this.runner; + this.runner = undefined; + await (running?.stop() ?? Promise.resolve()); + this.state = 'stopped'; + } + + public async resume(): Promise { + // Never re-consume a one-time runner configuration on resume. + return true; + } +} diff --git a/lambdas/services/microvm-lifecycle-hooks/src/payload.test.ts b/lambdas/services/microvm-lifecycle-hooks/src/payload.test.ts new file mode 100644 index 0000000000..3efa3814b6 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/payload.test.ts @@ -0,0 +1,124 @@ +import { HookRequestError, parseRunRequest } from './payload'; + +const MICROVM_ID = 'microvm-bdd2d536-3d87-35e4-8b40-18664608ebc1'; +const SSM_STORAGE = { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', + SSM_TOKEN_PATH: '/github-action-runners/tenant/token', +} as const; +const DYNAMODB_STORAGE = { + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'github-runner-config', + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb', +} as const; + +function request( + payload: object = { + runnerConfigSsmPath: '/github-action-runners/tenant/token', + version: 1, + }, + microvmId = MICROVM_ID, +): string { + return JSON.stringify({ + microvmId, + runHookPayload: JSON.stringify(payload), + }); +} + +describe('parseRunRequest', () => { + it('maps the strict version 1 payload to the allowlisted SSM storage environment', () => { + expect(parseRunRequest(request())).toEqual({ + microvmId: MICROVM_ID, + storage: SSM_STORAGE, + }); + }); + + it('preserves version 1 trailing-slash normalization', () => { + expect( + parseRunRequest( + request({ + runnerConfigSsmPath: '/github-action-runners/tenant/token/', + version: 1, + }), + ).storage, + ).toEqual({ + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', + SSM_TOKEN_PATH: '/github-action-runners/tenant/token', + }); + }); + + it.each([SSM_STORAGE, DYNAMODB_STORAGE])( + 'accepts a strict version 2 $RUNNER_CONFIG_STORAGE_PROVIDER storage context', + (storage) => { + expect( + parseRunRequest( + request({ + context: { storage }, + version: 2, + }), + ), + ).toEqual({ microvmId: MICROVM_ID, storage }); + }, + ); + + it('accepts opaque path-safe MicroVM identifiers up to 256 characters', () => { + expect(parseRunRequest(request(undefined, 'a'.repeat(256))).microvmId).toHaveLength(256); + expect(parseRunRequest(request(undefined, 'future_id.example-01')).microvmId).toBe('future_id.example-01'); + }); + + it.each([ + ['invalid outer JSON', '{'], + ['an invalid MicroVM identifier', request(undefined, '../vm')], + ['an overlong MicroVM identifier', request(undefined, 'a'.repeat(257))], + ['an unversioned payload', request({ runnerConfigSsmPath: '/runner/token' })], + ['a relative legacy SSM path', request({ runnerConfigSsmPath: 'runner/token', version: 1 })], + ['a root legacy SSM path', request({ runnerConfigSsmPath: '/', version: 1 })], + ['repeated legacy SSM slashes', request({ runnerConfigSsmPath: '/runner//token', version: 1 })], + ['legacy SSM traversal', request({ runnerConfigSsmPath: '/runner/../token', version: 1 })], + [ + 'extra version 1 fields', + request({ encodedJitConfig: 'not-a-real-secret', runnerConfigSsmPath: '/runner/token', version: 1 }), + ], + [ + 'missing version 1 fields', + request({ context: { storage: SSM_STORAGE }, runnerConfigSsmPath: '/runner/token', version: 1 }), + ], + ['missing version 2 context', request({ version: 2 })], + ['missing version 2 storage', request({ context: {}, version: 2 })], + ['extra version 2 fields', request({ context: { storage: SSM_STORAGE }, unexpected: true, version: 2 })], + ['extra version 2 context fields', request({ context: { storage: SSM_STORAGE, unexpected: true }, version: 2 })], + [ + 'an unknown storage provider', + request({ + context: { + storage: { RUNNER_CONFIG_STORAGE_PROVIDER: 'unknown', SSM_TOKEN_PATH: '/runner/token' }, + }, + version: 2, + }), + ], + [ + 'provider-incompatible storage fields', + request({ + context: { + storage: { + ...SSM_STORAGE, + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'runner-config', + }, + }, + version: 2, + }), + ], + [ + 'typed provider fields in the environment map', + request({ context: { storage: { provider: 'aws_ssm', tokenPath: '/runner/token' } }, version: 2 }), + ], + [ + 'AWS credential injection', + request({ context: { storage: { ...SSM_STORAGE, AWS_ACCESS_KEY_ID: 'not-a-real-key' } }, version: 2 }), + ], + [ + 'timeout override injection', + request({ context: { storage: { ...SSM_STORAGE, RUNNER_CONFIG_TIMEOUT_SECONDS: '60' } }, version: 2 }), + ], + ])('rejects %s', (_name, body) => { + expect(() => parseRunRequest(body)).toThrow(HookRequestError); + }); +}); diff --git a/lambdas/services/microvm-lifecycle-hooks/src/payload.ts b/lambdas/services/microvm-lifecycle-hooks/src/payload.ts new file mode 100644 index 0000000000..5dcc60daf9 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/payload.ts @@ -0,0 +1,101 @@ +import { + parseRunnerConfigStorageContext, + type RunnerConfigStorageContext, +} from '@aws-github-runner/storage-providers/runner-config-consumer'; + +import type { RunContext } from './contracts'; + +const MICROVM_ID_PATTERN = /^[A-Za-z0-9_.-]{1,256}$/; + +export const MAX_REQUEST_BYTES = 20 * 1024; + +export class HookRequestError extends Error { + public constructor(message: string) { + super(message); + this.name = 'HookRequestError'; + } +} + +interface LambdaRunRequest { + microvmId?: unknown; + runHookPayload?: unknown; +} + +interface VersionedRunPayload { + version?: unknown; + runnerConfigSsmPath?: unknown; + context?: unknown; +} + +interface VersionTwoContext { + storage?: unknown; +} + +function parseObject(value: string, errorMessage: string): T { + let parsed: unknown; + try { + parsed = JSON.parse(value); + } catch { + throw new HookRequestError(errorMessage); + } + if (parsed === null || typeof parsed !== 'object' || Array.isArray(parsed)) { + throw new HookRequestError(errorMessage); + } + return parsed as T; +} + +function hasExactKeys(value: object, expected: readonly string[]): boolean { + const keys = Object.keys(value); + return keys.length === expected.length && keys.every((key) => expected.includes(key)); +} + +function isObject(value: unknown): value is object { + return value !== null && typeof value === 'object' && !Array.isArray(value); +} + +function parseStorageContext(value: unknown): RunnerConfigStorageContext { + try { + return parseRunnerConfigStorageContext(value); + } catch { + // Storage validation details are deliberately not reflected to the hook caller. + throw new HookRequestError('runner configuration storage context is missing or invalid'); + } +} + +export function parseRunRequest(body: string): RunContext { + const request = parseObject(body, 'request body must be a JSON object'); + if (typeof request.microvmId !== 'string' || !MICROVM_ID_PATTERN.test(request.microvmId)) { + throw new HookRequestError('microvmId is missing or invalid'); + } + if (typeof request.runHookPayload !== 'string') { + throw new HookRequestError('runHookPayload must be a JSON string'); + } + + const payload = parseObject(request.runHookPayload, 'runHookPayload must contain valid JSON'); + if (payload.version === 1) { + if (!hasExactKeys(payload, ['version', 'runnerConfigSsmPath'])) { + throw new HookRequestError('version 1 runHookPayload contains unsupported or missing fields'); + } + return { + microvmId: request.microvmId, + storage: parseStorageContext({ + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', + SSM_TOKEN_PATH: payload.runnerConfigSsmPath, + }), + }; + } + if (payload.version === 2) { + if (!hasExactKeys(payload, ['version', 'context'])) { + throw new HookRequestError('version 2 runHookPayload contains unsupported or missing fields'); + } + if (!isObject(payload.context) || !hasExactKeys(payload.context, ['storage'])) { + throw new HookRequestError('version 2 context contains unsupported or missing fields'); + } + const context = payload.context as VersionTwoContext; + return { + microvmId: request.microvmId, + storage: parseStorageContext(context.storage), + }; + } + throw new HookRequestError('runHookPayload version must be 1 or 2'); +} diff --git a/lambdas/services/microvm-lifecycle-hooks/src/processes.test.ts b/lambdas/services/microvm-lifecycle-hooks/src/processes.test.ts new file mode 100644 index 0000000000..01837af8ed --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/processes.test.ts @@ -0,0 +1,109 @@ +import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; + +import { RunnerEntrypointLauncher } from './processes'; + +afterEach(() => { + vi.unstubAllEnvs(); +}); + +describe('RunnerEntrypointLauncher', () => { + it('passes the MicroVM id and one-time JIT only through stdin', async () => { + const directory = await mkdtemp(join(tmpdir(), 'microvm-entrypoint-')); + const entrypoint = join(directory, 'entrypoint.sh'); + const output = join(directory, 'output'); + const environmentOutput = join(directory, 'environment-output'); + + await writeFile( + entrypoint, + `#!/bin/sh +set -eu +case "$1" in + run) + cat > "$TEST_ENTRYPOINT_OUTPUT" + printf '%s|%s|%s|%s|%s|%s|%s|%s|%s|%s|%s|%s|%s' \ + "\${ENCODED_JIT_CONFIG-unset}" \ + "\${AWS_ACCESS_KEY_ID-unset}" \ + "\${AWS_SESSION_TOKEN-unset}" \ + "\${AWS_CONTAINER_CREDENTIALS_FULL_URI-unset}" \ + "\${AWS_PROFILE-unset}" \ + "\${AWS_DEFAULT_PROFILE-unset}" \ + "\${AWS_CONFIG_FILE-unset}" \ + "\${AWS_SHARED_CREDENTIALS_FILE-unset}" \ + "\${AWS_CREDENTIAL_EXPIRATION-unset}" \ + "\${RUNNER_CONFIG_STORAGE_PROVIDER-unset}" \ + "\${SSM_TOKEN_PATH-unset}" \ + "\${RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME-unset}" \ + "\${RUNNER_ALLOW_RUNASROOT-unset}" > "$TEST_ENTRYPOINT_ENV_OUTPUT" + printf 'ready\n' >&3 + ;; + *) exit 2 ;; +esac +`, + { mode: 0o700 }, + ); + + vi.stubEnv('RUNNER_ENTRYPOINT', entrypoint); + vi.stubEnv('TEST_ENTRYPOINT_OUTPUT', output); + vi.stubEnv('TEST_ENTRYPOINT_ENV_OUTPUT', environmentOutput); + vi.stubEnv('ENCODED_JIT_CONFIG', 'test-value'); + vi.stubEnv('AWS_ACCESS_KEY_ID', 'test-value'); + vi.stubEnv('AWS_SESSION_TOKEN', 'test-value'); + vi.stubEnv('AWS_CONTAINER_CREDENTIALS_FULL_URI', 'http://127.0.0.1/credentials'); + vi.stubEnv('AWS_PROFILE', 'test-profile'); + vi.stubEnv('AWS_DEFAULT_PROFILE', 'test-profile'); + vi.stubEnv('AWS_CONFIG_FILE', '/tmp/test-config'); + vi.stubEnv('AWS_SHARED_CREDENTIALS_FILE', '/tmp/test-credentials'); + vi.stubEnv('AWS_CREDENTIAL_EXPIRATION', '2099-01-01T00:00:00Z'); + vi.stubEnv('RUNNER_CONFIG_STORAGE_PROVIDER', 'aws_dynamodb'); + vi.stubEnv('SSM_TOKEN_PATH', '/runner/token'); + vi.stubEnv('RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME', 'runner-config'); + vi.stubEnv('RUNNER_ALLOW_RUNASROOT', '1'); + try { + const processHandle = new RunnerEntrypointLauncher().launch({ jitConfig: 'encoded-jit' }, 'mvm-1234'); + + await processHandle.ready; + await expect(processHandle.exit).resolves.toBe(0); + expect(JSON.parse(await readFile(output, 'utf8'))).toEqual({ + jitConfig: 'encoded-jit', + microvmId: 'mvm-1234', + version: 1, + }); + expect(await readFile(environmentOutput, 'utf8')).toBe( + 'unset|unset|unset|unset|unset|unset|unset|unset|unset|unset|unset|unset|unset', + ); + } finally { + await rm(directory, { force: true, recursive: true }); + } + }); + + it('requires the entrypoint to signal readiness before it exits', async () => { + const directory = await mkdtemp(join(tmpdir(), 'microvm-entrypoint-')); + const entrypoint = join(directory, 'entrypoint.sh'); + await writeFile( + entrypoint, + `#!/bin/sh +set -eu +case "$1" in + run) + cat >/dev/null + exit 7 + ;; + *) exit 2 ;; +esac +`, + { mode: 0o700 }, + ); + + vi.stubEnv('RUNNER_ENTRYPOINT', entrypoint); + try { + const processHandle = new RunnerEntrypointLauncher().launch({ jitConfig: 'encoded-jit' }, 'mvm-1234'); + + await expect(processHandle.ready).rejects.toThrow('exited before signaling readiness'); + await expect(processHandle.exit).resolves.toBe(7); + } finally { + await rm(directory, { force: true, recursive: true }); + } + }); +}); diff --git a/lambdas/services/microvm-lifecycle-hooks/src/processes.ts b/lambdas/services/microvm-lifecycle-hooks/src/processes.ts new file mode 100644 index 0000000000..881d58211d --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/processes.ts @@ -0,0 +1,188 @@ +import { type ChildProcess, spawn } from 'node:child_process'; +import { Readable } from 'node:stream'; + +import type { ManagedProcess, RunnerBootstrap, RunnerLauncher } from './contracts'; +import { delay } from './timing'; + +const CREDENTIAL_ENVIRONMENT_VARIABLES = [ + 'AWS_ACCESS_KEY_ID', + 'AWS_CONFIG_FILE', + 'AWS_CONTAINER_AUTHORIZATION_TOKEN', + 'AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE', + 'AWS_CONTAINER_CREDENTIALS_FULL_URI', + 'AWS_CONTAINER_CREDENTIALS_RELATIVE_URI', + 'AWS_CREDENTIAL_EXPIRATION', + 'AWS_DEFAULT_PROFILE', + 'AWS_PROFILE', + 'AWS_ROLE_ARN', + 'AWS_SECRET_ACCESS_KEY', + 'AWS_SECURITY_TOKEN', + 'AWS_SHARED_CREDENTIALS_FILE', + 'AWS_SESSION_TOKEN', + 'AWS_WEB_IDENTITY_TOKEN_FILE', + 'ENCODED_JIT_CONFIG', + 'RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME', + 'RUNNER_CONFIG_STORAGE_PROVIDER', + 'RUNNER_ALLOW_RUNASROOT', + 'SSM_TOKEN_PATH', +] as const; + +function signalProcessGroup(child: ChildProcess, signal: NodeJS.Signals): void { + if (child.pid === undefined || child.exitCode !== null || child.signalCode !== null) { + return; + } + try { + process.kill(-child.pid, signal); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'ESRCH') { + child.kill(signal); + } + } +} + +export class NodeManagedProcess implements ManagedProcess { + public readonly ready: Promise; + public readonly exit: Promise; + + public constructor( + private readonly child: ChildProcess, + readiness: Promise, + private readonly defaultStopGraceMs: number, + ) { + this.ready = readiness; + this.exit = new Promise((resolve) => { + child.once('exit', (code) => resolve(code)); + child.once('error', () => resolve(null)); + }); + } + + public get exited(): boolean { + return this.child.exitCode !== null || this.child.signalCode !== null; + } + + public async stop(graceMs = this.defaultStopGraceMs): Promise { + if (this.exited) { + return; + } + signalProcessGroup(this.child, 'SIGTERM'); + const exitedGracefully = await Promise.race([this.exit.then(() => true), delay(graceMs).then(() => false)]); + if (!exitedGracefully && !this.exited) { + signalProcessGroup(this.child, 'SIGKILL'); + await Promise.race([this.exit, delay(5_000)]); + } + } +} + +function entrypointEnvironment(microvmId: string): NodeJS.ProcessEnv { + const environment = { ...process.env }; + for (const variable of CREDENTIAL_ENVIRONMENT_VARIABLES) { + delete environment[variable]; + } + return { + ...environment, + MICROVM_ID: microvmId, + }; +} + +function waitForEntrypointReady(child: ChildProcess): Promise { + const candidate = child.stdio[3]; + if (!(candidate instanceof Readable)) { + return Promise.reject(new Error('runner entrypoint readiness pipe is unavailable')); + } + const readinessStream: Readable = candidate; + readinessStream.setEncoding('utf8'); + + return new Promise((resolve, reject) => { + let buffer = ''; + let settled = false; + + function cleanup(): void { + readinessStream.off('data', onData); + readinessStream.off('end', onEnd); + readinessStream.off('error', onError); + child.off('error', onError); + } + + function succeed(): void { + if (!settled) { + settled = true; + cleanup(); + resolve(); + } + } + + function fail(error: Error): void { + if (!settled) { + settled = true; + cleanup(); + reject(error); + } + } + + function onData(chunk: string | Buffer): void { + buffer += chunk.toString(); + if (buffer === 'ready\n') { + succeed(); + } else if (buffer.includes('\n') || buffer.length > 64) { + fail(new Error('runner entrypoint emitted an invalid readiness signal')); + } + } + + function onEnd(): void { + fail(new Error('runner entrypoint exited before signaling readiness')); + } + + function onError(error: Error): void { + fail(error); + } + + readinessStream.on('data', onData); + readinessStream.once('end', onEnd); + readinessStream.once('error', onError); + child.once('error', onError); + }); +} + +/** + * Sends the one-time JIT document through stdin to an image-specific supervisor. + * Neither the JIT document nor storage-provider credentials are exported to the runner. + */ +export class RunnerEntrypointLauncher implements RunnerLauncher { + private readonly entrypoint = process.env.RUNNER_ENTRYPOINT ?? '/opt/microvm/entrypoint.sh'; + + public constructor(private readonly stopGraceMs = 30_000) {} + + public launch(bootstrap: RunnerBootstrap, microvmId: string): ManagedProcess { + const child = spawn(this.entrypoint, ['run'], { + detached: true, + env: entrypointEnvironment(microvmId), + stdio: ['pipe', 'inherit', 'inherit', 'pipe'], + }); + const entrypointReady = waitForEntrypointReady(child); + const inputWritten = new Promise((resolve, reject) => { + const fail = (error: Error): void => reject(error); + child.once('error', fail); + child.once('spawn', () => { + if (child.stdin === null) { + reject(new Error('runner entrypoint stdin is unavailable')); + return; + } + child.stdin.once('error', fail); + child.stdin.end( + JSON.stringify({ + jitConfig: bootstrap.jitConfig, + microvmId, + version: 1, + }), + () => { + child.removeListener('error', fail); + child.stdin?.removeListener('error', fail); + resolve(); + }, + ); + }); + }); + const ready = Promise.all([inputWritten, entrypointReady]).then(() => undefined); + return new NodeManagedProcess(child, ready, this.stopGraceMs); + } +} diff --git a/lambdas/services/microvm-lifecycle-hooks/src/public.ts b/lambdas/services/microvm-lifecycle-hooks/src/public.ts new file mode 100644 index 0000000000..32b385b15c --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/public.ts @@ -0,0 +1,26 @@ +export type { + ConsumeOptions, + JitConfigSource, + Logger, + ManagedProcess, + RunContext, + RunnerBootstrap, + RunnerLauncher, +} from './contracts'; +export { consoleLogger } from './contracts'; +export { RunnerLifecycle } from './lifecycle'; +export { HookRequestError, MAX_REQUEST_BYTES, parseRunRequest } from './payload'; +export { NodeManagedProcess, RunnerEntrypointLauncher } from './processes'; +export { + createHookExitRequester, + createDefaultLifecycle, + createHookServer, + HOOK_PREFIX, + main, + parsePositiveInteger, + shutdownHookServer, + watchRunnerCompletion, +} from './server'; +export type { HookLifecycle, HookServerOptions } from './server'; +export { StorageJitConfigSource } from './storage'; +export type { StorageJitConfigSourceOptions } from './storage'; diff --git a/lambdas/services/microvm-lifecycle-hooks/src/server.test.ts b/lambdas/services/microvm-lifecycle-hooks/src/server.test.ts new file mode 100644 index 0000000000..823955d755 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/server.test.ts @@ -0,0 +1,210 @@ +import type { AddressInfo } from 'node:net'; + +import type { Logger } from './contracts'; +import { + createHookExitRequester, + createHookServer, + type HookLifecycle, + HOOK_PREFIX, + parsePositiveInteger, + shutdownHookServer, + watchRunnerCompletion, +} from './server'; + +const quietLogger: Logger = { + error: () => undefined, + info: () => undefined, + warn: () => undefined, +}; + +const idleLifecycle: HookLifecycle = { + resume: async () => true, + start: async () => true, + stop: async () => undefined, +}; + +async function listen(server: ReturnType): Promise { + await new Promise((resolve, reject) => { + server.once('error', reject); + server.listen(0, '127.0.0.1', resolve); + }); + const address = server.address() as AddressInfo; + return `http://127.0.0.1:${address.port}`; +} + +async function close(server: ReturnType): Promise { + await new Promise((resolve, reject) => { + server.close((error) => (error === undefined ? resolve() : reject(error))); + server.closeAllConnections(); + }); +} + +describe('hook server', () => { + it('rejects invalid and out-of-range positive integer values', () => { + expect(parsePositiveInteger(undefined, 8080, 65_535)).toBe(8080); + expect(parsePositiveInteger('', 8080, 65_535)).toBe(8080); + expect(parsePositiveInteger('0', 8080, 65_535)).toBe(8080); + expect(parsePositiveInteger('-1', 8080, 65_535)).toBe(8080); + expect(parsePositiveInteger('1.5', 8080, 65_535)).toBe(8080); + expect(parsePositiveInteger('8080http', 8080, 65_535)).toBe(8080); + expect(parsePositiveInteger('65536', 8080, 65_535)).toBe(8080); + expect(parsePositiveInteger('9007199254740992', 8080)).toBe(8080); + expect(parsePositiveInteger('9090', 8080, 65_535)).toBe(9090); + }); + + it('configures bounded request, header, connection, and socket limits', () => { + const server = createHookServer(idleLifecycle, quietLogger, { + headersTimeoutMs: 2_000, + keepAliveTimeoutMs: 3_000, + requestTimeoutMs: 4_000, + }); + + expect(server.headersTimeout).toBe(2_000); + expect(server.keepAliveTimeout).toBe(3_000); + expect(server.requestTimeout).toBe(4_000); + expect(server.maxConnections).toBe(128); + expect(server.maxHeadersCount).toBe(64); + expect(server.maxRequestsPerSocket).toBe(100); + }); + + it('acknowledges build hooks without starting a runner', async () => { + const lifecycle: HookLifecycle = { + resume: vi.fn(), + start: vi.fn(), + stop: vi.fn(), + }; + const server = createHookServer(lifecycle, quietLogger); + const baseUrl = await listen(server); + try { + const ready = await fetch(`${baseUrl}${HOOK_PREFIX}/ready`, { method: 'POST' }); + const validate = await fetch(`${baseUrl}${HOOK_PREFIX}/validate`, { method: 'POST' }); + + expect(ready.status).toBe(200); + await expect(ready.json()).resolves.toEqual({ status: 'ready' }); + expect(validate.status).toBe(200); + await expect(validate.json()).resolves.toEqual({ status: 'validated' }); + expect(lifecycle.start).not.toHaveBeenCalled(); + expect(lifecycle.stop).not.toHaveBeenCalled(); + } finally { + await close(server); + } + }); + + it('rejects oversized request bodies before invoking the lifecycle', async () => { + const lifecycle: HookLifecycle = { + ...idleLifecycle, + start: vi.fn(), + }; + const server = createHookServer(lifecycle, quietLogger); + const baseUrl = await listen(server); + try { + const response = await fetch(`${baseUrl}${HOOK_PREFIX}/run`, { + body: 'x'.repeat(20 * 1024 + 1), + method: 'POST', + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual({ error: 'request body is too large' }); + expect(lifecycle.start).not.toHaveBeenCalled(); + } finally { + await close(server); + } + }); + + it('does not reflect or log secret-bearing internal errors', async () => { + const messages: unknown[] = []; + const logger: Logger = { + error: (...values) => messages.push(...values), + info: () => undefined, + warn: () => undefined, + }; + const lifecycle: HookLifecycle = { + ...idleLifecycle, + start: async () => { + const error = new Error('encoded-jit-secret'); + error.name = 'encoded-jit-secret'; + throw error; + }, + }; + const server = createHookServer(lifecycle, logger); + const baseUrl = await listen(server); + try { + const response = await fetch(`${baseUrl}${HOOK_PREFIX}/run`, { + body: '{}', + method: 'POST', + }); + + expect(response.status).toBe(500); + await expect(response.json()).resolves.toEqual({ error: 'lifecycle hook failed' }); + expect(JSON.stringify(messages)).not.toContain('encoded-jit-secret'); + } finally { + await close(server); + } + }); + + it('waits for lifecycle cleanup before closing active connections', async () => { + const events: string[] = []; + let finishCleanup = (): void => undefined; + const cleanup = new Promise((resolve) => { + finishCleanup = resolve; + }); + const server = { + close: () => events.push('stop-accepting'), + closeAllConnections: () => events.push('close-connections'), + }; + const lifecycle = { + async stop(): Promise { + events.push('cleanup-started'); + await cleanup; + events.push('cleanup-finished'); + }, + }; + + const shutdown = shutdownHookServer(server, lifecycle); + await new Promise((resolve) => setImmediate(resolve)); + expect(events).toEqual(['stop-accepting', 'cleanup-started']); + + finishCleanup(); + await shutdown; + expect(events).toEqual(['stop-accepting', 'cleanup-started', 'cleanup-finished', 'close-connections']); + }); + + it.each([ + { expectedExitCode: 0, runnerExitCode: 0 }, + { expectedExitCode: 1, runnerExitCode: 7 }, + { expectedExitCode: 1, runnerExitCode: null }, + ])('requests hook exit $expectedExitCode after runner status $runnerExitCode', async (testCase) => { + const requestExit = vi.fn(); + + watchRunnerCompletion({ completion: Promise.resolve(testCase.runnerExitCode) }, quietLogger, requestExit); + + await Promise.resolve(); + expect(requestExit).not.toHaveBeenCalled(); + await new Promise((resolve) => setImmediate(resolve)); + expect(requestExit).toHaveBeenCalledOnce(); + expect(requestExit).toHaveBeenCalledWith(testCase.expectedExitCode); + }); + + it('closes the hook exactly once before publishing its process exit code', async () => { + const events: string[] = []; + const requestExit = createHookExitRequester( + { + close: () => events.push('stop-accepting'), + closeAllConnections: () => events.push('close-connections'), + }, + { + async stop(): Promise { + events.push('stop-runner'); + }, + }, + quietLogger, + (exitCode) => events.push(`exit:${exitCode}`), + ); + + requestExit(0); + requestExit(1); + await new Promise((resolve) => setImmediate(resolve)); + + expect(events).toEqual(['stop-accepting', 'stop-runner', 'close-connections', 'exit:0']); + }); +}); diff --git a/lambdas/services/microvm-lifecycle-hooks/src/server.ts b/lambdas/services/microvm-lifecycle-hooks/src/server.ts new file mode 100644 index 0000000000..fe735653e7 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/server.ts @@ -0,0 +1,275 @@ +import http, { type IncomingMessage, type ServerResponse } from 'node:http'; + +import type { Logger } from './contracts'; +import { consoleLogger } from './contracts'; +import { RunnerLifecycle } from './lifecycle'; +import { HookRequestError, MAX_REQUEST_BYTES } from './payload'; +import { RunnerEntrypointLauncher } from './processes'; +import { StorageJitConfigSource } from './storage'; + +export const HOOK_PREFIX = '/aws/lambda-microvms/runtime/v1'; + +const MAX_TIMER_SECONDS = 2_147_483; + +export interface HookLifecycle { + start(body: string): Promise; + stop(): Promise; + resume(): Promise; +} + +export interface HookServerOptions { + headersTimeoutMs?: number; + keepAliveTimeoutMs?: number; + requestTimeoutMs?: number; +} + +export function parsePositiveInteger( + value: string | undefined, + fallback: number, + maximum = Number.MAX_SAFE_INTEGER, +): number { + if (value === undefined || !/^\d+$/.test(value)) { + return fallback; + } + const parsed = Number(value); + return Number.isSafeInteger(parsed) && parsed > 0 && parsed <= maximum ? parsed : fallback; +} + +function timeoutMilliseconds(variable: string, fallbackSeconds: number, maximumSeconds = 60): number { + return ( + parsePositiveInteger(process.env[variable], fallbackSeconds, Math.min(maximumSeconds, MAX_TIMER_SECONDS)) * 1_000 + ); +} + +function respond(response: ServerResponse, status: number, payload: object): void { + const body = Buffer.from(JSON.stringify(payload)); + response.writeHead(status, { + 'Cache-Control': 'no-store', + 'Content-Length': body.length, + 'Content-Type': 'application/json', + }); + response.end(body); +} + +function readBody(request: IncomingMessage): Promise { + return new Promise((resolve, reject) => { + const contentLength = request.headers['content-length']; + let declaredLength: number | undefined; + if (contentLength !== undefined) { + declaredLength = Number(contentLength); + if (!Number.isInteger(declaredLength) || declaredLength < 0) { + reject(new HookRequestError('Content-Length is invalid')); + request.resume(); + return; + } + if (declaredLength > MAX_REQUEST_BYTES) { + reject(new HookRequestError('request body is too large')); + request.resume(); + return; + } + } + + const chunks: Buffer[] = []; + let size = 0; + let settled = false; + + const fail = (error: Error): void => { + if (settled) { + return; + } + settled = true; + reject(error); + }; + request.on('data', (chunk: Buffer) => { + if (settled) { + return; + } + size += chunk.length; + if (size > MAX_REQUEST_BYTES) { + fail(new HookRequestError('request body is too large')); + request.destroy(); + return; + } + chunks.push(chunk); + }); + request.once('end', () => { + if (settled) { + return; + } + if (declaredLength !== undefined && declaredLength !== size) { + fail(new HookRequestError('Content-Length does not match the request body')); + return; + } + settled = true; + resolve(Buffer.concat(chunks).toString('utf8')); + }); + request.once('aborted', () => fail(new HookRequestError('request body was interrupted'))); + request.once('error', (error) => fail(error)); + }); +} + +export function createHookServer( + lifecycle: HookLifecycle, + logger: Logger = consoleLogger, + options: HookServerOptions = {}, +): http.Server { + const requestTimeout = options.requestTimeoutMs ?? timeoutMilliseconds('HOOK_REQUEST_TIMEOUT_SECONDS', 10); + const headersTimeout = Math.min( + options.headersTimeoutMs ?? timeoutMilliseconds('HOOK_HEADERS_TIMEOUT_SECONDS', 5), + requestTimeout, + ); + const keepAliveTimeout = options.keepAliveTimeoutMs ?? timeoutMilliseconds('HOOK_KEEP_ALIVE_TIMEOUT_SECONDS', 5); + + const server = http.createServer( + { + headersTimeout, + keepAliveTimeout, + maxHeaderSize: 16 * 1024, + requestTimeout, + }, + async (request, response) => { + const path = request.url ?? ''; + if (request.method !== 'POST') { + request.resume(); + respond(response, 405, { error: 'method not allowed' }); + return; + } + + try { + // Consume every POST body so all lifecycle endpoints share the same bounded request handling. + const body = await readBody(request); + if (path === `${HOOK_PREFIX}/ready`) { + respond(response, 200, { status: 'ready' }); + return; + } + if (path === `${HOOK_PREFIX}/validate`) { + respond(response, 200, { status: 'validated' }); + return; + } + if (path === `${HOOK_PREFIX}/run`) { + const started = await lifecycle.start(body); + respond(response, 200, { status: started ? 'started' : 'already-started' }); + return; + } + if (path === `${HOOK_PREFIX}/terminate`) { + await lifecycle.stop(); + respond(response, 200, { status: 'stopped' }); + return; + } + if (path === `${HOOK_PREFIX}/resume`) { + const ready = await lifecycle.resume(); + respond(response, ready ? 200 : 503, { status: ready ? 'ready' : 'not-ready' }); + return; + } + if (path === `${HOOK_PREFIX}/suspend`) { + respond(response, 200, { status: 'ok' }); + return; + } + respond(response, 404, { error: 'unknown lifecycle hook' }); + } catch (error) { + if (error instanceof HookRequestError) { + logger.warn('Rejected invalid lifecycle hook request'); + respond(response, 400, { error: error.message }); + return; + } + // Parse and provider errors can contain credentials in both message and name. + logger.error('Lifecycle hook failed'); + respond(response, 500, { error: 'lifecycle hook failed' }); + } + }, + ); + server.maxConnections = 128; + server.maxHeadersCount = 64; + server.maxRequestsPerSocket = 100; + return server; +} + +export function createDefaultLifecycle(logger: Logger = consoleLogger): RunnerLifecycle { + return new RunnerLifecycle(new StorageJitConfigSource(), new RunnerEntrypointLauncher(), logger); +} + +interface ClosableServer { + close(): unknown; + closeAllConnections(): void; +} + +interface StoppableLifecycle { + stop(): Promise; +} + +export async function shutdownHookServer(server: ClosableServer, lifecycle: StoppableLifecycle): Promise { + server.close(); + try { + await lifecycle.stop(); + } finally { + server.closeAllConnections(); + } +} + +function hookExitCode(runnerExitCode: number | null): number { + return runnerExitCode === 0 ? 0 : 1; +} + +export function watchRunnerCompletion( + lifecycle: Pick, + logger: Logger, + requestExit: (exitCode: number) => void, +): void { + void lifecycle.completion.then((runnerExitCode) => { + const exitCode = hookExitCode(runnerExitCode); + if (exitCode === 0) { + logger.info('GitHub Actions runner exited with status %s', runnerExitCode); + } else { + logger.error('GitHub Actions runner exited unexpectedly with status %s', runnerExitCode ?? 'signal'); + } + // Let the /run handler flush its acknowledgement if the runner exits immediately after readiness. + setImmediate(() => requestExit(exitCode)); + }); +} + +export function createHookExitRequester( + server: ClosableServer, + lifecycle: StoppableLifecycle, + logger: Logger, + setExitCode: (exitCode: number) => void = (exitCode) => { + // Let Node exit naturally after lifecycle cleanup and log streams have drained. + process.exitCode = exitCode; + }, +): (exitCode: number) => void { + let exiting = false; + return (exitCode: number): void => { + if (exiting) { + return; + } + exiting = true; + void shutdownHookServer(server, lifecycle).then( + () => setExitCode(exitCode), + () => { + logger.error('Lifecycle hook shutdown failed'); + setExitCode(1); + }, + ); + }; +} + +export async function main(): Promise { + const logger = consoleLogger; + const lifecycle = createDefaultLifecycle(logger); + const server = createHookServer(lifecycle, logger); + const port = parsePositiveInteger(process.env.HOOK_PORT, 8080, 65_535); + + const requestExit = createHookExitRequester(server, lifecycle, logger); + process.once('SIGINT', () => requestExit(0)); + process.once('SIGTERM', () => requestExit(0)); + watchRunnerCompletion(lifecycle, logger, requestExit); + + await new Promise((resolve, reject) => { + const onError = (): void => reject(new Error('lifecycle hook server could not listen')); + server.once('error', onError); + server.listen(port, '0.0.0.0', () => { + server.off('error', onError); + logger.info('Lambda MicroVM lifecycle hooks listening on port %d', port); + resolve(); + }); + }); +} diff --git a/lambdas/services/microvm-lifecycle-hooks/src/storage.test.ts b/lambdas/services/microvm-lifecycle-hooks/src/storage.test.ts new file mode 100644 index 0000000000..78083db1ab --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/storage.test.ts @@ -0,0 +1,88 @@ +import type { + RunnerConfigConsumer, + RunnerConfigStorageContext, +} from '@aws-github-runner/storage-providers/runner-config-consumer'; + +import { StorageJitConfigSource } from './storage'; + +const DYNAMODB_STORAGE: RunnerConfigStorageContext = { + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'github-runner-config', + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb', +}; + +describe('StorageJitConfigSource', () => { + it('exports the allowlisted context once before resolving and consuming from the environment', async () => { + const events: string[] = []; + const environment: NodeJS.ProcessEnv = {}; + const consumer: RunnerConfigConsumer = { + consume: vi.fn(async () => { + events.push('consume'); + return 'encoded-jit'; + }), + }; + const exportEnvironment = vi.fn((context: RunnerConfigStorageContext, target: NodeJS.ProcessEnv) => { + events.push('export'); + Object.assign(target, context); + }); + const createConsumer = vi.fn((target: NodeJS.ProcessEnv) => { + events.push('create'); + expect(target).toBe(environment); + expect(target).toMatchObject(DYNAMODB_STORAGE); + return consumer; + }); + const source = new StorageJitConfigSource({ createConsumer, environment, exportEnvironment }); + const signal = new AbortController().signal; + + await expect( + source.consume({ microvmId: 'microvm-1234', storage: DYNAMODB_STORAGE }, { deadlineMs: 123_456, signal }), + ).resolves.toEqual({ jitConfig: 'encoded-jit' }); + await expect( + source.consume( + { + microvmId: 'microvm-1234', + storage: { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb', + RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'github-runner-config', + }, + }, + { deadlineMs: 123_457, signal }, + ), + ).resolves.toEqual({ jitConfig: 'encoded-jit' }); + + expect(events).toEqual(['export', 'create', 'consume', 'create', 'consume']); + expect(exportEnvironment).toHaveBeenCalledOnce(); + expect(createConsumer).toHaveBeenCalledTimes(2); + expect(consumer.consume).toHaveBeenNthCalledWith(1, 'microvm-1234', { + deadlineMs: 123_456, + signal, + }); + }); + + it('rejects storage context changes after the one-time environment export', async () => { + const environment: NodeJS.ProcessEnv = {}; + const consumer: RunnerConfigConsumer = { consume: vi.fn().mockResolvedValue('encoded-jit') }; + const exportEnvironment = vi.fn((context: RunnerConfigStorageContext, target: NodeJS.ProcessEnv) => { + Object.assign(target, context); + }); + const createConsumer = vi.fn().mockReturnValue(consumer); + const source = new StorageJitConfigSource({ createConsumer, environment, exportEnvironment }); + const options = { deadlineMs: 123_456, signal: new AbortController().signal }; + + await source.consume({ microvmId: 'microvm-1234', storage: DYNAMODB_STORAGE }, options); + await expect( + source.consume( + { + microvmId: 'microvm-1234', + storage: { + RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', + SSM_TOKEN_PATH: '/github-action-runners/tenant/token', + }, + }, + options, + ), + ).rejects.toThrow('storage context cannot change'); + + expect(exportEnvironment).toHaveBeenCalledOnce(); + expect(createConsumer).toHaveBeenCalledOnce(); + }); +}); diff --git a/lambdas/services/microvm-lifecycle-hooks/src/storage.ts b/lambdas/services/microvm-lifecycle-hooks/src/storage.ts new file mode 100644 index 0000000000..4e50a4c9b0 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/storage.ts @@ -0,0 +1,49 @@ +import { + createRunnerConfigConsumerFromEnvironment, + exportRunnerConfigStorageEnvironment, + type RunnerConfigConsumer, + type RunnerConfigStorageContext, +} from '@aws-github-runner/storage-providers/runner-config-consumer'; + +import type { ConsumeOptions, JitConfigSource, RunContext, RunnerBootstrap } from './contracts'; + +type RunnerConfigConsumerFactory = typeof createRunnerConfigConsumerFromEnvironment; +type RunnerConfigStorageExporter = typeof exportRunnerConfigStorageEnvironment; + +export interface StorageJitConfigSourceOptions { + createConsumer?: RunnerConfigConsumerFactory; + environment?: NodeJS.ProcessEnv; + exportEnvironment?: RunnerConfigStorageExporter; +} + +function storageContextFingerprint(context: RunnerConfigStorageContext): string { + return JSON.stringify(Object.entries(context).sort(([left], [right]) => left.localeCompare(right))); +} + +/** Adapts the shared provider registry to the lifecycle's one-time bootstrap contract. */ +export class StorageJitConfigSource implements JitConfigSource { + private readonly createConsumer: RunnerConfigConsumerFactory; + private readonly environment: NodeJS.ProcessEnv; + private readonly exportEnvironment: RunnerConfigStorageExporter; + private exportedStorageFingerprint?: string; + + public constructor(options: StorageJitConfigSourceOptions = {}) { + this.createConsumer = options.createConsumer ?? createRunnerConfigConsumerFromEnvironment; + this.environment = options.environment ?? process.env; + this.exportEnvironment = options.exportEnvironment ?? exportRunnerConfigStorageEnvironment; + } + + public async consume(context: RunContext, options: ConsumeOptions): Promise { + const fingerprint = storageContextFingerprint(context.storage); + if (this.exportedStorageFingerprint === undefined) { + this.exportEnvironment(context.storage, this.environment); + this.exportedStorageFingerprint = fingerprint; + } else if (this.exportedStorageFingerprint !== fingerprint) { + throw new Error('runner configuration storage context cannot change after initialization'); + } + + const consumer: RunnerConfigConsumer = this.createConsumer(this.environment); + const jitConfig = await consumer.consume(context.microvmId, options); + return { jitConfig }; + } +} diff --git a/lambdas/services/microvm-lifecycle-hooks/src/timing.test.ts b/lambdas/services/microvm-lifecycle-hooks/src/timing.test.ts new file mode 100644 index 0000000000..b62d8af038 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/timing.test.ts @@ -0,0 +1,32 @@ +import { beforeDeadlineOrAbort, delay } from './timing'; + +describe('timing helpers', () => { + it('removes the delay abort listener after resolving', async () => { + const signal = new AbortController().signal; + const remove = vi.spyOn(signal, 'removeEventListener'); + + await delay(1, signal); + + expect(remove).toHaveBeenCalledOnce(); + }); + + it('removes the delay abort listener after cancellation', async () => { + const controller = new AbortController(); + const remove = vi.spyOn(controller.signal, 'removeEventListener'); + const pending = delay(1_000, controller.signal); + + controller.abort(); + + await expect(pending).rejects.toThrow('operation was cancelled'); + expect(remove).toHaveBeenCalledOnce(); + }); + + it('rejects immediately when an operation is already aborted', async () => { + const controller = new AbortController(); + controller.abort(); + + await expect( + beforeDeadlineOrAbort(Promise.resolve('unused'), Date.now() + 1_000, controller.signal), + ).rejects.toThrow('runner start was cancelled'); + }); +}); diff --git a/lambdas/services/microvm-lifecycle-hooks/src/timing.ts b/lambdas/services/microvm-lifecycle-hooks/src/timing.ts new file mode 100644 index 0000000000..fc791d2e29 --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/src/timing.ts @@ -0,0 +1,66 @@ +export function delay(milliseconds: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + let settled = false; + const cleanup = (): void => signal?.removeEventListener('abort', cancel); + const finish = (): void => { + if (settled) { + return; + } + settled = true; + cleanup(); + resolve(); + }; + const timer = setTimeout(finish, milliseconds); + const cancel = (): void => { + if (settled) { + return; + } + settled = true; + clearTimeout(timer); + cleanup(); + reject(new Error('operation was cancelled')); + }; + signal?.addEventListener('abort', cancel, { once: true }); + if (signal?.aborted) { + cancel(); + } + }); +} + +export async function beforeDeadline(promise: Promise, deadlineMs: number): Promise { + const remaining = deadlineMs - Date.now(); + if (remaining <= 0) { + throw new Error('run-hook deadline elapsed'); + } + let timer: NodeJS.Timeout | undefined; + const deadline = new Promise((_resolve, reject) => { + timer = setTimeout(() => reject(new Error('run-hook deadline elapsed')), remaining); + }); + try { + return await Promise.race([promise, deadline]); + } finally { + if (timer !== undefined) { + clearTimeout(timer); + } + } +} + +export async function beforeDeadlineOrAbort( + promise: Promise, + deadlineMs: number, + signal: AbortSignal, +): Promise { + if (signal.aborted) { + throw new Error('runner start was cancelled'); + } + let cancel = (): void => undefined; + const cancelled = new Promise((_resolve, reject) => { + cancel = (): void => reject(new Error('runner start was cancelled')); + signal.addEventListener('abort', cancel, { once: true }); + }); + try { + return await beforeDeadline(Promise.race([promise, cancelled]), deadlineMs); + } finally { + signal.removeEventListener('abort', cancel); + } +} diff --git a/lambdas/services/microvm-lifecycle-hooks/tsconfig.json b/lambdas/services/microvm-lifecycle-hooks/tsconfig.json new file mode 100644 index 0000000000..714aa27b6b --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.json", + "compilerOptions": { + "types": ["node", "vitest/globals"] + }, + "include": ["src/**/*"], + "exclude": ["src/**/*.test.ts"] +} diff --git a/lambdas/services/microvm-lifecycle-hooks/vitest.config.ts b/lambdas/services/microvm-lifecycle-hooks/vitest.config.ts new file mode 100644 index 0000000000..e3c59146ee --- /dev/null +++ b/lambdas/services/microvm-lifecycle-hooks/vitest.config.ts @@ -0,0 +1,12 @@ +import { mergeConfig } from 'vitest/config'; + +import defaultConfig from '../../vitest.base.config'; + +export default mergeConfig(defaultConfig, { + test: { + coverage: { + include: ['src/**/*.ts'], + exclude: ['src/**/*.test.ts'], + }, + }, +}); diff --git a/lambdas/yarn.lock b/lambdas/yarn.lock index c80fecbd3a..75a5ceaa94 100644 --- a/lambdas/yarn.lock +++ b/lambdas/yarn.lock @@ -209,6 +209,16 @@ __metadata: languageName: unknown linkType: soft +"@aws-github-runner/microvm-lifecycle-hooks@workspace:services/microvm-lifecycle-hooks": + version: 0.0.0-use.local + resolution: "@aws-github-runner/microvm-lifecycle-hooks@workspace:services/microvm-lifecycle-hooks" + dependencies: + "@aws-github-runner/storage-providers": "npm:*" + "@types/node": "npm:^22.19.3" + "@vercel/ncc": "npm:^0.38.4" + 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"