From 6263f15b61a67e3f3d45c17b63bb96c2b7678a9a Mon Sep 17 00:00:00 2001
From: Jeroen Holthof <jeroen.holthof@kelenn.be>
Date: Wed, 27 Jan 2021 15:57:25 +0100
Subject: [PATCH] Add support for the customRoleArn (#1125)

The application can specify the role to assume, which will override/bypass the configured logic in the Identity Pool.

https://docs.aws.amazon.com/cognito/latest/developerguide/role-based-access-control.html#using-tokens-to-assign-roles-to-users
---
 packages/auth/src/Auth.ts        | 2 ++
 packages/auth/src/types/Auth.ts  | 2 ++
 packages/core/src/Credentials.ts | 7 +++++--
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/packages/auth/src/Auth.ts b/packages/auth/src/Auth.ts
index e90c2e763ba..b144e4c7b74 100644
--- a/packages/auth/src/Auth.ts
+++ b/packages/auth/src/Auth.ts
@@ -151,6 +151,7 @@ export class AuthClass {
 			identityPoolRegion,
 			clientMetadata,
 			endpoint,
+			customRoleArn,
 		} = this._config;
 
 		if (!this._config.storage) {
@@ -195,6 +196,7 @@ export class AuthClass {
 			identityPoolId,
 			refreshHandlers,
 			storage: this._storage,
+			customRoleArn,
 		});
 
 		// initiailize cognitoauth client if hosted ui options provided
diff --git a/packages/auth/src/types/Auth.ts b/packages/auth/src/types/Auth.ts
index 060f54454ed..226581698a4 100644
--- a/packages/auth/src/types/Auth.ts
+++ b/packages/auth/src/types/Auth.ts
@@ -15,6 +15,7 @@ import {
 	ICookieStorageData,
 	ICognitoStorage,
 	CognitoUserAttribute,
+	CognitoIdToken,
 } from 'amazon-cognito-identity-js';
 
 /**
@@ -51,6 +52,7 @@ export interface AuthOptions {
 	identityPoolRegion?: string;
 	clientMetadata?: any;
 	endpoint?: string;
+	customRoleArn?: (token : CognitoIdToken | string) => string,
 }
 
 export enum CognitoHostedUIIdentityProvider {
diff --git a/packages/core/src/Credentials.ts b/packages/core/src/Credentials.ts
index e518db3d61e..c5a0f3f4531 100644
--- a/packages/core/src/Credentials.ts
+++ b/packages/core/src/Credentials.ts
@@ -362,7 +362,7 @@ export class CredentialsClass {
 		const logins = {};
 		logins[domain] = token;
 
-		const { identityPoolId, region } = this._config;
+		const { identityPoolId, region, customRoleArn } = this._config;
 		if (!identityPoolId) {
 			logger.debug('No Cognito Federated Identity pool provided');
 			return Promise.reject('No Cognito Federated Identity pool provided');
@@ -385,6 +385,7 @@ export class CredentialsClass {
 				identityId: identity_id,
 				logins,
 				client: cognitoClient,
+				customRoleArn: customRoleArn && customRoleArn(token),
 			};
 			credentials = fromCognitoIdentity(cognitoIdentityParams)();
 		} else {
@@ -392,6 +393,7 @@ export class CredentialsClass {
 				logins,
 				identityPoolId,
 				client: cognitoClient,
+				customRoleArn: customRoleArn && customRoleArn(token),
 			};
 			credentials = fromCognitoIdentityPool(cognitoIdentityParams)();
 		}
@@ -401,7 +403,7 @@ export class CredentialsClass {
 	private _setCredentialsFromSession(session): Promise<ICredentials> {
 		logger.debug('set credentials from session');
 		const idToken = session.getIdToken().getJwtToken();
-		const { region, userPoolId, identityPoolId } = this._config;
+		const { region, userPoolId, identityPoolId, customRoleArn } = this._config;
 		if (!identityPoolId) {
 			logger.debug('No Cognito Federated Identity pool provided');
 			return Promise.reject('No Cognito Federated Identity pool provided');
@@ -440,6 +442,7 @@ export class CredentialsClass {
 				client: cognitoClient,
 				logins,
 				identityId: IdentityId,
+				customRoleArn: customRoleArn && customRoleArn(idToken)
 			};
 
 			const credentialsFromCognitoIdentity = fromCognitoIdentity(