Async cognito support (#874)

Co-authored-by: Bret Ambrose <[email protected]>

Author: bretambrose

crt/aws-c-auth

@@ -9,6 +9,8 @@
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
 
 import software.amazon.awssdk.crt.http.HttpHeader;
 import software.amazon.awssdk.crt.http.HttpProxyOptions;
@@ -47,6 +49,7 @@ static public class CognitoCredentialsProviderBuilder {
         private String identity;
         private String customRoleArn;
         private ArrayList<CognitoLoginTokenPair> logins = new ArrayList<CognitoLoginTokenPair>();
+        private CognitoLoginTokenSource loginTokenSource;
 
         private TlsContext tlsContext;
         private ClientBootstrap clientBootstrap;
@@ -148,6 +151,23 @@ public CognitoCredentialsProviderBuilder withHttpProxyOptions(HttpProxyOptions h
 
         HttpProxyOptions getHttpProxyOptions() { return httpProxyOptions; }
 
+        /**
+         * Sets a login token source for the credentials provider.  The login token source will be used to
+         * gather additional login tokens to submit as part of the HTTP request sent to Cognito.  A login token source
+         * allows you to dynamically add login tokens on a per-request basis.  Using a login token source requires
+         * you to follow certain requirements in order to avoid undesirable behavior.  See the documentation for
+         * `CognitoLoginTokenSource` for further details.
+         *
+         * @param loginTokenSource object to source login tokens from before every HTTP request to Cognito
+         * @return The current builder
+         */
+        public CognitoCredentialsProviderBuilder withLoginTokenSource(CognitoLoginTokenSource loginTokenSource) {
+            this.loginTokenSource = loginTokenSource;
+
+            return this;
+        }
+
+        CognitoLoginTokenSource getLoginTokenSource() { return loginTokenSource; }
 
         /**
          * Creates a new Cognito credentials provider, based on this builder's configuration
@@ -213,14 +233,15 @@ private CognitoCredentialsProvider(CognitoCredentialsProviderBuilder builder) {
                 proxyTlsContextHandle,
                 proxyAuthorizationType,
                 proxyAuthorizationUsername != null ? proxyAuthorizationUsername.getBytes(UTF8) : null,
-                proxyAuthorizationPassword != null ? proxyAuthorizationPassword.getBytes(UTF8) : null);
+                proxyAuthorizationPassword != null ? proxyAuthorizationPassword.getBytes(UTF8) : null,
+                builder.loginTokenSource);
 
         acquireNativeHandle(nativeHandle);
         addReferenceTo(clientBootstrap);
         addReferenceTo(tlsContext);
     }
 
-    private void writeLengthPrefixedBytesSafe(ByteBuffer buffer, byte[] bytes) {
+    private static void writeLengthPrefixedBytesSafe(ByteBuffer buffer, byte[] bytes) {
         if (bytes != null) {
             buffer.putInt(bytes.length);
             buffer.put(bytes);
@@ -229,7 +250,7 @@ private void writeLengthPrefixedBytesSafe(ByteBuffer buffer, byte[] bytes) {
         }
     }
 
-    private byte[] marshalLoginsForJni(ArrayList<CognitoLoginTokenPair> logins) {
+    private static byte[] marshalLoginsForJni(List<CognitoLoginTokenPair> logins) {
         int size = 0;
 
         for (CognitoLoginTokenPair login : logins) {
@@ -256,6 +277,16 @@ private byte[] marshalLoginsForJni(ArrayList<CognitoLoginTokenPair> logins) {
         return buffer.array();
     }
 
+    private static CompletableFuture<List<CognitoLoginTokenPair>> createChainedFuture(long invocationHandle, CompletableFuture<List<CognitoLoginTokenPair>> baseFuture) {
+        return baseFuture.whenComplete((token_pairs, ex) -> {
+            if (ex == null) {
+                completeLoginTokenFetch(invocationHandle, marshalLoginsForJni(token_pairs), null);
+            } else {
+                completeLoginTokenFetch(invocationHandle, null, ex);
+            }
+        });
+    }
+
     /*******************************************************************************
      * Native methods
      ******************************************************************************/
@@ -273,5 +304,8 @@ private static native long cognitoCredentialsProviderNew(CognitoCredentialsProvi
                                                           long proxyTlsContext,
                                                           int proxyAuthorizationType,
                                                           byte[] proxyAuthorizationUsername,
-                                                          byte[] proxyAuthorizationPassword);
+                                                          byte[] proxyAuthorizationPassword,
+                                                         CognitoLoginTokenSource loginTokenSource);
+
+    private static native void completeLoginTokenFetch(long invocationHandle, byte[] marshalledLogins, Throwable ex);
 }

src/main/java/software/amazon/awssdk/crt/auth/credentials/CognitoCredentialsProvider.java

@@ -0,0 +1,46 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package software.amazon.awssdk.crt.auth.credentials;
+
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Interface to allow for dynamic sourcing (i.e. per fetch-credentials request submitted to Cognito) of Cognito login
+ * token pairs.  It is *critical* to follow the guidance given in the documentation for `startLoginTokenFetch`
+ */
+public interface CognitoLoginTokenSource {
+
+    /**
+     * Method that a Cognito credentials provider will invoke before sending a fetch credentials
+     * request to Cognito.  The CognitoLoginTokenPairs that the future gets completed with are joined
+     * with the (static) CognitoLoginTokenPairs that were specified in the credential provider configuration
+     * on construction.  The merged set of CognitoLoginTokenPairs are added to the HTTP request sent
+     * to Cognito that sources credentials.
+     *
+     * You must follow several guidelines to properly use this feature; not following these guidelines can result
+     * in deadlocks, poor performance, or other undesirable behavior.
+     *
+     * 1. If you use this feature, you must complete the future or the underlying connection attempt will hang forever.
+     * Credentials sourcing is halted until the future gets completed.  If something goes wrong during
+     * login token sourcing, complete the future exceptionally.
+     *
+     * 2. You must not block or wait for asynchronous operations in this function.  This function is invoked from a CRT
+     * event loop thread, and the event loop is halted until this function is returned from.  If you need to perform
+     * an asynchronous or non-trivial operation in order to source the necessary login token pairs, then you must
+     * ensure that sourcing task executes on another thread.  The easiest way to do this would be to pass the future
+     * to a sourcing task that runs on an external executor.
+     *
+     * 3. No attempt is made to de-duplicate login keys.  If the final, unioned set of login token pairs contains
+     * multiple pairs with the same key, then which one of the duplicates gets used is not well-defined.  For correct
+     * behavior, you must ensure there can be no duplicates.
+     *
+     * @param tokenFuture future to complete with dynamically sourced login token pairs in order to continue the
+     *                    credentials fetching process
+     */
+    void startLoginTokenFetch(CompletableFuture<List<CognitoCredentialsProvider.CognitoLoginTokenPair>> tokenFuture);
+
+}

src/main/java/software/amazon/awssdk/crt/auth/credentials/CognitoLoginTokenSource.java

@@ -6,6 +6,7 @@
 
 import java.util.concurrent.CompletableFuture;
 import software.amazon.awssdk.crt.CrtResource;
+import software.amazon.awssdk.crt.CrtRuntimeException;
 import software.amazon.awssdk.crt.Log;
 
 /**
@@ -41,11 +42,13 @@ public CompletableFuture<Credentials> getCredentials() {
      * @param future the future that the credentials should be applied to
      * @param credentials the fetched credentials, if successful
      */
-    private void onGetCredentialsComplete(CompletableFuture<Credentials> future, Credentials credentials) {
+    private void onGetCredentialsComplete(CompletableFuture<Credentials> future, int errorCode, Credentials credentials) {
         if (credentials != null) {
             future.complete(credentials);
+        } else if (errorCode != 0) {
+            future.completeExceptionally(new CrtRuntimeException(errorCode));
         } else {
-            future.completeExceptionally(new RuntimeException("Failed to get a valid set of credentials"));
+            future.completeExceptionally(new CrtRuntimeException("Failed to get a valid set of credentials"));
         }
     }
 

src/main/java/software/amazon/awssdk/crt/auth/credentials/CredentialsProvider.java

@@ -172,6 +172,10 @@
   {
     "name": "java.util.concurrent.CompletableFuture",
     "methods": [
+      {
+        "name": "<init>",
+        "parameterTypes": []
+      },
       {
         "name": "complete",
         "parameterTypes": [
@@ -295,6 +299,29 @@
       }
     ]
   },
+  {
+    "name": "software.amazon.awssdk.crt.auth.credentials.CognitoCredentialsProvider",
+    "methods": [
+      {
+        "name": "createChainedFuture",
+        "parameterTypes": [
+          "long",
+          "java.util.concurrent.CompletableFuture"
+        ]
+      }
+    ]
+  },
+  {
+    "name": "software.amazon.awssdk.crt.auth.credentials.CognitoLoginTokenSource",
+    "methods": [
+      {
+        "name": "startLoginTokenFetch",
+        "parameterTypes": [
+          "java.util.concurrent.CompletableFuture"
+        ]
+      }
+    ]
+  },
   {
     "name": "software.amazon.awssdk.crt.auth.credentials.Credentials",
     "fields": [
@@ -325,6 +352,7 @@
         "name": "onGetCredentialsComplete",
         "parameterTypes": [
           "java.util.concurrent.CompletableFuture",
+          "int",
           "software.amazon.awssdk.crt.auth.credentials.Credentials"
         ]
       },