From 20f504d51c25830fc37c201bae2c0c37df3fdd99 Mon Sep 17 00:00:00 2001
From: Yang Xia <55853655+xiazcy@users.noreply.github.com>
Date: Fri, 12 Jan 2024 09:47:47 -0800
Subject: [PATCH] Updated sigV4 signing library for gremlin and downgrade
 gremlin driver to Neptune supported version

---
 athena-neptune/pom.xml                         |  3 ++-
 .../connectors/neptune/NeptuneConnection.java  | 18 ++++++++++++++++--
 .../NeptuneGremlinConnection.java              | 18 ++++++++++++++++--
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/athena-neptune/pom.xml b/athena-neptune/pom.xml
index 4a0c17ec3c..1e8a0c9f94 100644
--- a/athena-neptune/pom.xml
+++ b/athena-neptune/pom.xml
@@ -9,7 +9,8 @@
     <artifactId>athena-neptune</artifactId>
     <version>2022.47.1</version>
     <properties>
-        <gremlinDriverVersion>3.7.1</gremlinDriverVersion>
+        <!-- make sure gremlin driver version stays within the Neptune supported range -->
+        <gremlinDriverVersion>3.6.5</gremlinDriverVersion>
         <neptune.sigv4.signer.version>2.4.0</neptune.sigv4.signer.version>
     </properties>
     <dependencies>
diff --git a/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneConnection.java b/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneConnection.java
index 5f6bf5e76f..ef1ba9ef02 100644
--- a/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneConnection.java
+++ b/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneConnection.java
@@ -19,9 +19,11 @@
  */
 package com.amazonaws.athena.connectors.neptune;
 
+import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
+import com.amazonaws.neptune.auth.NeptuneNettyHttpSigV4Signer;
+import com.amazonaws.neptune.auth.NeptuneSigV4SignerException;
 import org.apache.tinkerpop.gremlin.driver.Client;
 import org.apache.tinkerpop.gremlin.driver.Cluster;
-import org.apache.tinkerpop.gremlin.driver.SigV4WebSocketChannelizer;
 import org.apache.tinkerpop.gremlin.driver.remote.DriverRemoteConnection;
 import org.apache.tinkerpop.gremlin.process.traversal.AnonymousTraversalSource;
 import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
@@ -43,7 +45,19 @@ protected NeptuneConnection(String neptuneEndpoint, String neptunePort, boolean
                .enableSsl(true);
                
         if (enabledIAM) {
-            builder = builder.channelizer(SigV4WebSocketChannelizer.class);
+            builder.handshakeInterceptor(r ->
+                    {
+                        try {
+                            NeptuneNettyHttpSigV4Signer sigV4Signer =
+                                    new NeptuneNettyHttpSigV4Signer(region, new DefaultAWSCredentialsProviderChain());
+                            sigV4Signer.signRequest(r);
+                        }
+                        catch (NeptuneSigV4SignerException e) {
+                            throw new RuntimeException("Exception occurred while signing the request", e);
+                        }
+                        return r;
+                    }
+            );
         }
         
         cluster = builder.create();
diff --git a/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/propertygraph/NeptuneGremlinConnection.java b/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/propertygraph/NeptuneGremlinConnection.java
index 6ce130defd..734f8aac72 100644
--- a/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/propertygraph/NeptuneGremlinConnection.java
+++ b/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/propertygraph/NeptuneGremlinConnection.java
@@ -20,9 +20,11 @@
 package com.amazonaws.athena.connectors.neptune.propertygraph;
 
 import com.amazonaws.athena.connectors.neptune.NeptuneConnection;
+import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
+import com.amazonaws.neptune.auth.NeptuneNettyHttpSigV4Signer;
+import com.amazonaws.neptune.auth.NeptuneSigV4SignerException;
 import org.apache.tinkerpop.gremlin.driver.Client;
 import org.apache.tinkerpop.gremlin.driver.Cluster;
-import org.apache.tinkerpop.gremlin.driver.SigV4WebSocketChannelizer;
 import org.apache.tinkerpop.gremlin.driver.remote.DriverRemoteConnection;
 import org.apache.tinkerpop.gremlin.process.traversal.AnonymousTraversalSource;
 import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
@@ -40,7 +42,19 @@ public NeptuneGremlinConnection(String neptuneEndpoint, String neptunePort, bool
                .enableSsl(true);
                
         if (enabledIAM) {
-            builder = builder.channelizer(SigV4WebSocketChannelizer.class);
+            builder.handshakeInterceptor(r ->
+                    {
+                        try {
+                            NeptuneNettyHttpSigV4Signer sigV4Signer =
+                                    new NeptuneNettyHttpSigV4Signer(region, new DefaultAWSCredentialsProviderChain());
+                            sigV4Signer.signRequest(r);
+                        }
+                        catch (NeptuneSigV4SignerException e) {
+                            throw new RuntimeException("Exception occurred while signing the request", e);
+                        }
+                        return r;
+                    }
+            );
         }
         
         cluster = builder.create();