From e133c817c6bbec3e5f735c29a9b42a50d1f1a3f9 Mon Sep 17 00:00:00 2001
From: Devin Smith <devinsmith@deephaven.io>
Date: Mon, 6 Feb 2023 08:54:34 -0800
Subject: [PATCH] Add extraHeaders support for java-client (#3406)

This adds the ability for construction of a java-client with extra headers set on every request. This is potentially useful for cases where a proxy needs additional information to be able to route a request. It has been plumbed through the java client example utilities. In the future, we may wish to more formally support these use cases through DeephavenTarget / DH URIs.
---
 .../io/deephaven/client/examples/ConnectOptions.java |  8 ++++++++
 .../java/io/deephaven/client/impl/ChannelHelper.java | 12 ++++++++++++
 .../java/io/deephaven/client/impl/ClientConfig.java  | 11 ++++++++++-
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/java-client/example-utilities/src/main/java/io/deephaven/client/examples/ConnectOptions.java b/java-client/example-utilities/src/main/java/io/deephaven/client/examples/ConnectOptions.java
index cdeb943be7a..77a02c2384c 100644
--- a/java-client/example-utilities/src/main/java/io/deephaven/client/examples/ConnectOptions.java
+++ b/java-client/example-utilities/src/main/java/io/deephaven/client/examples/ConnectOptions.java
@@ -12,6 +12,8 @@
 import io.grpc.ManagedChannel;
 import picocli.CommandLine.Option;
 
+import java.util.Map;
+
 public class ConnectOptions {
 
     public static final String DEFAULT_HOST = "localhost";
@@ -46,6 +48,9 @@ public static ManagedChannel open(ConnectOptions options) {
     @Option(names = {"--ssl"}, description = "The optional ssl config file.", converter = SSLConverter.class)
     SSLConfig ssl;
 
+    @Option(names = {"--header"})
+    Map<String, String> headers;
+
     private ClientConfig config() {
         final Builder builder = ClientConfig.builder().target(target);
         if (userAgent != null) {
@@ -57,6 +62,9 @@ private ClientConfig config() {
         if (ssl != null) {
             builder.ssl(ssl);
         }
+        if (headers != null) {
+            builder.putAllExtraHeaders(headers);
+        }
         return builder.build();
     }
 
diff --git a/java-client/session/src/main/java/io/deephaven/client/impl/ChannelHelper.java b/java-client/session/src/main/java/io/deephaven/client/impl/ChannelHelper.java
index 25c2417e102..08910d3acc6 100644
--- a/java-client/session/src/main/java/io/deephaven/client/impl/ChannelHelper.java
+++ b/java-client/session/src/main/java/io/deephaven/client/impl/ChannelHelper.java
@@ -7,14 +7,17 @@
 import io.deephaven.ssl.config.TrustJdk;
 import io.deephaven.ssl.config.impl.KickstartUtils;
 import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
 import io.grpc.netty.GrpcSslContexts;
 import io.grpc.netty.NettyChannelBuilder;
+import io.grpc.stub.MetadataUtils;
 import io.netty.handler.ssl.SslContextBuilder;
 import io.netty.handler.ssl.SupportedCipherSuiteFilter;
 import nl.altindag.ssl.SSLFactory;
 import nl.altindag.ssl.util.NettySslUtils;
 
 import javax.net.ssl.SSLException;
+import java.util.Map;
 
 public final class ChannelHelper {
 
@@ -62,6 +65,15 @@ public static ManagedChannel channel(ClientConfig clientConfig) {
             channelBuilder.usePlaintext();
         }
         clientConfig.userAgent().ifPresent(channelBuilder::userAgent);
+        if (!clientConfig.extraHeaders().isEmpty()) {
+            channelBuilder.intercept(MetadataUtils.newAttachHeadersInterceptor(of(clientConfig.extraHeaders())));
+        }
         return channelBuilder.build();
     }
+
+    private static Metadata of(Map<String, String> map) {
+        final Metadata metadata = new Metadata();
+        map.forEach((k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v));
+        return metadata;
+    }
 }
diff --git a/java-client/session/src/main/java/io/deephaven/client/impl/ClientConfig.java b/java-client/session/src/main/java/io/deephaven/client/impl/ClientConfig.java
index 69abc34fd1c..66149c1c18d 100644
--- a/java-client/session/src/main/java/io/deephaven/client/impl/ClientConfig.java
+++ b/java-client/session/src/main/java/io/deephaven/client/impl/ClientConfig.java
@@ -6,10 +6,10 @@
 import io.deephaven.annotations.BuildableStyle;
 import io.deephaven.ssl.config.SSLConfig;
 import io.deephaven.uri.DeephavenTarget;
-import org.immutables.value.Value.Check;
 import org.immutables.value.Value.Default;
 import org.immutables.value.Value.Immutable;
 
+import java.util.Map;
 import java.util.Optional;
 
 /**
@@ -40,6 +40,11 @@ public static Builder builder() {
      */
     public abstract Optional<String> userAgent();
 
+    /**
+     * The extra headers.
+     */
+    public abstract Map<String, String> extraHeaders();
+
     /**
      * The maximum inbound message size. Defaults to 100MiB.
      */
@@ -56,6 +61,10 @@ public interface Builder {
 
         Builder userAgent(String userAgent);
 
+        Builder putExtraHeaders(String key, String value);
+
+        Builder putAllExtraHeaders(Map<String, ? extends String> entries);
+
         Builder maxInboundMessageSize(int maxInboundMessageSize);
 
         ClientConfig build();