From 28ca06e56c350230cae6444348bbbcb25a962f80 Mon Sep 17 00:00:00 2001 From: Tofik Hasanov Date: Mon, 8 Jun 2026 14:53:37 -0400 Subject: [PATCH] fix(cloud-security): retry transient AssumeRole failures (first-scan "could not assume role") MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Customer-reported: the first AWS scan fails with a "Could not assume AWS role" finding; switching the scan engine and back, then re-running, works — on every account, with no change in AWS. Customer correctly concluded it's not the role (it assumes fine the second time). Root cause (verified): assuming the customer's cross-account role is a two-hop STS flow with NO application-level retry, and the AWS SDK's default retry only covers throttling/5xx/network — NOT the classes that actually occur here: - IAM/STS eventual consistency on a new or just-edited role/trust-policy returns AccessDenied on the first assume and succeeds seconds later (no AWS change). - The base/roleAssumer session can briefly return ExpiredToken. So a single transient failure surfaced as a sticky finding. The "switch switch" is INCIDENTAL — updateMode only writes metadata (idempotent, no cache/cred refresh); what actually fixes it is the elapsed time before the retry. The visible finding is written by the integration-platform checks path (resolveAwsSessionOrFail in manifests/aws/checks/shared.ts), not the Cloud Tests scan (which returns success:false with no finding) — so both assume implementations are fixed. - new shared helper retryAssume / isRetryableAssumeError (manifests/aws/checks/assume-retry.ts): bounded exponential backoff + jitter, retries only transient / eventual-consistency classes (AccessDenied, ExpiredToken, IDPCommunicationError, throttling, 5xx, network). Hard config errors (bad ARN, ValidationError) and a persistently-denied role still propagate; backoff capped so a genuinely broken role fails within seconds. - wrap both AssumeRole hops in shared.ts (checks path) and aws-security.service.ts assumeRole (Cloud Tests scan path). - export retryAssume from the package barrel for apps/api. - tests: retry-then-succeed, no-retry on config errors, give-up after N, classifier. Note: not yet live-verified; the exact STS class (AccessDenied vs ExpiredToken) is already captured in evidence.error on the failing finding — the fix covers both. cloud-security suite green (311); helper tests green; typecheck clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../providers/aws-security.service.ts | 31 ++--- packages/integration-platform/src/index.ts | 7 ++ .../aws/checks/__tests__/assume-retry.test.ts | 89 +++++++++++++++ .../src/manifests/aws/checks/assume-retry.ts | 106 ++++++++++++++++++ .../src/manifests/aws/checks/shared.ts | 31 ++--- 5 files changed, 238 insertions(+), 26 deletions(-) create mode 100644 packages/integration-platform/src/manifests/aws/checks/__tests__/assume-retry.test.ts create mode 100644 packages/integration-platform/src/manifests/aws/checks/assume-retry.ts diff --git a/apps/api/src/cloud-security/providers/aws-security.service.ts b/apps/api/src/cloud-security/providers/aws-security.service.ts index de43d89555..e09ef7c5ab 100644 --- a/apps/api/src/cloud-security/providers/aws-security.service.ts +++ b/apps/api/src/cloud-security/providers/aws-security.service.ts @@ -1,5 +1,6 @@ import { Injectable, Logger } from '@nestjs/common'; import { AssumeRoleCommand, STSClient } from '@aws-sdk/client-sts'; +import { retryAssume } from '@trycompai/integration-platform'; import { CostExplorerClient, GetCostAndUsageCommand, @@ -514,12 +515,14 @@ export class AWSSecurityService { region, credentials: getAwsBaseCredentials(partition), }); - const roleAssumerResp = await baseSts.send( - new AssumeRoleCommand({ - RoleArn: roleAssumerArn, - RoleSessionName: 'CompRoleAssumer', - DurationSeconds: 3600, - }), + const roleAssumerResp = await retryAssume(() => + baseSts.send( + new AssumeRoleCommand({ + RoleArn: roleAssumerArn, + RoleSessionName: 'CompRoleAssumer', + DurationSeconds: 3600, + }), + ), ); const roleAssumerCreds = roleAssumerResp.Credentials; @@ -541,13 +544,15 @@ export class AWSSecurityService { this.logger.log(`Assuming customer role ${roleArn} in region ${region}`); - const customerResp = await roleAssumerSts.send( - new AssumeRoleCommand({ - RoleArn: roleArn, - ExternalId: externalId, - RoleSessionName: sessionName ?? 'CompSecurityAudit', - DurationSeconds: 3600, - }), + const customerResp = await retryAssume(() => + roleAssumerSts.send( + new AssumeRoleCommand({ + RoleArn: roleArn, + ExternalId: externalId, + RoleSessionName: sessionName ?? 'CompSecurityAudit', + DurationSeconds: 3600, + }), + ), ); const customerCreds = customerResp.Credentials; diff --git a/packages/integration-platform/src/index.ts b/packages/integration-platform/src/index.ts index 6131578d9c..b50fed9c72 100644 --- a/packages/integration-platform/src/index.ts +++ b/packages/integration-platform/src/index.ts @@ -145,6 +145,13 @@ export { } from './manifests/aws/credentials'; export type { AwsEnvironment } from './manifests/aws/credentials'; +// Shared AWS STS AssumeRole retry (transient / IAM-eventual-consistency safe), +// reused by the Cloud Tests scanner in apps/api. +export { + retryAssume, + isRetryableAssumeError, +} from './manifests/aws/checks/assume-retry'; + // API Response types (for frontend and API type sharing) export type { diff --git a/packages/integration-platform/src/manifests/aws/checks/__tests__/assume-retry.test.ts b/packages/integration-platform/src/manifests/aws/checks/__tests__/assume-retry.test.ts new file mode 100644 index 0000000000..bb2eed1e8c --- /dev/null +++ b/packages/integration-platform/src/manifests/aws/checks/__tests__/assume-retry.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, it } from 'bun:test'; +import { isRetryableAssumeError, retryAssume } from '../assume-retry'; + +// Fast, deterministic options for tests — never actually sleep. +const fastOpts = { + sleep: async () => {}, + random: () => 0, + baseDelayMs: 1, + maxDelayMs: 1, +}; + +function awsError(name: string, extra: Record = {}): Error { + return Object.assign(new Error(`${name} occurred`), { name, ...extra }); +} + +describe('isRetryableAssumeError', () => { + it('retries IAM/STS eventual-consistency AccessDenied (the customer bug)', () => { + expect(isRetryableAssumeError(awsError('AccessDenied'))).toBe(true); + expect(isRetryableAssumeError(awsError('AccessDeniedException'))).toBe(true); + }); + + it('retries expired-token, throttling, 5xx and network errors', () => { + expect(isRetryableAssumeError(awsError('ExpiredToken'))).toBe(true); + expect(isRetryableAssumeError(awsError('ThrottlingException'))).toBe(true); + expect( + isRetryableAssumeError({ name: 'X', $metadata: { httpStatusCode: 503 } }), + ).toBe(true); + expect(isRetryableAssumeError(new Error('socket hang up'))).toBe(true); + }); + + it('does NOT retry hard configuration errors', () => { + expect(isRetryableAssumeError(awsError('ValidationError'))).toBe(false); + expect( + isRetryableAssumeError(new Error('Invalid IAM Role ARN format')), + ).toBe(false); + expect(isRetryableAssumeError(null)).toBe(false); + expect(isRetryableAssumeError(undefined)).toBe(false); + }); +}); + +describe('retryAssume', () => { + it('retries a transient AccessDenied then succeeds (no AWS change needed)', async () => { + let calls = 0; + const result = await retryAssume(async () => { + calls++; + if (calls < 3) throw awsError('AccessDenied'); + return 'creds'; + }, fastOpts); + expect(result).toBe('creds'); + expect(calls).toBe(3); + }); + + it('does not retry a non-transient error (single attempt)', async () => { + let calls = 0; + await expect( + retryAssume(async () => { + calls++; + throw awsError('ValidationError'); + }, fastOpts), + ).rejects.toThrow(); + expect(calls).toBe(1); + }); + + it('gives up after the configured attempts on a persistent transient error', async () => { + let calls = 0; + await expect( + retryAssume( + async () => { + calls++; + throw awsError('AccessDenied'); + }, + { ...fastOpts, attempts: 4 }, + ), + ).rejects.toThrow(); + expect(calls).toBe(4); + }); + + it('succeeds on the first attempt without sleeping', async () => { + let slept = false; + const result = await retryAssume(async () => 'ok', { + ...fastOpts, + sleep: async () => { + slept = true; + }, + }); + expect(result).toBe('ok'); + expect(slept).toBe(false); + }); +}); diff --git a/packages/integration-platform/src/manifests/aws/checks/assume-retry.ts b/packages/integration-platform/src/manifests/aws/checks/assume-retry.ts new file mode 100644 index 0000000000..d81fe7795b --- /dev/null +++ b/packages/integration-platform/src/manifests/aws/checks/assume-retry.ts @@ -0,0 +1,106 @@ +/** + * Bounded retry for AWS STS `AssumeRole` calls. + * + * Why this exists: assuming a customer's cross-account role is a two-hop STS + * flow, and the FIRST attempt right after a connection is set up (or its trust + * policy edited) can fail transiently with NO change in AWS — IAM/STS is + * eventually consistent, so a brand-new/edited role can return `AccessDenied` + * on the first assume and then succeed seconds later. The base/roleAssumer + * session can likewise briefly return `ExpiredToken`. None of these classes are + * retried by the AWS SDK's default strategy (it only retries throttling/5xx/ + * network), so a single transient failure surfaced to customers as a sticky + * "Could not assume AWS role" finding that "fixed itself" on the next run. + * + * This helper retries only those transient/eventually-consistent classes with + * capped exponential backoff + jitter. Hard configuration errors (bad ARN, + * ValidationError) and a role that is *persistently* denied still propagate — + * the backoff is capped so a genuinely broken role fails within a few seconds. + */ + +const RETRYABLE_ASSUME_ERROR_NAMES = new Set([ + // IAM/STS eventual consistency on a new or just-edited cross-account role. + 'AccessDenied', + 'AccessDeniedException', + // Base/roleAssumer session momentarily expired, or an STS hiccup. + 'ExpiredToken', + 'ExpiredTokenException', + 'IDPCommunicationError', + 'RegionDisabledException', + // Standard transient classes (belt-and-suspenders alongside the SDK). + 'Throttling', + 'ThrottlingException', + 'TooManyRequestsException', + 'RequestLimitExceeded', + 'ServiceUnavailable', + 'ServiceUnavailableException', + 'InternalError', + 'InternalFailure', +]); + +export function isRetryableAssumeError(error: unknown): boolean { + if (!error || typeof error !== 'object') return false; + const e = error as { + name?: string; + Code?: string; + code?: string; + __type?: string; + message?: string; + $metadata?: { httpStatusCode?: number }; + }; + const name = e.name ?? e.Code ?? e.code ?? e.__type ?? ''; + if (RETRYABLE_ASSUME_ERROR_NAMES.has(name)) return true; + + const status = e.$metadata?.httpStatusCode; + if (typeof status === 'number' && status >= 500) return true; + + return /ECONNRESET|ETIMEDOUT|ENOTFOUND|EAI_AGAIN|socket hang up|network|timeout/i.test( + e.message ?? '', + ); +} + +export interface RetryAssumeOptions { + /** Total attempts including the first (default 4). */ + attempts?: number; + /** Base backoff in ms; grows exponentially, capped by maxDelayMs (default 500). */ + baseDelayMs?: number; + /** Maximum backoff per attempt in ms (default 5000). */ + maxDelayMs?: number; + /** Injectable sleep so tests don't actually wait. */ + sleep?: (ms: number) => Promise; + /** Injectable jitter source in [0, 1); defaults to Math.random. */ + random?: () => number; +} + +const defaultSleep = (ms: number): Promise => + new Promise((resolve) => setTimeout(resolve, ms)); + +/** + * Run `fn` (an AssumeRole call) with bounded exponential backoff + jitter, + * retrying only transient/eventually-consistent failures. Non-transient errors + * propagate immediately; transient ones that never recover propagate after the + * attempts are exhausted. + */ +export async function retryAssume( + fn: () => Promise, + options: RetryAssumeOptions = {}, +): Promise { + const attempts = Math.max(1, options.attempts ?? 4); + const baseDelayMs = options.baseDelayMs ?? 500; + const maxDelayMs = options.maxDelayMs ?? 5000; + const sleep = options.sleep ?? defaultSleep; + const random = options.random ?? Math.random; + + let lastError: unknown; + for (let attempt = 0; attempt < attempts; attempt++) { + try { + return await fn(); + } catch (error) { + lastError = error; + const isLastAttempt = attempt === attempts - 1; + if (isLastAttempt || !isRetryableAssumeError(error)) throw error; + const backoff = Math.min(maxDelayMs, baseDelayMs * 2 ** attempt); + await sleep(backoff + Math.floor(backoff * 0.5 * random())); + } + } + throw lastError; +} diff --git a/packages/integration-platform/src/manifests/aws/checks/shared.ts b/packages/integration-platform/src/manifests/aws/checks/shared.ts index 8a7e48e6eb..c95735ef94 100644 --- a/packages/integration-platform/src/manifests/aws/checks/shared.ts +++ b/packages/integration-platform/src/manifests/aws/checks/shared.ts @@ -1,5 +1,6 @@ import { AssumeRoleCommand, STSClient } from '@aws-sdk/client-sts'; import type { CheckContext, FindingSeverity } from '../../../types'; +import { retryAssume } from './assume-retry'; export interface AwsSession { credentials: { @@ -113,12 +114,14 @@ export async function assumeAwsSession( region, credentials: awsBaseCredentials(partition), }); - const assumerResp = await baseSts.send( - new AssumeRoleCommand({ - RoleArn: roleAssumerArn, - RoleSessionName: 'CompRoleAssumer', - DurationSeconds: 3600, - }), + const assumerResp = await retryAssume(() => + baseSts.send( + new AssumeRoleCommand({ + RoleArn: roleAssumerArn, + RoleSessionName: 'CompRoleAssumer', + DurationSeconds: 3600, + }), + ), ); const assumer = assumerResp.Credentials; if ( @@ -139,13 +142,15 @@ export async function assumeAwsSession( sessionToken: assumer.SessionToken, }, }); - const res = await assumerSts.send( - new AssumeRoleCommand({ - RoleArn: roleArn, - ExternalId: externalId, - RoleSessionName: 'CompEvidenceCheck', - DurationSeconds: 3600, - }), + const res = await retryAssume(() => + assumerSts.send( + new AssumeRoleCommand({ + RoleArn: roleArn, + ExternalId: externalId, + RoleSessionName: 'CompEvidenceCheck', + DurationSeconds: 3600, + }), + ), ); const c = res.Credentials; if (!c?.AccessKeyId || !c.SecretAccessKey || !c.SessionToken) return null;