Caching caller identity

packages/sns/lib/utils/stsUtils.spec.ts

@@ -1,8 +1,9 @@
 import type { SNSClient } from '@aws-sdk/client-sns'
 import type { STSClient } from '@aws-sdk/client-sts'
-import { beforeAll, beforeEach, describe, expect, it } from 'vitest'
+import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
 import { registerDependencies } from '../../test/utils/testContext'
 import { assertTopic, deleteTopic } from './snsUtils'
+import { buildTopicArn, clearCachedCallerIdentity } from './stsUtils'
 
 describe('stsUtils', () => {
   let stsClient: STSClient
@@ -19,17 +20,37 @@ describe('stsUtils', () => {
 
     beforeEach(async () => {
       await deleteTopic(snsClient, stsClient, topicName)
+      clearCachedCallerIdentity()
     })
 
     it('build ARN for topic', async () => {
-      const expectedTopicArn = `arn:aws:sns:eu-west-1:000000000000:${topicName}`
-      expect(expectedTopicArn).toMatchInlineSnapshot(
+      const buildedTopicArn = await buildTopicArn(stsClient, topicName)
+      expect(buildedTopicArn).toMatchInlineSnapshot(
         `"arn:aws:sns:eu-west-1:000000000000:my-test-topic"`,
       )
 
       // creating real topic to make sure arn is correct
       const topicArn = await assertTopic(snsClient, stsClient, { Name: topicName })
-      expect(topicArn).toBe(expectedTopicArn)
+      expect(topicArn).toBe(buildedTopicArn)
+    })
+
+    it('should be able to handle parallel calls getting caller identity only once', async () => {
+      const stsClientSpy = vi.spyOn(stsClient, 'send')
+
+      const result = await Promise.all([
+        buildTopicArn(stsClient, topicName),
+        buildTopicArn(stsClient, topicName),
+        buildTopicArn(stsClient, topicName),
+      ])
+
+      expect(result).toMatchInlineSnapshot(`
+        [
+          "arn:aws:sns:eu-west-1:000000000000:my-test-topic",
+          "arn:aws:sns:eu-west-1:000000000000:my-test-topic",
+          "arn:aws:sns:eu-west-1:000000000000:my-test-topic",
+        ]
+      `)
+      expect(stsClientSpy).toHaveBeenCalledOnce()
     })
   })
 })

packages/sns/lib/utils/stsUtils.ts

@@ -1,14 +1,43 @@
-import { GetCallerIdentityCommand, type STSClient } from '@aws-sdk/client-sts'
+import {
+  GetCallerIdentityCommand,
+  type GetCallerIdentityCommandOutput,
+  type STSClient,
+} from '@aws-sdk/client-sts'
+
+let callerIdentityPromise: Promise<void> | undefined
+let callerIdentityCached: GetCallerIdentityCommandOutput | undefined
 
 /**
  * Manually builds the ARN of a topic based on the current AWS account and the topic name.
  * It follows the following pattern: arn:aws:sns:<region>:<account-id>:<topic-name>
  * Doc -> https://docs.aws.amazon.com/IAM/latest/UserGuide/reference-arns.html
  */
 export const buildTopicArn = async (client: STSClient, topicName: string) => {
-  const identityResponse = await client.send(new GetCallerIdentityCommand({}))
+  const identityResponse = await getAndCacheCallerIdentity(client)
   const region =
     typeof client.config.region === 'string' ? client.config.region : await client.config.region()
 
   return `arn:aws:sns:${region}:${identityResponse.Account}:${topicName}`
 }
+
+export const clearCachedCallerIdentity = () => {
+  callerIdentityPromise = undefined
+  callerIdentityCached = undefined
+}
+
+const getAndCacheCallerIdentity = async (
+  client: STSClient,
+): Promise<GetCallerIdentityCommandOutput> => {
+  if (!callerIdentityCached) {
+    if (!callerIdentityPromise) {
+      callerIdentityPromise = client.send(new GetCallerIdentityCommand({})).then((response) => {
+        callerIdentityCached = response
+      })
+    }
+
+    await callerIdentityPromise
+    callerIdentityPromise = undefined
+  }
+
+  return callerIdentityCached!
+}