From 0f768383926da7a02c3ef6aa8825ea9817e88b9c Mon Sep 17 00:00:00 2001 From: Md Mozammil Khan Date: Wed, 29 Jul 2026 16:20:15 +0530 Subject: [PATCH] feat(aws-kinesis): cache assumed-role credentials to avoid STS throttling Under high TPS, aws-kinesis called STS AssumeRole on every request (twice: intermediary role + target role), causing IAM AssumeRole throttling errors. Add an in-memory TTL cache in the shared AWS sts lib (used only by aws-kinesis), keyed by role ARN + external id + region: - TTL derived from the actual STS credential Expiration minus a 5-minute safety buffer, falling back to 55 minutes when STS omits an expiration. - In-flight refresh de-duplication so a burst of concurrent cache misses collapses into a single STS refresh (avoids thundering herd). This reduces STS calls from two per request to at most two per role per TTL window. assumeRole's public signature is unchanged. Co-Authored-By: Claude Opus 4.8 --- .../src/lib/AWS/__test__/index.test.ts | 96 +++++++++++++++++- .../destination-actions/src/lib/AWS/sts.ts | 97 +++++++++++++++++-- 2 files changed, 186 insertions(+), 7 deletions(-) diff --git a/packages/destination-actions/src/lib/AWS/__test__/index.test.ts b/packages/destination-actions/src/lib/AWS/__test__/index.test.ts index 0bfa92a758..5a93fdd114 100644 --- a/packages/destination-actions/src/lib/AWS/__test__/index.test.ts +++ b/packages/destination-actions/src/lib/AWS/__test__/index.test.ts @@ -1,5 +1,5 @@ import { STSClient, AssumeRoleCommand } from '@aws-sdk/client-sts' -import { assumeRole } from '../sts' +import { assumeRole, __clearAssumedRoleCacheForTests } from '../sts' import { ErrorCodes } from '@segment/actions-core' // Mock dependencies @@ -24,6 +24,8 @@ describe('assumeRole', () => { beforeEach(() => { jest.clearAllMocks() + // Reset the in-memory credentials cache so cases don't leak assumed roles into one another. + __clearAssumedRoleCacheForTests() ;(STSClient as jest.Mock).mockImplementation(() => ({ send: mockSend })) @@ -108,4 +110,96 @@ describe('assumeRole', () => { }) ) }) + + const mockBothRoleCalls = () => { + mockSend + .mockResolvedValueOnce({ + Credentials: { + AccessKeyId: 'AKIA_INTERMEDIARY', + SecretAccessKey: 'SECRET_INTERMEDIARY', + SessionToken: 'TOKEN_INTERMEDIARY' + } + }) + .mockResolvedValueOnce({ + Credentials: { + AccessKeyId: 'AKIA_FINAL', + SecretAccessKey: 'SECRET_FINAL', + SessionToken: 'TOKEN_FINAL' + } + }) + } + + describe('credentials caching', () => { + it('reuses cached credentials for the same role/externalId/region without calling STS again', async () => { + mockBothRoleCalls() + + const first = await assumeRole('arn:aws:iam::222222222222:role/TargetRole', 'external-id', 'us-east-1') + expect(STSClient).toHaveBeenCalledTimes(2) + + const second = await assumeRole('arn:aws:iam::222222222222:role/TargetRole', 'external-id', 'us-east-1') + // No additional STS calls: served entirely from the in-memory cache. + expect(STSClient).toHaveBeenCalledTimes(2) + expect(second).toEqual(first) + }) + + it('does not share cache entries across different roles/regions', async () => { + mockBothRoleCalls() + await assumeRole('arn:aws:iam::222222222222:role/TargetRole', 'external-id', 'us-east-1') + expect(STSClient).toHaveBeenCalledTimes(2) + + mockBothRoleCalls() + await assumeRole('arn:aws:iam::999999999999:role/OtherRole', 'external-id', 'us-east-1') + // Different role -> cache miss -> another pair of STS calls. + expect(STSClient).toHaveBeenCalledTimes(4) + }) + + it('collapses concurrent cache misses into a single STS refresh', async () => { + mockBothRoleCalls() + + const [a, b] = await Promise.all([ + assumeRole('arn:aws:iam::222222222222:role/TargetRole', 'external-id', 'us-east-1'), + assumeRole('arn:aws:iam::222222222222:role/TargetRole', 'external-id', 'us-east-1') + ]) + + // Both concurrent callers share one refresh (intermediary + target = 2 STS calls total). + expect(STSClient).toHaveBeenCalledTimes(2) + expect(a).toEqual(b) + }) + + it('refreshes after cached credentials expire', async () => { + const nowSpy = jest.spyOn(Date, 'now') + try { + // STS returns an expiration 1 hour out; cache TTL = expiration - 5min buffer. + const t0 = 1_000_000_000_000 + nowSpy.mockReturnValue(t0) + mockSend + .mockResolvedValueOnce({ + Credentials: { + AccessKeyId: 'AKIA_INTERMEDIARY', + SecretAccessKey: 'SECRET_INTERMEDIARY', + SessionToken: 'TOKEN_INTERMEDIARY' + } + }) + .mockResolvedValueOnce({ + Credentials: { + AccessKeyId: 'AKIA_FINAL', + SecretAccessKey: 'SECRET_FINAL', + SessionToken: 'TOKEN_FINAL', + Expiration: new Date(t0 + 60 * 60 * 1000) + } + }) + + await assumeRole('arn:aws:iam::222222222222:role/TargetRole', 'external-id', 'us-east-1') + expect(STSClient).toHaveBeenCalledTimes(2) + + // Advance past the cached expiry (1h - 5min buffer). + nowSpy.mockReturnValue(t0 + 60 * 60 * 1000) + mockBothRoleCalls() + await assumeRole('arn:aws:iam::222222222222:role/TargetRole', 'external-id', 'us-east-1') + expect(STSClient).toHaveBeenCalledTimes(4) + } finally { + nowSpy.mockRestore() + } + }) + }) }) diff --git a/packages/destination-actions/src/lib/AWS/sts.ts b/packages/destination-actions/src/lib/AWS/sts.ts index 814660ad98..c01950e344 100644 --- a/packages/destination-actions/src/lib/AWS/sts.ts +++ b/packages/destination-actions/src/lib/AWS/sts.ts @@ -51,6 +51,37 @@ const awsCredentialsCache: AWSCredentialsCache = { credentials: { accessKeyId: '', secretAccessKey: '', sessionToken: '' } } +type AssumedRoleCacheEntry = { + credentials: AWSCredentials + expiresAt: number +} + +// In-memory cache of assumed-role credentials, keyed by role ARN + external id + region. +// Under high TPS, calling STS AssumeRole on every request causes IAM throttling. Caching the +// assumed credentials until shortly before they expire keeps STS calls to (at most) one refresh +// per role per TTL window instead of one (well, two - intermediary + target) per request. +const assumedRoleCache = new Map() + +// De-duplicates concurrent refreshes for the same key. Without this, a burst of requests that +// all miss the cache at the same time would each trigger a fresh set of STS calls (thundering +// herd). Instead they all await the single in-flight refresh. +const inflightRoleRefreshes = new Map>() + +// Refresh a little before the credentials actually expire so in-flight requests aren't handed +// credentials that expire mid-use. +const CREDENTIALS_EXPIRY_BUFFER_MS = 5 * 60 * 1000 // 5 minutes +// Fallback TTL used when STS does not return an expiration. AssumeRole sessions default to 1 hour. +const DEFAULT_CREDENTIALS_TTL_MS = 55 * 60 * 1000 // 55 minutes + +const buildAssumedRoleCacheKey = (roleArn: string, externalId: string, region: string): string => + `${roleArn}|${externalId}|${region}` + +// Exposed for tests to reset the in-memory caches between cases. +export const __clearAssumedRoleCacheForTests = (): void => { + assumedRoleCache.clear() + inflightRoleRefreshes.clear() +} + function getToken(): string { const tokenFilepath = process.env['AWS_WEB_IDENTITY_TOKEN_FILE'] || '/var/run/secrets/kubernetes.io/serviceaccount/token' @@ -133,13 +164,64 @@ export async function getAWSCredentialsFromEKS(request: RequestClient): Promise< } export const assumeRole = async (roleArn: string, externalId: string, region: string): Promise => { + const cacheKey = buildAssumedRoleCacheKey(roleArn, externalId, region) + + const cached = assumedRoleCache.get(cacheKey) + if (cached && cached.expiresAt > Date.now()) { + return cached.credentials + } + + // Collapse concurrent refreshes for the same key into a single STS round-trip. + let inflight = inflightRoleRefreshes.get(cacheKey) + if (!inflight) { + inflight = refreshAssumedRoleCredentials(roleArn, externalId, region) + .then((entry) => { + assumedRoleCache.set(cacheKey, entry) + return entry + }) + .finally(() => { + inflightRoleRefreshes.delete(cacheKey) + }) + inflightRoleRefreshes.set(cacheKey, inflight) + } + + const entry = await inflight + return entry.credentials +} + +const refreshAssumedRoleCredentials = async ( + roleArn: string, + externalId: string, + region: string +): Promise => { const intermediaryARN = process.env.AMAZON_KINESIS_ACTIONS_ROLE_ADDRESS as string const intermediaryExternalId = process.env.AMAZON_KINESIS_ACTIONS_EXTERNAL_ID as string - const intermediaryCreds = await getSTSCredentials(intermediaryARN, intermediaryExternalId, region) - return getSTSCredentials(roleArn, externalId, region, intermediaryCreds) + const intermediary = await getSTSCredentials(intermediaryARN, intermediaryExternalId, region) + const target = await getSTSCredentials(roleArn, externalId, region, intermediary.credentials) + + // Prefer the actual STS expiration (minus a safety buffer) so the cache stays valid for the + // full session lifetime; fall back to a conservative default if STS omits it. + const ttl = target.expiration + ? target.expiration.getTime() - Date.now() - CREDENTIALS_EXPIRY_BUFFER_MS + : DEFAULT_CREDENTIALS_TTL_MS + + return { + credentials: target.credentials, + expiresAt: Date.now() + Math.max(ttl, 0) + } +} + +type STSCredentialsResult = { + credentials: AWSCredentials + expiration?: Date } -const getSTSCredentials = async (roleId: string, externalId: string, region: string, credentials?: AWSCredentials) => { +const getSTSCredentials = async ( + roleId: string, + externalId: string, + region: string, + credentials?: AWSCredentials +): Promise => { const options = { credentials, region: region } const stsClient = new STSClient(options) const roleSessionName: string = uuidv4() @@ -159,8 +241,11 @@ const getSTSCredentials = async (roleId: string, externalId: string, region: str } return { - accessKeyId: result.Credentials.AccessKeyId, - secretAccessKey: result.Credentials.SecretAccessKey, - sessionToken: result.Credentials.SessionToken + credentials: { + accessKeyId: result.Credentials.AccessKeyId, + secretAccessKey: result.Credentials.SecretAccessKey, + sessionToken: result.Credentials.SessionToken + }, + expiration: result.Credentials.Expiration } }