Merge pull request #8616 from aws-amplify/lambda-auth

Lambda authorizers support (API, DataStore, PubSub) Doing a normal PR merge instead of squashing in this case because I want API, PubSub, and DataStore to all get minor version bumps as a result of merging this PR.
aws-amplify · Jul 26, 2021 · e268efc · e268efc
2 parents 908cd16 + 52d43cc
commit e268efc
Show file tree

Hide file tree

Showing 14 changed files with 709 additions and 1,192 deletions.
diff --git a/packages/api-graphql/__tests__/GraphQLAPI-test.ts b/packages/api-graphql/__tests__/GraphQLAPI-test.ts
diff --git a/packages/api-graphql/src/GraphQLAPI.ts b/packages/api-graphql/src/GraphQLAPI.ts
@@ -32,9 +32,14 @@ const USER_AGENT_HEADER = 'x-amz-user-agent';
 
 const logger = new Logger('GraphQLAPI');
 
-export const graphqlOperation = (query, variables = {}) => ({
+export const graphqlOperation = (
+	query,
+	variables = {},
+	authToken?: string
+) => ({
 	query,
 	variables,
+	authToken,
 });
 
 /**
@@ -113,7 +118,10 @@ export class GraphQLAPIClass {
 		}
 	}
 
-	private async _headerBasedAuth(defaultAuthenticationType?) {
+	private async _headerBasedAuth(
+		defaultAuthenticationType?,
+		additionalHeaders: { [key: string]: string } = {}
+	) {
 		const {
 			aws_appsync_authenticationType,
 			aws_appsync_apiKey: apiKey,
@@ -171,6 +179,14 @@ export class GraphQLAPIClass {
 					throw new Error(GraphQLAuthError.NO_CURRENT_USER);
 				}
 				break;
+			case 'AWS_LAMBDA':
+				if (!additionalHeaders.Authorization) {
+					throw new Error(GraphQLAuthError.NO_AUTH_TOKEN);
+				}
+				headers = {
+					Authorization: additionalHeaders.Authorization,
+				};
+				break;
 			default:
 				headers = {
 					Authorization: null,
@@ -202,7 +218,7 @@ export class GraphQLAPIClass {
 	 * @returns {Promise<GraphQLResult> | Observable<object>}
 	 */
 	graphql(
-		{ query: paramQuery, variables = {}, authMode }: GraphQLOptions,
+		{ query: paramQuery, variables = {}, authMode, authToken }: GraphQLOptions,
 		additionalHeaders?: { [key: string]: string }
 	) {
 		const query =
@@ -217,14 +233,21 @@ export class GraphQLAPIClass {
 			operation: operationType,
 		} = operationDef as OperationDefinitionNode;
 
+		const headers = additionalHeaders || {};
+
+		// if an authorization header is set, have the explicit authToken take precedence
+		if (authToken) {
+			headers.Authorization = authToken;
+		}
+
 		switch (operationType) {
 			case 'query':
 			case 'mutation':
 				const cancellableToken = this._api.getCancellableToken();
 				const initParams = { cancellableToken };
 				const responsePromise = this._graphql(
 					{ query, variables, authMode },
-					additionalHeaders,
+					headers,
 					initParams
 				);
 				this._api.updateRequestToBeCancellable(
@@ -233,10 +256,7 @@ export class GraphQLAPIClass {
 				);
 				return responsePromise;
 			case 'subscription':
-				return this._graphqlSubscribe(
-					{ query, variables, authMode },
-					additionalHeaders
-				);
+				return this._graphqlSubscribe({ query, variables, authMode }, headers);
 		}
 
 		throw new Error(`invalid operation type: ${operationType}`);
@@ -260,10 +280,11 @@ export class GraphQLAPIClass {
 		} = this._options;
 
 		const headers = {
-			...(!customGraphqlEndpoint && (await this._headerBasedAuth(authMode))),
+			...(!customGraphqlEndpoint &&
+				(await this._headerBasedAuth(authMode, additionalHeaders))),
 			...(customGraphqlEndpoint &&
 				(customEndpointRegion
-					? await this._headerBasedAuth(authMode)
+					? await this._headerBasedAuth(authMode, additionalHeaders)
 					: { Authorization: null })),
 			...(await graphql_headers({ query, variables })),
 			...additionalHeaders,
@@ -344,7 +365,12 @@ export class GraphQLAPIClass {
 	}
 
 	private _graphqlSubscribe(
-		{ query, variables, authMode: defaultAuthenticationType }: GraphQLOptions,
+		{
+			query,
+			variables,
+			authMode: defaultAuthenticationType,
+			authToken,
+		}: GraphQLOptions,
 		additionalHeaders = {}
 	): Observable<any> {
 		const {
@@ -368,6 +394,7 @@ export class GraphQLAPIClass {
 				variables,
 				graphql_headers,
 				additionalHeaders,
+				authToken,
 			});
 		} else {
 			logger.debug('No pubsub module applied for subscription');

diff --git a/packages/api-graphql/src/types/index.ts b/packages/api-graphql/src/types/index.ts
@@ -18,13 +18,15 @@ export interface GraphQLOptions {
 	query: string | DocumentNode;
 	variables?: object;
 	authMode?: GRAPHQL_AUTH_MODE;
+	authToken?: string;
 }
 
 export enum GRAPHQL_AUTH_MODE {
 	API_KEY = 'API_KEY',
 	AWS_IAM = 'AWS_IAM',
 	OPENID_CONNECT = 'OPENID_CONNECT',
 	AMAZON_COGNITO_USER_POOLS = 'AMAZON_COGNITO_USER_POOLS',
+	AWS_LAMBDA = 'AWS_LAMBDA',
 }
 
 export interface GraphQLResult<T = object> {
@@ -40,4 +42,5 @@ export enum GraphQLAuthError {
 	NO_CURRENT_USER = 'No current user',
 	NO_CREDENTIALS = 'No credentials',
 	NO_FEDERATED_JWT = 'No federated jwt',
+	NO_AUTH_TOKEN = 'No auth token specified',
 }
diff --git a/packages/datastore/__tests__/authStrategies.test.ts b/packages/datastore/__tests__/authStrategies.test.ts
@@ -10,6 +10,11 @@ import { NAMESPACES } from '../src/util';
 describe('Auth Strategies', () => {
 	describe('multiAuthStrategy', () => {
 		const rules = {
+			function: {
+				provider: ModelAttributeAuthProvider.FUNCTION,
+				allow: ModelAttributeAuthAllow.CUSTOM,
+				operations: ['create', 'update', 'delete', 'read'],
+			},
 			owner: {
 				provider: ModelAttributeAuthProvider.USER_POOLS,
 				ownerField: 'owner',
@@ -68,6 +73,21 @@ describe('Auth Strategies', () => {
 			},
 		};
 
+		test('function', async () => {
+			const authRules = [rules.function];
+			await testMultiAuthStrategy({
+				authRules,
+				hasAuthenticatedUser: true,
+				result: ['AWS_LAMBDA'],
+			});
+
+			await testMultiAuthStrategy({
+				authRules,
+				hasAuthenticatedUser: false,
+				result: ['AWS_LAMBDA'],
+			});
+		});
+
 		test('owner', async () => {
 			const authRules = [rules.owner];
 			await testMultiAuthStrategy({
@@ -353,6 +373,31 @@ describe('Auth Strategies', () => {
 			});
 		});
 
+		test('function/owner/public IAM/API key', async () => {
+			const authRules = [
+				rules.function,
+				rules.owner,
+				rules.publicIAM,
+				rules.publicAPIKeyExplicit,
+			];
+			await testMultiAuthStrategy({
+				authRules,
+				hasAuthenticatedUser: true,
+				result: [
+					'AWS_LAMBDA',
+					'AMAZON_COGNITO_USER_POOLS',
+					'AWS_IAM',
+					'API_KEY',
+				],
+			});
+
+			await testMultiAuthStrategy({
+				authRules,
+				hasAuthenticatedUser: false,
+				result: ['AWS_LAMBDA', 'AWS_IAM', 'API_KEY'],
+			});
+		});
+
 		test('duplicates', async () => {
 			const authRules = [
 				rules.owner,

diff --git a/packages/datastore/__tests__/helpers.ts b/packages/datastore/__tests__/helpers.ts
@@ -1,4 +1,10 @@
-import { ModelInit, MutableModel, Schema, InternalSchema } from '../src/types';
+import {
+	ModelInit,
+	MutableModel,
+	Schema,
+	InternalSchema,
+	SchemaModel,
+} from '../src/types';
 
 export declare class Model {
 	public readonly id: string;