feat: Re-use Cognito identity id (#437)

Author: qhanam

src/dispatch/CognitoIdentityClient.ts

@@ -1,6 +1,8 @@
+/* eslint-disable no-underscore-dangle */
 import { HttpHandler, HttpRequest } from '@aws-sdk/protocol-http';
 import { Credentials } from '@aws-sdk/types';
 import { responseToJson } from './utils';
+import { IDENTITY_KEY } from '../utils/constants';
 
 const METHOD = 'POST';
 const CONTENT_TYPE = 'application/x-amz-json-1.1';
@@ -62,16 +64,40 @@ export class CognitoIdentityClient {
     }
 
     public getId = async (request: { IdentityPoolId: string }) => {
+        let getIdResponse: GetIdResponse | null = null;
+
+        try {
+            getIdResponse = JSON.parse(
+                localStorage.getItem(IDENTITY_KEY)!
+            ) as GetIdResponse | null;
+        } catch (e) {
+            // Ignore -- we will get a new identity Id from Cognito
+        }
+
+        if (getIdResponse && getIdResponse.IdentityId) {
+            return Promise.resolve(getIdResponse);
+        }
+
         try {
             const requestPayload = JSON.stringify(request);
             const idRequest = this.getHttpRequest(
                 GET_ID_TARGET,
                 requestPayload
             );
-            const { response } = await this.fetchRequestHandler.handle(
-                idRequest
-            );
-            return (await responseToJson(response)) as GetIdResponse;
+            const getIdResponse = (await responseToJson(
+                (
+                    await this.fetchRequestHandler.handle(idRequest)
+                ).response
+            )) as GetIdResponse;
+            try {
+                localStorage.setItem(
+                    IDENTITY_KEY,
+                    JSON.stringify({ IdentityId: getIdResponse.IdentityId })
+                );
+            } catch (e) {
+                // Ignore
+            }
+            return getIdResponse;
         } catch (e) {
             throw new Error(`CWR: Failed to retrieve Cognito identity: ${e}`);
         }
@@ -107,9 +133,11 @@ export class CognitoIdentityClient {
             const { response } = await this.fetchRequestHandler.handle(
                 credentialRequest
             );
-            const { Credentials } = (await responseToJson(
+            const credentialsResponse = (await responseToJson(
                 response
             )) as CredentialsResponse;
+            this.validateCredenentialsResponse(credentialsResponse);
+            const Credentials = credentialsResponse.Credentials;
             const { AccessKeyId, Expiration, SecretKey, SessionToken } =
                 Credentials;
             return {
@@ -125,6 +153,22 @@ export class CognitoIdentityClient {
         }
     };
 
+    private validateCredenentialsResponse = (cr: any) => {
+        if (
+            cr &&
+            cr.__type &&
+            (cr.__type === 'ResourceNotFoundException' ||
+                cr.__type === 'ValidationException')
+        ) {
+            // The request may have failed because of ValidationException or
+            // ResourceNotFoundException, which means the identity Id is bad. In
+            // any case, we invalidate the identity Id so the entire process can
+            // be re-tried.
+            localStorage.removeItem(IDENTITY_KEY);
+            throw new Error(`${cr.__type}: ${cr.message}`);
+        }
+    };
+
     private getHttpRequest = (target: string, payload: string) =>
         new HttpRequest({
             method: METHOD,

src/dispatch/__tests__/CognitoIdentityClient.test.ts

@@ -4,6 +4,7 @@ import { advanceTo } from 'jest-date-mock';
 import { CognitoIdentityClient } from '../CognitoIdentityClient';
 import { Credentials } from '@aws-sdk/types';
 import { getReadableStream } from '../../test-utils/test-utils';
+import { IDENTITY_KEY } from '../../utils/constants';
 
 const mockCredentials =
     '{ "IdentityId": "a", "Credentials": { "AccessKeyId": "x", "SecretKey": "y", "SessionToken": "z" } }';
@@ -22,6 +23,7 @@ describe('CognitoIdentityClient tests', () => {
     beforeEach(() => {
         advanceTo(0);
         fetchHandler.mockClear();
+        localStorage.clear();
 
         // @ts-ignore
         FetchHttpHandler.mockImplementation(() => {
@@ -74,7 +76,7 @@ describe('CognitoIdentityClient tests', () => {
         });
 
         // Assert
-        return expect(
+        await expect(
             client.getCredentialsForIdentity('my-fake-identity-id')
         ).rejects.toEqual(expected);
     });
@@ -175,4 +177,106 @@ describe('CognitoIdentityClient tests', () => {
             })
         ).rejects.toEqual(expected);
     });
+
+    test('when identity Id is retrieved from Cognito then next identity Id is retrieved from localStorage', async () => {
+        fetchHandler.mockResolvedValueOnce({
+            response: {
+                body: getReadableStream(mockIdCommand)
+            }
+        });
+
+        // Init
+        const client: CognitoIdentityClient = new CognitoIdentityClient({
+            fetchRequestHandler: new FetchHttpHandler(),
+            region: Utils.AWS_RUM_REGION
+        });
+
+        // Run
+        await client.getId({ IdentityPoolId: 'my-fake-identity-pool-id' });
+        const idCommand = await client.getId({
+            IdentityPoolId: 'my-fake-identity-pool-id'
+        });
+
+        // Assert
+        expect(fetchHandler).toHaveBeenCalledTimes(1);
+        expect(idCommand).toMatchObject({
+            IdentityId: 'mockId'
+        });
+    });
+
+    test('when getCredentialsForIdentity returns a ResourceNotFoundException then an error is thrown', async () => {
+        fetchHandler.mockResolvedValueOnce({
+            response: {
+                body: getReadableStream(
+                    '{"__type": "ResourceNotFoundException", "message": ""}'
+                )
+            }
+        });
+        const expected: Error = new Error(
+            `CWR: Failed to retrieve credentials for Cognito identity: Error: ResourceNotFoundException: `
+        );
+
+        // Init
+        const client: CognitoIdentityClient = new CognitoIdentityClient({
+            fetchRequestHandler: new FetchHttpHandler(),
+            region: Utils.AWS_RUM_REGION
+        });
+
+        // Assert
+        await expect(
+            client.getCredentialsForIdentity('my-fake-identity-id')
+        ).rejects.toEqual(expected);
+    });
+
+    test('when getCredentialsForIdentity returns a ValidationException then an error is thrown', async () => {
+        fetchHandler.mockResolvedValueOnce({
+            response: {
+                body: getReadableStream(
+                    '{"__type": "ValidationException", "message": ""}'
+                )
+            }
+        });
+        const expected: Error = new Error(
+            `CWR: Failed to retrieve credentials for Cognito identity: Error: ValidationException: `
+        );
+
+        // Init
+        const client: CognitoIdentityClient = new CognitoIdentityClient({
+            fetchRequestHandler: new FetchHttpHandler(),
+            region: Utils.AWS_RUM_REGION
+        });
+
+        // Assert
+        await expect(
+            client.getCredentialsForIdentity('my-fake-identity-id')
+        ).rejects.toEqual(expected);
+    });
+
+    test('when getCredentialsForIdentity returns a ResourceNotFoundException then identity id is removed from localStorage ', async () => {
+        localStorage.setItem(IDENTITY_KEY, 'my-fake-identity-id');
+
+        fetchHandler.mockResolvedValueOnce({
+            response: {
+                body: getReadableStream(
+                    '{"__type": "ResourceNotFoundException", "message": ""}'
+                )
+            }
+        });
+
+        // Init
+        const client: CognitoIdentityClient = new CognitoIdentityClient({
+            fetchRequestHandler: new FetchHttpHandler(),
+            region: Utils.AWS_RUM_REGION
+        });
+
+        // Run
+        try {
+            await client.getCredentialsForIdentity('my-fake-identity-id');
+        } catch (e) {
+            // Ignore
+        }
+
+        // Assert
+        expect(localStorage.getItem(IDENTITY_KEY)).toBe(null);
+    });
 });

src/utils/constants.ts

@@ -1,4 +1,5 @@
 export const CRED_KEY = 'cwr_c';
+export const IDENTITY_KEY = 'cwr_i';
 export const SESSION_COOKIE_NAME = 'cwr_s';
 export const USER_COOKIE_NAME = 'cwr_u';
 export const CRED_RENEW_MS = 30000;