feat: add support for AWS_LAMBDA auth type (#1412)

Author: richardmcclellan

aws-api-appsync/src/main/java/com/amplifyframework/api/aws/AuthorizationType.java

@@ -56,6 +56,11 @@ public enum AuthorizationType {
      */
     AMAZON_COGNITO_USER_POOLS,
 
+    /**
+     * Control access by implementing your own API authorization logic within an AWS Lambda function.
+     */
+    AWS_LAMBDA,
+
     /**
      * No authorization.
      */
@@ -107,6 +112,8 @@ public static AuthorizationType from(AuthStrategy.Provider authRuleProvider) {
                 return AWS_IAM;
             case API_KEY:
                 return API_KEY;
+            case FUNCTION:
+                return AWS_LAMBDA;
             default:
                 throw new IllegalArgumentException("No such authorization type: " + authRuleProvider.name());
         }

aws-api/src/main/java/com/amplifyframework/api/aws/AWSApiPlugin.java

@@ -25,6 +25,7 @@
 import com.amplifyframework.api.ApiPlugin;
 import com.amplifyframework.api.aws.auth.ApiRequestDecoratorFactory;
 import com.amplifyframework.api.aws.auth.AuthRuleRequestDecorator;
+import com.amplifyframework.api.aws.auth.RequestDecorator;
 import com.amplifyframework.api.aws.operation.AWSRestOperation;
 import com.amplifyframework.api.events.ApiEndpointStatusChangeEvent;
 import com.amplifyframework.api.events.ApiEndpointStatusChangeEvent.ApiEndpointStatus;
@@ -144,26 +145,35 @@ public void configure(
                 configurator.applyConfiguration(okHttpClientBuilder);
             }
 
+            final ApiRequestDecoratorFactory requestDecoratorFactory = new ApiRequestDecoratorFactory(authProvider,
+                    apiConfiguration.getAuthorizationType(),
+                    apiConfiguration.getRegion(),
+                    apiConfiguration.getEndpointType(),
+                    apiConfiguration.getApiKey());
+
             ClientDetails clientDetails = null;
             if (EndpointType.REST.equals(endpointType)) {
-                final InterceptorFactory interceptorFactory =
-                    new AppSyncSigV4SignerInterceptorFactory(authProvider);
                 if (apiConfiguration.getAuthorizationType() != AuthorizationType.NONE) {
-                    okHttpClientBuilder.addInterceptor(interceptorFactory.create(apiConfiguration));
+                    AuthorizationType authorizationType = apiConfiguration.getAuthorizationType();
+                    RequestDecorator decorator = requestDecoratorFactory.forAuthType(authorizationType);
+                    okHttpClientBuilder.addInterceptor(chain -> {
+                        try {
+                            return chain.proceed(decorator.decorate(chain.request()));
+                        } catch (ApiException.ApiAuthException apiAuthException) {
+                            throw new IOException("Failed to decorate request for authorization.", apiAuthException);
+                        }
+                    });
                 }
-                clientDetails = new ClientDetails(apiConfiguration, okHttpClientBuilder.build(), null, null);
+                clientDetails = new ClientDetails(apiConfiguration,
+                                                  okHttpClientBuilder.build(),
+                                                  null,
+                                                  requestDecoratorFactory);
                 restApis.add(apiName);
             } else if (EndpointType.GRAPHQL.equals(endpointType)) {
                 final SubscriptionAuthorizer subscriptionAuthorizer =
                     new SubscriptionAuthorizer(apiConfiguration, authProvider);
                 final SubscriptionEndpoint subscriptionEndpoint =
                     new SubscriptionEndpoint(apiConfiguration, gqlResponseFactory, subscriptionAuthorizer);
-                final ApiRequestDecoratorFactory requestDecoratorFactory =
-                    new ApiRequestDecoratorFactory(authProvider,
-                                                   apiConfiguration.getAuthorizationType(),
-                                                   apiConfiguration.getRegion(),
-                                                   apiConfiguration.getApiKey());
-
                 clientDetails = new ClientDetails(apiConfiguration,
                                                   okHttpClientBuilder.build(),
                                                   subscriptionEndpoint,

aws-api/src/main/java/com/amplifyframework/api/aws/ApiAuthProviders.java

@@ -19,6 +19,7 @@
 
 import com.amplifyframework.api.aws.sigv4.ApiKeyAuthProvider;
 import com.amplifyframework.api.aws.sigv4.CognitoUserPoolsAuthProvider;
+import com.amplifyframework.api.aws.sigv4.FunctionAuthProvider;
 import com.amplifyframework.api.aws.sigv4.OidcAuthProvider;
 
 import com.amazonaws.auth.AWSCredentialsProvider;
@@ -34,12 +35,14 @@ public final class ApiAuthProviders {
     private final AWSCredentialsProvider awsCredentialsProvider;
     private final CognitoUserPoolsAuthProvider cognitoUserPoolsAuthProvider;
     private final OidcAuthProvider oidcAuthProvider;
+    private final FunctionAuthProvider functionAuthProvider;
 
     private ApiAuthProviders(Builder builder) {
         this.apiKeyAuthProvider = builder.getApiKeyAuthProvider();
         this.awsCredentialsProvider = builder.getAWSCredentialsProvider();
         this.oidcAuthProvider = builder.getOidcAuthProvider();
         this.cognitoUserPoolsAuthProvider = builder.getCognitoUserPoolsAuthProvider();
+        this.functionAuthProvider = builder.getFunctionAuthProvider();
     }
 
     /**
@@ -74,6 +77,14 @@ public CognitoUserPoolsAuthProvider getCognitoUserPoolsAuthProvider() {
         return this.cognitoUserPoolsAuthProvider;
     }
 
+    /**
+     * Gets the custom auth provider.
+     * @return an implementation of {@link FunctionAuthProvider}
+     */
+    public FunctionAuthProvider getFunctionAuthProvider() {
+        return this.functionAuthProvider;
+    }
+
     /**
      * Statically gets the builder for conveniently
      * configuring an immutable instance of {@link ApiAuthProviders}.
@@ -101,6 +112,7 @@ public static final class Builder {
         private AWSCredentialsProvider awsCredentialsProvider;
         private CognitoUserPoolsAuthProvider cognitoUserPoolsAuthProvider;
         private OidcAuthProvider oidcAuthProvider;
+        private FunctionAuthProvider functionAuthProvider;
 
         /**
          * Assigns an API key Auth provider.
@@ -142,6 +154,16 @@ public ApiAuthProviders.Builder oidcAuthProvider(@NonNull OidcAuthProvider provi
             return ApiAuthProviders.Builder.this;
         }
 
+        /**
+         * Assigns a serverless function (e.g. AWS Lambda) auth provider.
+         * @param provider an instance of {@link FunctionAuthProvider}
+         * @return this builder object for chaining
+         */
+        public ApiAuthProviders.Builder functionAuthProvider(@NonNull FunctionAuthProvider provider) {
+            ApiAuthProviders.Builder.this.functionAuthProvider = Objects.requireNonNull(provider);
+            return ApiAuthProviders.Builder.this;
+        }
+
         /**
          * Creates an immutable instance of {@link ApiAuthProviders}
          * configured to this builder instance.
@@ -166,5 +188,9 @@ CognitoUserPoolsAuthProvider getCognitoUserPoolsAuthProvider() {
         OidcAuthProvider getOidcAuthProvider() {
             return ApiAuthProviders.Builder.this.oidcAuthProvider;
         }
+
+        FunctionAuthProvider getFunctionAuthProvider() {
+            return ApiAuthProviders.Builder.this.functionAuthProvider;
+        }
     }
 }

aws-api/src/main/java/com/amplifyframework/api/aws/AppSyncSigV4SignerInterceptorFactory.java

@@ -24,6 +24,7 @@
 import com.amplifyframework.api.aws.sigv4.AppSyncV4Signer;
 import com.amplifyframework.api.aws.sigv4.CognitoUserPoolsAuthProvider;
 import com.amplifyframework.api.aws.sigv4.DefaultCognitoUserPoolsAuthProvider;
+import com.amplifyframework.api.aws.sigv4.FunctionAuthProvider;
 import com.amplifyframework.api.aws.sigv4.OidcAuthProvider;
 import com.amplifyframework.api.graphql.GraphQLRequest;
 import com.amplifyframework.core.Amplify;
@@ -105,6 +106,19 @@ private JSONObject createHeaders(GraphQLRequest<?> request,
                     };
                 }
                 return forOidc(oidcProvider);
+            case AWS_LAMBDA:
+                FunctionAuthProvider functionAuthProvider = authProviders.getFunctionAuthProvider();
+                if (functionAuthProvider == null) {
+                    functionAuthProvider = () -> {
+                        throw new ApiAuthException(
+                                "FunctionAuthProvider interface is not implemented.",
+                                "Please implement FunctionAuthProvider interface to return " +
+                                        "appropriate token from the appropriate service."
+                        );
+                    };
+                }
+                return forAwsLambda(functionAuthProvider);
+
             case NONE:
             default:
                 return new JSONObject();
@@ -154,6 +168,20 @@ private JSONObject forOidc(OidcAuthProvider oidcProvider) throws ApiException {
         }
     }
 
+    private JSONObject forAwsLambda(FunctionAuthProvider functionAuthProvider) throws ApiException {
+        try {
+            return new JSONObject()
+                    .put("host", getHost())
+                    .put("Authorization", functionAuthProvider.getLatestAuthToken());
+        } catch (JSONException jsonException) {
+            // This error should never be thrown
+            throw new ApiException(
+                    "Error constructing the authorization json for the AWS_LAMBDA auth type.",
+                    jsonException, AmplifyException.REPORT_BUG_TO_AWS_SUGGESTION
+            );
+        }
+    }
+
     private JSONObject forIam(
             AWSCredentialsProvider credentialsProvider,
             GraphQLRequest<?> request,