fix(cloud-security): retry transient AssumeRole failures (first-scan "could not assume role")

Merged
TO
5mo ago
+238-26

Summary

Customer (with multiple AWS accounts, cross-account role auth) 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 nothing changed in AWS. They correctly concluded it isn't the role (it assumes fine the second time) — it's our bug.

Root cause (investigated + adversarially verified)

Assuming the customer's cross-account role is a two-hop STS flow (task role → Comp roleAssumer → customer role) with no application-level retry, and the AWS SDK's default retry strategy only covers throttling / 5xx / network — not the classes that actually occur here:

  • IAM/STS eventual consistency: a new or just-edited role/trust-policy returns AccessDenied on the first assume and succeeds seconds later, with no AWS change.
  • The base/roleAssumer session can briefly return ExpiredToken.

A single transient failure surfaced as a sticky finding, written by the integration-platform checks path (resolveAwsSessionOrFail → ctx.fail({ title: 'Could not assume AWS role' }) in manifests/aws/checks/shared.ts) — not the Cloud Tests scan (which returns success:false with no finding). So both assume implementations are fixed.

The "switch switch" is incidental, not causal. updateMode only writes metadata.awsScanMode (idempotent, no cache or credential refresh) and both modes share the same assume path. What actually makes the retry succeed is the elapsed time before re-running.

Fix

  • New shared helper retryAssume / isRetryableAssumeError (manifests/aws/checks/assume-retry.ts): bounded exponential backoff + jitter, retrying 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 is capped so a genuinely broken role fails within a few seconds.
  • Wrapped both AssumeRole hops in shared.ts (checks path — writes the finding) and aws-security.service.ts assumeRole (Cloud Tests scan path).
  • Exported retryAssume from the package barrel for apps/api.

Tests

  • assume-retry.test.ts: retry-then-succeed, no-retry on config errors, give-up after N attempts, and the classifier (AccessDenied/ExpiredToken/throttling/5xx/network = retryable; ValidationError/bad-ARN = not). 7 passing.
  • Full cloud-security jest suite green (311); typecheck clean.

Honest notes

  • Not yet live-verified. The exact STS class for this org (AccessDenied vs ExpiredToken) is already captured in evidence.error on the failing finding — the fix covers both. Pulling that error string would confirm which hop/class it was.
  • Safety: pure additive retry — happy path (assume succeeds first try) is byte-identical. Bounded attempts so a broken role can't hang the scan.

🤖 Generated with Claude Code

Summary by cubic

Retries transient AWS STS AssumeRole failures with bounded backoff to stop first‑scan “Could not assume AWS role” findings. Applies to both the checks path and the Cloud Tests scanner.

  • Bug Fixes
    • Added shared retryAssume and isRetryableAssumeError with capped exponential backoff + jitter; retries AccessDenied, ExpiredToken, throttling, 5xx, and network errors; skips config errors (e.g., ValidationError).
    • Wrapped both AssumeRole hops in the checks (shared.ts) and scanner (aws-security.service.ts) paths to ensure transient failures don’t surface as findings.
    • Exported from @trycompai/integration-platform for reuse in apps/api; added unit tests for retry behavior; no change to the success path.