fix: Add getId to enhanced authflow (#433)

aws-observability · Aug 18, 2023 · 8b95de0 · 8b95de0
1 parent ac47136
commit 8b95de0
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 67 deletions.
diff --git a/src/dispatch/CognitoIdentityClient.ts b/src/dispatch/CognitoIdentityClient.ts
@@ -29,6 +29,7 @@ interface CognitoCredentials {
     AccessKeyId: string;
     Expiration: number;
     SecretAccessKey: string;
+    SecretKey: string;
     SessionToken: string;
 }
 
@@ -46,12 +47,6 @@ interface GetIdResponse {
     IdentityId: string;
 }
 
-export const fromCognitoIdentityPool = (
-    params: CognitoProviderParameters
-): (() => Promise<Credentials>) => {
-    return () => params.client.getCredentialsForIdentity(params.identityPoolId);
-};
-
 export declare type CognitoIdentityClientConfig = {
     fetchRequestHandler: HttpHandler;
     region?: string;
@@ -115,11 +110,11 @@ export class CognitoIdentityClient {
             const { Credentials } = (await responseToJson(
                 response
             )) as CredentialsResponse;
-            const { AccessKeyId, Expiration, SecretAccessKey, SessionToken } =
+            const { AccessKeyId, Expiration, SecretKey, SessionToken } =
                 Credentials;
             return {
                 accessKeyId: AccessKeyId as string,
-                secretAccessKey: SecretAccessKey as string,
+                secretAccessKey: SecretKey as string,
                 sessionToken: SessionToken as string,
                 expiration: new Date(Expiration * 1000)
             };

diff --git a/src/dispatch/EnhancedAuthentication.ts b/src/dispatch/EnhancedAuthentication.ts
@@ -1,9 +1,6 @@
-import {
-    CognitoIdentityClient,
-    fromCognitoIdentityPool
-} from './CognitoIdentityClient';
+import { CognitoIdentityClient } from './CognitoIdentityClient';
 import { Config } from '../orchestration/Orchestration';
-import { CredentialProvider, Credentials } from '@aws-sdk/types';
+import { Credentials } from '@aws-sdk/types';
 import { FetchHttpHandler } from '@aws-sdk/fetch-http-handler';
 import { CRED_KEY, CRED_RENEW_MS } from '../utils/constants';
 
@@ -96,7 +93,6 @@ export class EnhancedAuthentication {
                     // The credentials have expired.
                     return reject();
                 }
-                this.credentials = credentials;
                 resolve(credentials);
             });
         };
@@ -111,21 +107,26 @@ export class EnhancedAuthentication {
      */
     private AnonymousCognitoCredentialsProvider =
         async (): Promise<Credentials> => {
-            const credentialProvider: CredentialProvider =
-                fromCognitoIdentityPool({
-                    client: this.cognitoIdentityClient,
-                    identityPoolId: this.config.identityPoolId as string
-                });
+            return this.cognitoIdentityClient
+                .getId({ IdentityPoolId: this.config.identityPoolId as string })
+                .then((getIdResponse) =>
+                    this.cognitoIdentityClient.getCredentialsForIdentity(
+                        getIdResponse.IdentityId
+                    )
+                )
+                .then((credentials: Credentials) => {
+                    this.credentials = credentials;
+                    try {
+                        localStorage.setItem(
+                            CRED_KEY,
+                            JSON.stringify(credentials)
+                        );
+                    } catch (e) {
+                        // Ignore
+                    }
 
-            return credentialProvider().then((credentials) => {
-                this.credentials = credentials;
-                try {
-                    localStorage.setItem(CRED_KEY, JSON.stringify(credentials));
-                } catch (e) {
-                    // Ignore
-                }
-                return credentials;
-            });
+                    return credentials;
+                });
         };
 
     private renewCredentials(): boolean {

diff --git a/src/dispatch/__tests__/CognitoIdentityClient.test.ts b/src/dispatch/__tests__/CognitoIdentityClient.test.ts
@@ -6,7 +6,7 @@ import { Credentials } from '@aws-sdk/types';
 import { getReadableStream } from '../../test-utils/test-utils';
 
 const mockCredentials =
-    '{ "IdentityId": "a", "Credentials": { "AccessKeyId": "x", "SecretAccessKey": "y", "SessionToken": "z" } }';
+    '{ "IdentityId": "a", "Credentials": { "AccessKeyId": "x", "SecretKey": "y", "SessionToken": "z" } }';
 const mockToken = '{"IdentityId": "mockId", "Token": "mockToken"}';
 const mockIdCommand = '{"IdentityId": "mockId"}';
 

diff --git a/src/dispatch/__tests__/EnhancedAuthentication.test.ts b/src/dispatch/__tests__/EnhancedAuthentication.test.ts
@@ -1,14 +1,11 @@
 import { CRED_KEY } from '../../utils/constants';
-import { Credentials } from '@aws-sdk/types';
 import { EnhancedAuthentication } from '../EnhancedAuthentication';
-import { fromCognitoIdentityPool } from '../CognitoIdentityClient';
 import { DEFAULT_CONFIG } from '../../test-utils/test-utils';
 
 const mockGetId = jest.fn();
 const getCredentials = jest.fn();
 
 jest.mock('../CognitoIdentityClient', () => ({
-    fromCognitoIdentityPool: jest.fn(),
     CognitoIdentityClient: jest.fn().mockImplementation(() => ({
         getId: mockGetId,
         getCredentialsForIdentity: getCredentials
@@ -31,18 +28,6 @@ describe('EnhancedAuthentication tests', () => {
             sessionToken: 'z',
             expiration: new Date(Date.now() + 3600 * 1000)
         });
-        (fromCognitoIdentityPool as any).mockReset();
-        (fromCognitoIdentityPool as any).mockReturnValue(
-            () =>
-                new Promise<Credentials>((resolve) =>
-                    resolve({
-                        accessKeyId: 'x',
-                        secretAccessKey: 'y',
-                        sessionToken: 'z',
-                        expiration: new Date(Date.now() + 3600 * 1000)
-                    })
-                )
-        );
         localStorage.removeItem(CRED_KEY);
     });
 
@@ -169,29 +154,19 @@ describe('EnhancedAuthentication tests', () => {
     test('when credential is retrieved from basic auth then next credential is retrieved from localStorage', async () => {
         // Init
         const expiration = new Date(Date.now() + 3600 * 1000);
-        (fromCognitoIdentityPool as any)
-            .mockReturnValueOnce(
-                () =>
-                    new Promise<Credentials>((resolve) =>
-                        resolve({
-                            accessKeyId: 'a',
-                            secretAccessKey: 'b',
-                            sessionToken: 'c',
-                            expiration
-                        })
-                    )
-            )
-            .mockReturnValueOnce(
-                () =>
-                    new Promise<Credentials>((resolve) =>
-                        resolve({
-                            accessKeyId: 'x',
-                            secretAccessKey: 'y',
-                            sessionToken: 'z',
-                            expiration
-                        })
-                    )
-            );
+        getCredentials
+            .mockResolvedValueOnce({
+                accessKeyId: 'a',
+                expiration,
+                secretAccessKey: 'b',
+                sessionToken: 'c'
+            })
+            .mockResolvedValueOnce({
+                accessKeyId: 'x',
+                expiration,
+                secretAccessKey: 'y',
+                sessionToken: 'z'
+            });
 
         const auth = new EnhancedAuthentication({
             ...DEFAULT_CONFIG,