|
| 1 | +import { createLogger } from '@sim/logger' |
| 2 | +import { type AWSCredentials, awsCredentialProvider } from '@trigger.dev/sdk/oidc/aws' |
1 | 3 | import { env } from '@/lib/core/config/env' |
2 | 4 |
|
3 | | -interface AwsCredentials { |
4 | | - accessKeyId: string |
5 | | - secretAccessKey: string |
| 5 | +const logger = createLogger('AwsCredentials') |
| 6 | + |
| 7 | +type AwsCredentialsProvider = () => Promise<AWSCredentials> |
| 8 | + |
| 9 | +/** |
| 10 | + * Refresh workload credentials this long before they expire. It must exceed the AWS SDK's own |
| 11 | + * five-minute refresh window, or a client keeps asking for credentials this module still |
| 12 | + * considers fresh. |
| 13 | + */ |
| 14 | +const REFRESH_BEFORE_EXPIRY_MS = 10 * 60_000 |
| 15 | + |
| 16 | +/** |
| 17 | + * At the start of a run, refresh workload credentials that expire within this window. A token can |
| 18 | + * only be minted while an attempt is executing, and work a run leaves behind (the hosted-key |
| 19 | + * metrics flush) still calls AWS after the attempt ends. |
| 20 | + */ |
| 21 | +const RUN_START_MIN_REMAINING_MS = 30 * 60_000 |
| 22 | + |
| 23 | +/** |
| 24 | + * One AWS session for the whole worker process, not one per run. Every run in a deployment |
| 25 | + * assumes the same role, and Sim's AWS clients are process-wide, so a run-named session would |
| 26 | + * attribute later runs' calls to the run that minted it. |
| 27 | + */ |
| 28 | +const WORKLOAD_ROLE_SESSION_NAME = 'sim-trigger-worker' |
| 29 | + |
| 30 | +let workloadProvider: AwsCredentialsProvider | null = null |
| 31 | +let currentWorkloadCredentials: AWSCredentials | null = null |
| 32 | +let pendingWorkloadCredentials: Promise<AWSCredentials> | null = null |
| 33 | + |
| 34 | +function getWorkloadProvider(roleArn: string): AwsCredentialsProvider { |
| 35 | + if (!workloadProvider) { |
| 36 | + if (!env.AWS_REGION) { |
| 37 | + throw new Error('TRIGGER_AWS_ROLE_ARN is set but AWS_REGION is not; set both.') |
| 38 | + } |
| 39 | + workloadProvider = awsCredentialProvider({ |
| 40 | + roleArn, |
| 41 | + region: env.AWS_REGION, |
| 42 | + roleSessionName: WORKLOAD_ROLE_SESSION_NAME, |
| 43 | + }) |
| 44 | + } |
| 45 | + return workloadProvider |
| 46 | +} |
| 47 | + |
| 48 | +/** Credentials this close to expiry are treated as expired, so a signed request never outlives them. */ |
| 49 | +const EXPIRY_SKEW_MS = 60_000 |
| 50 | + |
| 51 | +function remainingMs(credentials: AWSCredentials | null): number { |
| 52 | + return credentials ? credentials.expiration.getTime() - Date.now() : 0 |
| 53 | +} |
| 54 | + |
| 55 | +/** |
| 56 | + * The cached workload credentials while they have more than `minRemainingMs` left, otherwise a |
| 57 | + * fresh exchange. Concurrent callers share one exchange. When the exchange fails, still-valid |
| 58 | + * cached credentials are served instead: outside an executing attempt no token can be minted, and |
| 59 | + * work a run leaves behind (the metrics flush) must not lose credentials that have not expired. |
| 60 | + */ |
| 61 | +function resolveWorkloadCredentials( |
| 62 | + roleArn: string, |
| 63 | + minRemainingMs: number |
| 64 | +): Promise<AWSCredentials> { |
| 65 | + const current = currentWorkloadCredentials |
| 66 | + if (current && remainingMs(current) > minRemainingMs) { |
| 67 | + return Promise.resolve(current) |
| 68 | + } |
| 69 | + if (!pendingWorkloadCredentials) { |
| 70 | + pendingWorkloadCredentials = getWorkloadProvider(roleArn)() |
| 71 | + .then((credentials) => { |
| 72 | + currentWorkloadCredentials = credentials |
| 73 | + return credentials |
| 74 | + }) |
| 75 | + .catch((error: unknown) => { |
| 76 | + const cached = currentWorkloadCredentials |
| 77 | + if (!cached || remainingMs(cached) <= EXPIRY_SKEW_MS) throw error |
| 78 | + logger.warn('Workload-identity exchange failed; serving cached credentials until expiry', { |
| 79 | + error, |
| 80 | + remainingMs: remainingMs(cached), |
| 81 | + }) |
| 82 | + return cached |
| 83 | + }) |
| 84 | + .finally(() => { |
| 85 | + pendingWorkloadCredentials = null |
| 86 | + }) |
| 87 | + } |
| 88 | + return pendingWorkloadCredentials |
6 | 89 | } |
7 | 90 |
|
8 | 91 | /** |
9 | | - * Explicit AWS credentials from the environment, or `undefined` to defer to the |
10 | | - * default AWS provider chain (the ECS task role in our deployments). |
| 92 | + * Credentials for Sim's own AWS clients: file storage (S3), AppConfig, email (SES), hosted-key |
| 93 | + * metrics and hosted Bedrock. Never for a user's integration, which brings its own keys. |
11 | 94 | * |
12 | | - * Shared by every AWS SDK client (S3, AppConfig, …) so credential resolution is |
13 | | - * identical everywhere: explicit keys when both `AWS_ACCESS_KEY_ID` and |
14 | | - * `AWS_SECRET_ACCESS_KEY` are set (self-hosted, trigger.dev workers), otherwise |
15 | | - * the instance/task role. |
| 95 | + * With `TRIGGER_AWS_ROLE_ARN` set (deployed Trigger.dev runs), temporary credentials for that |
| 96 | + * role, exchanged for the run's Trigger.dev workload-identity token. It takes precedence over any |
| 97 | + * static keys in the environment, so a failed exchange surfaces instead of silently falling back |
| 98 | + * to a stored key. Otherwise `undefined`: the AWS SDK's default chain, which reads |
| 99 | + * `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY` and `AWS_SESSION_TOKEN`, then shared profiles, then |
| 100 | + * the ECS task role. |
| 101 | + */ |
| 102 | +export function getAwsCredentials(): AwsCredentialsProvider | undefined { |
| 103 | + const roleArn = env.TRIGGER_AWS_ROLE_ARN |
| 104 | + if (!roleArn) return undefined |
| 105 | + return () => resolveWorkloadCredentials(roleArn, REFRESH_BEFORE_EXPIRY_MS) |
| 106 | +} |
| 107 | + |
| 108 | +/** Whether Sim's own AWS clients have credentials of their own rather than the default chain. */ |
| 109 | +export function hasExplicitAwsCredentials(): boolean { |
| 110 | + return Boolean(env.TRIGGER_AWS_ROLE_ARN || (env.AWS_ACCESS_KEY_ID && env.AWS_SECRET_ACCESS_KEY)) |
| 111 | +} |
| 112 | + |
| 113 | +/** |
| 114 | + * Refreshes workload credentials at the start of a Trigger.dev run, while a token can be minted, |
| 115 | + * so they outlive the attempt. A no-op outside workload identity. Never throws: a failure is |
| 116 | + * logged here, and the run's first AWS call retries the exchange and fails with the cause. |
16 | 117 | */ |
17 | | -export function getAwsCredentialsFromEnv(): AwsCredentials | undefined { |
18 | | - return env.AWS_ACCESS_KEY_ID && env.AWS_SECRET_ACCESS_KEY |
19 | | - ? { |
20 | | - accessKeyId: env.AWS_ACCESS_KEY_ID, |
21 | | - secretAccessKey: env.AWS_SECRET_ACCESS_KEY, |
22 | | - } |
23 | | - : undefined |
| 118 | +export async function refreshAwsCredentialsForRun(): Promise<void> { |
| 119 | + const roleArn = env.TRIGGER_AWS_ROLE_ARN |
| 120 | + if (!roleArn) return |
| 121 | + try { |
| 122 | + await resolveWorkloadCredentials(roleArn, RUN_START_MIN_REMAINING_MS) |
| 123 | + } catch (error) { |
| 124 | + logger.error('Could not exchange the workload-identity token for AWS credentials', { error }) |
| 125 | + } |
24 | 126 | } |
0 commit comments