kenhuuu commented on code in PR #2601:
URL: https://github.com/apache/tinkerpop/pull/2601#discussion_r1600373395


##########
gremlin-driver/src/main/java/org/apache/tinkerpop/gremlin/driver/Auth.java:
##########
@@ -18,33 +18,258 @@
  */
 package org.apache.tinkerpop.gremlin.driver;
 
+import com.amazonaws.DefaultRequest;
+import com.amazonaws.SignableRequest;
+import com.amazonaws.auth.AWS4Signer;
+import com.amazonaws.auth.AWSCredentials;
+import com.amazonaws.auth.AWSCredentialsProvider;
+import com.amazonaws.auth.BasicSessionCredentials;
+import com.amazonaws.http.HttpMethodName;
+import com.amazonaws.util.SdkHttpUtils;
+import com.amazonaws.util.StringUtils;
+import io.netty.buffer.ByteBuf;
 import io.netty.handler.codec.http.FullHttpRequest;
 import io.netty.handler.codec.http.HttpHeaderNames;
+import io.netty.handler.codec.http.HttpHeaders;
+import org.apache.http.entity.StringEntity;
 
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UnsupportedEncodingException;
+import java.net.URI;
+import java.util.ArrayList;
 import java.util.Base64;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+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;
 
 public abstract class Auth implements RequestInterceptor {
 
+    private static final String NEPTUNE_SERVICE_NAME = "neptune-db";
+
     public static Auth basic(final String username, final String password) {
         return new Basic(username, password);
     }
 
+    public static Auth sigv4(final String regionName, final 
AWSCredentialsProvider awsCredentialsProvider) {
+        return new Sigv4(regionName, awsCredentialsProvider, 
NEPTUNE_SERVICE_NAME);
+    }
+
+    public static Auth sigv4(final String regionName, final 
AWSCredentialsProvider awsCredentialsProvider, final String serviceName) {
+        return new Sigv4(regionName, awsCredentialsProvider, serviceName);
+    }
+
     public static class Basic extends Auth {
 
         private final String username;
         private final String password;
 
-        private Basic(final String username, final String password ) {
+        private Basic(final String username, final String password) {
             this.username = username;
             this.password = password;
         }
 
         @Override
-        public FullHttpRequest apply(FullHttpRequest fullHttpRequest) {
+        public FullHttpRequest apply(final FullHttpRequest fullHttpRequest) {
             final String valueToEncode = username + ":" + password;
             fullHttpRequest.headers().add(HttpHeaderNames.AUTHORIZATION,
                     "Basic " + 
Base64.getEncoder().encodeToString(valueToEncode.getBytes()));
             return fullHttpRequest;
         }
     }
+
+    public static class Sigv4 extends Auth {
+        private final AWSCredentialsProvider awsCredentialsProvider;
+        private final AWS4Signer aws4Signer;
+
+        private Sigv4(final String regionName, final AWSCredentialsProvider 
awsCredentialsProvider, final String serviceName) {
+            this.awsCredentialsProvider = awsCredentialsProvider;
+
+            aws4Signer = new AWS4Signer();
+            aws4Signer.setRegionName(regionName);
+            aws4Signer.setServiceName(serviceName);
+        }
+
+        @Override
+        public FullHttpRequest apply(final FullHttpRequest fullHttpRequest) {
+            try {
+                // Convert Http request into an AWS SDK signable request
+                final SignableRequest<?> awsSignableRequest = 
toSignableRequest(fullHttpRequest);
+
+                // Sign the AWS SDK signable request (which internally adds 
some HTTP headers)
+                final AWSCredentials credentials = 
awsCredentialsProvider.getCredentials();
+                aws4Signer.sign(awsSignableRequest, credentials);
+
+                // extract session token if temporary credentials are provided
+                String sessionToken = "";
+                if ((credentials instanceof BasicSessionCredentials)) {
+                    sessionToken = ((BasicSessionCredentials) 
credentials).getSessionToken();
+                }
+
+                // todo: confirm is needed to replace header `Host` with `host`
+                fullHttpRequest.headers().remove(HttpHeaderNames.HOST);
+                fullHttpRequest.headers().add(HOST, 
awsSignableRequest.getHeaders().get(HOST));
+                fullHttpRequest.headers().add(X_AMZ_DATE, 
awsSignableRequest.getHeaders().get(X_AMZ_DATE));
+                fullHttpRequest.headers().add(AUTHORIZATION, 
awsSignableRequest.getHeaders().get(AUTHORIZATION));
+
+                if (!sessionToken.isEmpty()) {
+                    fullHttpRequest.headers().add(X_AMZ_SECURITY_TOKEN, 
sessionToken);
+                }
+            } catch (final Throwable t) {

Review Comment:
   We probably need slightly better handling of errors here rather than just 
RuntimeException.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscr...@tinkerpop.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org

Reply via email to