fix(@aws-amplify/datastore): add token to currentAuthenticatedUser fo…

…r OIDC (aws-amplify#6858) * add token to currentAuthenticatedUser * add api-graphql test for storage token
nubpro · Oct 2, 2020 · 4f66aac · 4f66aac
1 parent 4ed2ea6
commit 4f66aac
Show file tree

Hide file tree

Showing 6 changed files with 157 additions and 14 deletions.
diff --git a/packages/api-graphql/__tests__/GraphQLAPI-test.ts b/packages/api-graphql/__tests__/GraphQLAPI-test.ts
@@ -289,7 +289,7 @@ describe('API test', () => {
 			expect(spyon).toBeCalledWith(url, init);
 		});
 
-		test('happy-case-query-oidc', async () => {
+		test('happy-case-query-oidc with Cache token', async () => {
 			const spyonAuth = jest
 				.spyOn(Credentials, 'get')
 				.mockImplementationOnce(() => {
@@ -381,6 +381,109 @@ describe('API test', () => {
 			spyonCache.mockClear();
 		});
 
+		test('happy-case-query-oidc with auth storage federated token', async () => {
+			const spyonCredentials = jest
+				.spyOn(Credentials, 'get')
+				.mockImplementationOnce(() => {
+					return new Promise((res, rej) => {
+						res('cred');
+					});
+				});
+
+			const cache_config = {
+				capacityInBytes: 3000,
+				itemMaxSize: 800,
+				defaultTTL: 3000000,
+				defaultPriority: 5,
+				warningThreshold: 0.8,
+				storage: window.localStorage,
+			};
+
+			Cache.configure(cache_config);
+
+			const spyonCache = jest
+				.spyOn(Cache, 'getItem')
+				.mockImplementationOnce(() => {
+					return null;
+				});
+
+			const spyonAuth = jest
+				.spyOn(Auth, 'currentAuthenticatedUser')
+				.mockImplementationOnce(() => {
+					return new Promise((res, rej) => {
+						res({
+							name: 'federated user',
+							token: 'federated_token_from_storage',
+						});
+					});
+				});
+
+			const spyon = jest
+				.spyOn(RestClient.prototype, 'post')
+				.mockImplementationOnce((url, init) => {
+					return new Promise((res, rej) => {
+						res({});
+					});
+				});
+
+			const api = new API(config);
+			const url = 'https://appsync.amazonaws.com',
+				region = 'us-east-2',
+				variables = { id: '809392da-ec91-4ef0-b219-5238a8f942b2' };
+			api.configure({
+				aws_appsync_graphqlEndpoint: url,
+				aws_appsync_region: region,
+				aws_appsync_authenticationType: 'OPENID_CONNECT',
+			});
+			const GetEvent = `query GetEvent($id: ID! $nextToken: String) {
+				getEvent(id: $id) {
+					id
+					name
+					where
+					when
+					description
+					comments(nextToken: $nextToken) {
+						items {
+						commentId
+						content
+						createdAt
+						}
+					}
+				}
+			}`;
+
+			const doc = parse(GetEvent);
+			const query = print(doc);
+
+			const headers = {
+				Authorization: 'federated_token_from_storage',
+				'x-amz-user-agent': Constants.userAgent,
+			};
+
+			const body = {
+				query,
+				variables,
+			};
+
+			const init = {
+				headers,
+				body,
+				signerServiceInfo: {
+					service: 'appsync',
+					region,
+				},
+				cancellableToken: mockCancellableToken,
+			};
+
+			await api.graphql(graphqlOperation(GetEvent, variables));
+
+			expect(spyon).toBeCalledWith(url, init);
+
+			spyonCredentials.mockClear();
+			spyonCache.mockClear();
+			spyonAuth.mockClear();
+		});
+
 		test('multi-auth default case AWS_IAM, using API_KEY as auth mode', async () => {
 			expect.assertions(1);
 

diff --git a/packages/api-graphql/src/GraphQLAPI.ts b/packages/api-graphql/src/GraphQLAPI.ts
@@ -139,13 +139,22 @@ export class GraphQLAPIClass {
 				}
 				break;
 			case 'OPENID_CONNECT':
-				const federatedInfo = await this.Cache.getItem('federatedInfo');
-
-				if (!federatedInfo || !federatedInfo.token) {
+				let token;
+				// backwards compatibility
+				const federatedInfo = await Cache.getItem('federatedInfo');
+				if (federatedInfo) {
+					token = federatedInfo.token;
+				} else {
+					const currentUser = await Auth.currentAuthenticatedUser();
+					if (currentUser) {
+						token = currentUser.token;
+					}
+				}
+				if (!token) {
 					throw new Error('No federated jwt');
 				}
 				headers = {
-					Authorization: federatedInfo.token,
+					Authorization: token,
 				};
 				break;
 			case 'AMAZON_COGNITO_USER_POOLS':

diff --git a/packages/auth/__tests__/auth-unit-test.ts b/packages/auth/__tests__/auth-unit-test.ts
@@ -1473,7 +1473,10 @@ describe('auth unit test', () => {
 						setItem() {},
 						getItem() {
 							return JSON.stringify({
-								user: 'federated_user',
+								user: {
+									name: 'federated user',
+								},
+								token: '12345',
 							});
 						},
 						removeItem() {},
@@ -1487,7 +1490,10 @@ describe('auth unit test', () => {
 			});
 
 			expect.assertions(1);
-			expect(await auth.currentAuthenticatedUser()).toBe('federated_user');
+			expect(await auth.currentAuthenticatedUser()).toStrictEqual({
+				name: 'federated user',
+				token: '12345',
+			});
 
 			spyon.mockClear();
 		});

diff --git a/packages/auth/src/Auth.ts b/packages/auth/src/Auth.ts
@@ -1251,9 +1251,15 @@ export class AuthClass {
 		}
 
 		try {
-			federatedUser = JSON.parse(
+			const federatedInfo = JSON.parse(
 				this._storage.getItem('aws-amplify-federatedInfo')
-			).user;
+			);
+			if (federatedInfo) {
+				federatedUser = {
+					...federatedInfo.user,
+					token: federatedInfo.token,
+				};
+			}
 		} catch (e) {
 			logger.debug('cannot load federated user from auth storage');
 		}

diff --git a/packages/datastore/src/sync/processors/subscription.ts b/packages/datastore/src/sync/processors/subscription.ts
@@ -258,15 +258,25 @@ class SubscriptionProcessor {
 				}
 
 				try {
-					// retrieving token info from OIDC
+					let token;
+					// backwards compatibility
 					const federatedInfo = await Cache.getItem('federatedInfo');
-					const { token } = federatedInfo;
+					if (federatedInfo) {
+						token = federatedInfo.token;
+					} else {
+						const currentUser = await Auth.currentAuthenticatedUser();
+						if (currentUser) {
+							token = currentUser.token;
+						}
+					}
+
 					const payload = token.split('.')[1];
 
 					oidcTokenPayload = JSON.parse(
 						Buffer.from(payload, 'base64').toString('utf8')
 					);
 				} catch (err) {
+					logger.warn('error getting OIDC JWT', err);
 					// best effort to get oidc jwt
 				}
 

diff --git a/packages/pubsub/src/Providers/AWSAppSyncRealTimeProvider.ts b/packages/pubsub/src/Providers/AWSAppSyncRealTimeProvider.ts
@@ -780,13 +780,22 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 	}
 
 	private async _awsRealTimeOPENIDHeader({ host }) {
+		let token;
+		// backwards compatibility
 		const federatedInfo = await Cache.getItem('federatedInfo');
-
-		if (!federatedInfo || !federatedInfo.token) {
+		if (federatedInfo) {
+			token = federatedInfo.token;
+		} else {
+			const currentUser = await Auth.currentAuthenticatedUser();
+			if (currentUser) {
+				token = currentUser.token;
+			}
+		}
+		if (!token) {
 			throw new Error('No federated jwt');
 		}
 		return {
-			Authorization: federatedInfo.token,
+			Authorization: token,
 			host,
 		};
 	}