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


##########
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) {
+                throw new RuntimeException(t);
+            }
+            return fullHttpRequest;
+        }
+
+        private SignableRequest<?> toSignableRequest(final FullHttpRequest 
request)
+                throws Exception {
+
+            // make sure the request is not null and contains the minimal 
required set of information
+            checkNotNull(request, "The request must not be null");
+            checkNotNull(request.uri(), "The request URI must not be null");
+            checkNotNull(request.method(), "The request method must not be 
null");
+
+            // convert the headers to the internal API format
+            final HttpHeaders headers = request.headers();
+            final Map<String, String> headersInternal = new HashMap<>();
+
+            String hostName = "";
+
+            // we don't want to add the Host header as the Signer always adds 
the host header.
+            for (String header : headers.names()) {
+                // Skip adding the Host header as the signing process will add 
one.
+                if (!header.equalsIgnoreCase(HOST)) {
+                    headersInternal.put(header, headers.get(header));
+                } else {
+                    hostName = headers.get(header);
+                }
+            }
+
+            // convert the parameters to the internal API format
+            final URI uri = URI.create(request.uri());
+
+            final String queryStr = uri.getQuery();
+            final Map<String, List<String>> parametersInternal = new 
HashMap<>(extractParametersFromQueryString(queryStr));
+
+            // carry over the entity (or an empty entity, if no entity is 
provided)
+            final InputStream content;
+            final ByteBuf contentBuffer = request.content();
+            boolean hasContent = false;
+            try {
+                if (contentBuffer != null && contentBuffer.isReadable()) {
+                    hasContent = true;
+                    contentBuffer.retain();
+                    byte[] bytes = new byte[contentBuffer.readableBytes()];
+                    contentBuffer.getBytes(contentBuffer.readerIndex(), bytes);
+                    content = new ByteArrayInputStream(bytes);
+                } else {
+                    content = new StringEntity("").getContent();
+                }
+            } catch (UnsupportedEncodingException e) {
+                throw new Exception("Encoding of the input string failed", e);
+            } catch (IOException e) {
+                throw new Exception("IOException while accessing entity 
content", e);
+            } finally {
+                if (hasContent) {
+                    contentBuffer.release();
+                }
+            }
+
+            if (StringUtils.isNullOrEmpty(hostName)) {
+                // try to extract hostname from the uri since hostname was not 
provided in the header.
+                final String authority = uri.getAuthority();
+                if (authority == null) {
+                    throw new Exception("Unable to identify host information,"
+                            + " either hostname should be provided in the uri 
or should be passed as a header");
+                }
+
+                hostName = authority;
+            }
+
+            final URI endpointUri = URI.create("http://"; + hostName);
+
+            return convertToSignableRequest(
+                    request.method().name(),
+                    endpointUri,
+                    uri.getPath(),
+                    headersInternal,
+                    parametersInternal,
+                    content);
+        }
+
+        private Map<String, List<String>> 
extractParametersFromQueryString(final String queryStr) {

Review Comment:
   I don't know if we need to support this right now.



-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to