Replace AuthTokens.basic with NeptuneAuthToken in OpenCypher #259

This change updates the implementation of InternalAuthToken to NeptuneAuthToken, taken from AWS Neptune documentation, so that the signature can be automatically renewed even if the connection is broken.Replace AuthTokens.basic with NeptuneAuthToken in OpenCypher

src/main/java/software/aws/neptune/opencypher/OpenCypherIAMRequestGenerator.java

@@ -24,16 +24,23 @@
 import com.amazonaws.http.HttpMethodName;
 import com.google.gson.Gson;
 import org.neo4j.driver.AuthToken;
-import org.neo4j.driver.AuthTokens;
+import org.neo4j.driver.Values;
+import org.neo4j.driver.internal.security.InternalAuthToken;
 
 import java.net.URI;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
 import static com.amazonaws.auth.internal.SignerConstants.AUTHORIZATION;
 import static com.amazonaws.auth.internal.SignerConstants.HOST;
 import static com.amazonaws.auth.internal.SignerConstants.X_AMZ_DATE;
 import static com.amazonaws.auth.internal.SignerConstants.X_AMZ_SECURITY_TOKEN;
+import lombok.Builder;
+import lombok.Getter;
+import lombok.NonNull;
+import org.neo4j.driver.Value;
+import org.neo4j.driver.internal.value.StringValue;
 
 /**
  * Class to help with IAM authentication.
@@ -53,17 +60,7 @@ public class OpenCypherIAMRequestGenerator {
      * @return AuthToken for IAM authentication.
      */
     public static AuthToken createAuthToken(final String url, final String region) {
-        final Request<Void> request = new DefaultRequest<>(SERVICE_NAME);
-        request.setHttpMethod(HttpMethodName.GET);
-        request.setEndpoint(URI.create(url));
-        request.setResourcePath("/opencypher");
-
-        final AWS4Signer signer = new AWS4Signer();
-        signer.setRegionName(region);
-        signer.setServiceName(request.getServiceName());
-        signer.sign(request, AWS_CREDENTIALS_PROVIDER.getCredentials());
-
-        return AuthTokens.basic(DUMMY_USERNAME, getAuthInfoJson(request));
+        return new NeptuneAuthToken(region, url, AWS_CREDENTIALS_PROVIDER);
     }
 
     private static String getAuthInfoJson(final Request<Void> request) {
@@ -76,4 +73,63 @@ private static String getAuthInfoJson(final Request<Void> request) {
 
         return GSON.toJson(obj);
     }
+
+    /**
+     * This class is derived from the AWS documentation on accessing Amazon Neptune with IAM authentication.
+     * For more details and information, please refer to:
+     * <a href="https://docs.aws.amazon.com/neptune/latest/userguide/access-graph-opencypher-bolt.html#access-graph-opencypher-bolt-java-iam-auth">Using the Bolt protocol to make openCypher queries to Neptune</a>.
+     */
+
+    public static class NeptuneAuthToken extends InternalAuthToken {
+        private static final String SCHEME = "basic";
+        private static final String REALM = "realm";
+        private static final String SERVICE_NAME = "neptune-db";
+        private static final String DUMMY_USERNAME = "username";
+        @NonNull
+        private final String region;
+        @NonNull
+        @Getter
+        private final String url;
+        @NonNull
+        private final AWSCredentialsProvider credentialsProvider;
+
+        @Builder
+        private NeptuneAuthToken(
+                @NonNull final String region,
+                @NonNull final String url,
+                @NonNull final AWSCredentialsProvider credentialsProvider
+        ) {
+            // The superclass caches the result of toMap(), which we don't want
+            super(Collections.emptyMap());
+            this.region = region;
+            this.url = url;
+            this.credentialsProvider = credentialsProvider;
+        }
+
+        @Override
+        public Map<String, Value> toMap() {
+            final Map<String, Value> map = new HashMap<>();
+            map.put(SCHEME_KEY, Values.value(SCHEME));
+            map.put(PRINCIPAL_KEY, Values.value(DUMMY_USERNAME));
+            map.put(CREDENTIALS_KEY, new StringValue(getSignedHeader()));
+            map.put(REALM_KEY, Values.value(REALM));
+
+            return map;
+        }
+
+        private String getSignedHeader() {
+            final Request<Void> request = new DefaultRequest<>(SERVICE_NAME);
+            request.setHttpMethod(HttpMethodName.GET);
+            request.setEndpoint(URI.create(url));
+            // Comment out the following line if you're using an engine version older than 1.2.0.0
+            request.setResourcePath("/opencypher");
+
+            final AWS4Signer signer = new AWS4Signer();
+            signer.setRegionName(region);
+            signer.setServiceName(request.getServiceName());
+            signer.sign(request, credentialsProvider.getCredentials());
+
+            return getAuthInfoJson(request);
+        }
+    }
 }

src/test/java/software/aws/neptune/opencypher/OpenCypherIAMRequestGeneratorTest.java

@@ -99,8 +99,7 @@ private void verifyAuthToken(final boolean useTempCreds) {
 
         assertEquals(value("basic"), internalAuthToken.get(SCHEME_KEY));
         assertEquals(value(DUMMY_USERNAME), internalAuthToken.get(PRINCIPAL_KEY));
-
-        assertFalse(internalAuthToken.containsKey(REALM_KEY));
+        assertTrue(internalAuthToken.containsKey(REALM_KEY));
 
         assertTrue(internalAuthToken.containsKey(CREDENTIALS_KEY));
         final Value credentialsValue = internalAuthToken.get(CREDENTIALS_KEY);