feat(aws-kinesis): cache assumed-role credentials to avoid STS throttling

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()
+      }
+    })
+  })
 })

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<string, AssumedRoleCacheEntry>()
+
+// 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<string, Promise<AssumedRoleCacheEntry>>()
+
+// 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<AWSCredentials> => {
+  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<AssumedRoleCacheEntry> => {
   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<STSCredentialsResult> => {
   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
   }
 }