Support proxy endpoint and port

neptune-export/src/main/java/com/amazonaws/services/neptune/cli/CommonConnectionModule.java

@@ -15,12 +15,12 @@
 import com.amazonaws.services.neptune.AmazonNeptune;
 import com.amazonaws.services.neptune.cluster.ConnectionConfig;
 import com.amazonaws.services.neptune.cluster.NeptuneClusterMetadata;
+import com.amazonaws.services.neptune.cluster.ProxyConfig;
 import com.amazonaws.services.neptune.export.EndpointValidator;
 import com.github.rvesse.airline.annotations.Option;
 import com.github.rvesse.airline.annotations.restrictions.*;
 import org.apache.commons.lang.StringUtils;
 
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.function.Supplier;
@@ -55,19 +55,34 @@ public class CommonConnectionModule {
 
     @Option(name = {"--nlb-endpoint"}, description = "Network load balancer endpoint (optional: use only if connecting to an IAM DB enabled Neptune cluster through a network load balancer (NLB) – see https://github.com/aws-samples/aws-dbs-refarch-graph/tree/master/src/connecting-using-a-load-balancer#connecting-to-amazon-neptune-from-clients-outside-the-neptune-vpc-using-aws-network-load-balancer).")
     @Once
-    @MutuallyExclusiveWith(tag = "load-balancer")
+    @MutuallyExclusiveWith(tag = "proxy-endpoint")
     private String networkLoadBalancerEndpoint;
 
     @Option(name = {"--alb-endpoint"}, description = "Application load balancer endpoint (optional: use only if connecting to an IAM DB enabled Neptune cluster through an application load balancer (ALB) – see https://github.com/aws-samples/aws-dbs-refarch-graph/tree/master/src/connecting-using-a-load-balancer#connecting-to-amazon-neptune-from-clients-outside-the-neptune-vpc-using-aws-application-load-balancer).")
     @Once
-    @MutuallyExclusiveWith(tag = "load-balancer")
+    @MutuallyExclusiveWith(tag = "proxy-endpoint")
     private String applicationLoadBalancerEndpoint;
 
     @Option(name = {"--lb-port"}, description = "Load balancer port (optional, default 80).")
     @Port(acceptablePorts = {PortType.SYSTEM, PortType.USER})
     @Once
     private int loadBalancerPort = 80;
 
+    @Option(name = {"--proxy-endpoint"}, description = "Proxy endpoint (optional: use only if connecting to an IAM DB enabled Neptune cluster through a proxy such as a bastion host).")
+    @Once
+    @MutuallyExclusiveWith(tag = "proxy-endpoint")
+    private String proxyEndpoint;
+
+    @Option(name = {"--proxy-port"}, description = "Proxy port (optional, default 8182).")
+    @Port(acceptablePorts = {PortType.SYSTEM, PortType.USER})
+    @Once
+    private int proxyPort = 8182;
+
+    @Option(name = {"--proxy-remove-host-header"}, description = "Remove Host header after Sigv4 signing request to be forwarded via proxy.")
+    @Port(acceptablePorts = {PortType.SYSTEM, PortType.USER})
+    @Once
+    private boolean removeProxyHostHeader = false;
+
     private final Supplier<AmazonNeptune> amazonNeptuneClientSupplier;
 
     public CommonConnectionModule(Supplier<AmazonNeptune> amazonNeptuneClientSupplier) {
@@ -76,23 +91,32 @@ public CommonConnectionModule(Supplier<AmazonNeptune> amazonNeptuneClientSupplie
 
     public ConnectionConfig config() {
 
-        if (StringUtils.isNotEmpty(clusterId)){
+        if (StringUtils.isNotEmpty(clusterId)) {
             NeptuneClusterMetadata clusterMetadata = NeptuneClusterMetadata.createFromClusterId(clusterId, amazonNeptuneClientSupplier);
             endpoints.addAll(clusterMetadata.endpoints());
         }
 
-        if (endpoints.isEmpty()){
+        if (endpoints.isEmpty()) {
             throw new IllegalStateException("You must supply a cluster ID or one or more endpoints");
         }
 
+        ProxyConfig proxyConfig = null;
+
+        if (StringUtils.isNotEmpty(networkLoadBalancerEndpoint)) {
+            proxyConfig = new ProxyConfig(networkLoadBalancerEndpoint, loadBalancerPort, false);
+        } else if (StringUtils.isNotEmpty(applicationLoadBalancerEndpoint)) {
+            proxyConfig = new ProxyConfig(applicationLoadBalancerEndpoint, loadBalancerPort, true);
+        } else if (StringUtils.isNotEmpty(proxyEndpoint)) {
+            proxyConfig = new ProxyConfig(proxyEndpoint, proxyPort, removeProxyHostHeader);
+        }
+
         return new ConnectionConfig(
                 clusterId,
                 EndpointValidator.validate(endpoints),
                 port,
-                networkLoadBalancerEndpoint,
-                applicationLoadBalancerEndpoint,
-                loadBalancerPort,
                 useIamAuth,
-                !disableSsl);
+                !disableSsl,
+                proxyConfig
+        );
     }
 }

neptune-export/src/main/java/com/amazonaws/services/neptune/cluster/CloneCluster.java

@@ -80,11 +80,7 @@ public Cluster cloneCluster(ConnectionConfig connectionConfig, ConcurrencyConfig
                         targetClusterId,
                         targetClusterMetadata.endpoints(),
                         connectionConfig.port(),
-                        connectionConfig.nlbEndpoint(),
-                        connectionConfig.albEndpoint(),
-                        connectionConfig.lbPort(),
-                        targetClusterMetadata.isIAMDatabaseAuthenticationEnabled(),
-                        true
+                        targetClusterMetadata.isIAMDatabaseAuthenticationEnabled(), true, connectionConfig.proxyConfig()
                 ),
                 new ConcurrencyConfig(newConcurrency),
                 targetClusterMetadata,

neptune-export/src/main/java/com/amazonaws/services/neptune/cluster/ConnectionConfig.java

@@ -23,45 +23,35 @@ public class ConnectionConfig {
     private final String clusterId;
     private final Collection<String> neptuneEndpoints;
     private final int neptunePort;
-    private final String nlbEndpoint;
-    private final String albEndpoint;
-    private final int lbPort;
     private final boolean useIamAuth;
     private boolean useSsl;
+    private final ProxyConfig proxyConfig;
 
     public ConnectionConfig(String clusterId,
                             Collection<String> neptuneEndpoints,
                             int neptunePort,
-                            String nlbEndpoint,
-                            String albEndpoint,
-                            int lbPort,
-                            boolean useIamAuth,
-                            boolean useSsl) {
+                            boolean useIamAuth, boolean useSsl, ProxyConfig proxyConfig) {
         this.clusterId = clusterId;
         this.neptuneEndpoints = neptuneEndpoints;
         this.neptunePort = neptunePort;
-        this.nlbEndpoint = nlbEndpoint;
-        this.albEndpoint = albEndpoint;
-        this.lbPort = lbPort;
         this.useIamAuth = useIamAuth;
         this.useSsl = useSsl;
+        this.proxyConfig = proxyConfig;
     }
 
     public Collection<String> endpoints() {
         if (isDirectConnection()) {
             return neptuneEndpoints;
-        } else if (nlbEndpoint != null) {
-            return Collections.singletonList(nlbEndpoint);
         } else {
-            return Collections.singletonList(albEndpoint);
+            return Collections.singletonList(proxyConfig.endpoint());
         }
     }
 
     public int port() {
         if (isDirectConnection()) {
             return neptunePort;
         } else {
-            return lbPort;
+            return proxyConfig.port();
         }
     }
 
@@ -76,32 +66,21 @@ public boolean useSsl() {
     public HandshakeRequestConfig handshakeRequestConfig() {
         if (isDirectConnection()) {
             return new HandshakeRequestConfig(Collections.emptyList(), neptunePort, false);
-        } else if (nlbEndpoint != null) {
-            return new HandshakeRequestConfig(neptuneEndpoints, neptunePort, false);
         } else {
-            return new HandshakeRequestConfig(neptuneEndpoints, neptunePort, true);
+            return new HandshakeRequestConfig(neptuneEndpoints, neptunePort, proxyConfig.removeHostHeader());
         }
     }
 
     public boolean isDirectConnection() {
-        return nlbEndpoint == null && albEndpoint == null;
+        return proxyConfig == null;
     }
 
-    public String nlbEndpoint() {
-        return nlbEndpoint;
+    public ProxyConfig proxyConfig() {
+        return proxyConfig;
     }
 
-    public String albEndpoint() {
-        return albEndpoint;
-    }
-
-    public int lbPort() {
-        return lbPort;
-    }
-
-
     public String clusterId() {
-        return StringUtils.isNotEmpty(clusterId ) ?
+        return StringUtils.isNotEmpty(clusterId) ?
                 clusterId :
                 NeptuneClusterMetadata.clusterIdFromEndpoint(endpoints().iterator().next());
     }

neptune-export/src/main/java/com/amazonaws/services/neptune/cluster/ProxyConfig.java

@@ -0,0 +1,47 @@
+/*
+Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License").
+You may not use this file except in compliance with the License.
+A copy of the License is located at
+    http://www.apache.org/licenses/LICENSE-2.0
+or in the "license" file accompanying this file. This file is distributed
+on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+express or implied. See the License for the specific language governing
+permissions and limitations under the License.
+*/
+
+package com.amazonaws.services.neptune.cluster;
+
+public class ProxyConfig
+{
+    private final String endpoint;
+    private final int port;
+    private final boolean removeHostHeader;
+
+    public ProxyConfig(String endpoint, int port, boolean removeHostHeader) {
+        this.endpoint = endpoint;
+        this.port = port;
+        this.removeHostHeader = removeHostHeader;
+    }
+
+    public String endpoint() {
+        return endpoint;
+    }
+
+    public int port() {
+        return port;
+    }
+
+    public boolean removeHostHeader() {
+        return removeHostHeader;
+    }
+
+    @Override
+    public String toString() {
+        return "ProxyConfig{" +
+                "endpoint='" + endpoint + '\'' +
+                ", port=" + port +
+                ", removeHostHeader=" + removeHostHeader +
+                '}';
+    }
+}