feat: add context support for case insensitive and multivalue (#39)

* feat: add context support for case insensitive and multivalue

* refactor: update usages from deprecated versions

Author: aaron-steinfeld

grpc-client-rx-utils/build.gradle.kts

@@ -6,7 +6,7 @@ plugins {
 }
 
 dependencies {
-  api(platform("io.grpc:grpc-bom:1.45.1"))
+  api(platform("io.grpc:grpc-bom:1.50.0"))
   api("io.reactivex.rxjava3:rxjava:3.1.4")
   api("io.grpc:grpc-stub")
   api(project(":grpc-context-utils"))

grpc-client-utils/build.gradle.kts

@@ -7,7 +7,7 @@ plugins {
 
 dependencies {
 
-  api(platform("io.grpc:grpc-bom:1.47.0"))
+  api(platform("io.grpc:grpc-bom:1.50.0"))
   api("io.grpc:grpc-context")
   api("io.grpc:grpc-api")
   api(platform("io.netty:netty-bom:4.1.79.Final")) {

grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/RequestContextAsCreds.java

@@ -6,8 +6,8 @@
 import io.grpc.CallCredentials;
 import io.grpc.Metadata;
 import io.grpc.Status;
-import java.util.Map;
 import org.hypertrace.core.grpcutils.context.RequestContext;
+import org.hypertrace.core.grpcutils.context.RequestContext.RequestContextHeader;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -28,17 +28,17 @@ public void thisUsesUnstableApi() {}
   protected void applyRequestContext(MetadataApplier applier, RequestContext requestContext) {
     Metadata metadata = new Metadata();
     if (requestContext != null) {
-      for (Map.Entry<String, String> entry : requestContext.getAll().entrySet()) {
+      for (RequestContextHeader header : requestContext.getAllHeaders()) {
         // Exclude null headers
-        if (entry.getValue() != null) {
-          String key = entry.getKey();
+        if (header.getValue() != null) {
+          String key = header.getName();
           if (key.toLowerCase().endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
             metadata.put(
-                Metadata.Key.of(entry.getKey(), BINARY_BYTE_MARSHALLER),
-                entry.getValue().getBytes());
+                Metadata.Key.of(header.getName(), BINARY_BYTE_MARSHALLER),
+                header.getValue().getBytes());
           } else {
             metadata.put(
-                Metadata.Key.of(entry.getKey(), ASCII_STRING_MARSHALLER), entry.getValue());
+                Metadata.Key.of(header.getName(), ASCII_STRING_MARSHALLER), header.getValue());
           }
         }
       }

grpc-context-utils/build.gradle.kts

@@ -10,14 +10,17 @@ tasks.test {
 }
 
 dependencies {
-  api(platform("io.grpc:grpc-bom:1.45.1"))
+  api(platform("io.grpc:grpc-bom:1.50.0"))
   implementation("io.grpc:grpc-core")
 
   implementation("com.auth0:java-jwt:3.19.1")
   implementation("com.auth0:jwks-rsa:0.21.1")
   implementation("com.google.guava:guava:31.1-jre")
   implementation("org.slf4j:slf4j-api:1.7.36")
 
+  annotationProcessor("org.projectlombok:lombok:1.18.24")
+  compileOnly("org.projectlombok:lombok:1.18.24")
+
   constraints {
     implementation("com.fasterxml.jackson.core:jackson-databind:2.13.4.2") {
       because("https://nvd.nist.gov/vuln/detail/CVE-2022-42003")

grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/DefaultContextualKey.java

@@ -1,25 +1,22 @@
 package org.hypertrace.core.grpcutils.context;
 
+import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.Multimap;
 import java.util.Collection;
-import java.util.Map;
-import java.util.Map.Entry;
 import java.util.Objects;
-import java.util.Set;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
-import java.util.stream.Collectors;
 
 class DefaultContextualKey<T> implements ContextualKey<T> {
   private final RequestContext context;
   private final T data;
-  private final Map<String, String> cacheableContextHeaders;
+  private final Multimap<String, String> cacheableContextHeaders;
 
   DefaultContextualKey(RequestContext context, T data, Collection<String> cacheableHeaderNames) {
     this.context = context;
     this.data = data;
-    this.cacheableContextHeaders =
-        this.extractCacheableHeaders(context.getRequestHeaders(), cacheableHeaderNames);
+    this.cacheableContextHeaders = this.extractCacheableHeaders(context, cacheableHeaderNames);
   }
 
   @Override
@@ -76,14 +73,12 @@ public String toString() {
         + '}';
   }
 
-  private Map<String, String> extractCacheableHeaders(
-      Map<String, String> allHeaders, Collection<String> cacheableHeaderNames) {
-    Set<String> cacheableHeaderNameSet =
-        cacheableHeaderNames.stream()
-            .map(String::toLowerCase)
-            .collect(Collectors.toUnmodifiableSet());
-    return allHeaders.entrySet().stream()
-        .filter(entry -> cacheableHeaderNameSet.contains(entry.getKey().toLowerCase()))
-        .collect(Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue));
+  private Multimap<String, String> extractCacheableHeaders(
+      RequestContext requestContext, Collection<String> cacheableHeaderNames) {
+    return cacheableHeaderNames.stream()
+        .collect(
+            ImmutableListMultimap.flatteningToImmutableListMultimap(
+                Function.identity(),
+                headerName -> requestContext.getAllHeaderValues(headerName).stream()));
   }
 }

grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java

@@ -1,22 +1,28 @@
 package org.hypertrace.core.grpcutils.context;
 
 import static io.grpc.Metadata.ASCII_STRING_MARSHALLER;
+import static java.util.Objects.requireNonNull;
 import static org.hypertrace.core.grpcutils.context.RequestContextConstants.CACHE_MEANINGFUL_HEADERS;
 import static org.hypertrace.core.grpcutils.context.RequestContextConstants.TENANT_ID_HEADER_KEY;
 
-import com.google.common.collect.Maps;
+import com.google.common.collect.ListMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.MultimapBuilder;
+import com.google.common.collect.Multimaps;
 import io.grpc.Context;
 import io.grpc.Metadata;
 import io.grpc.Metadata.Key;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.Callable;
+import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import lombok.Value;
 
 /**
  * Context of the GRPC request that should be carried and can made available to the services so that
@@ -26,10 +32,9 @@ public class RequestContext {
   public static final Context.Key<RequestContext> CURRENT = Context.key("request_context");
 
   public static RequestContext forTenantId(String tenantId) {
-    RequestContext requestContext = new RequestContext();
-    requestContext.add(RequestContextConstants.TENANT_ID_HEADER_KEY, tenantId);
-    requestContext.add(RequestContextConstants.REQUEST_ID_HEADER_KEY, UUID.randomUUID().toString());
-    return requestContext;
+    return new RequestContext()
+        .put(RequestContextConstants.TENANT_ID_HEADER_KEY, tenantId)
+        .put(RequestContextConstants.REQUEST_ID_HEADER_KEY, UUID.randomUUID().toString());
   }
 
   public static RequestContext fromMetadata(Metadata metadata) {
@@ -53,19 +58,20 @@ public static RequestContext fromMetadata(Metadata metadata) {
               }
               // The value could be null or empty for some keys so validate that.
               if (value != null && !value.isEmpty()) {
-                requestContext.add(k, value);
+                requestContext.put(k, value);
               }
             });
 
     return requestContext;
   }
 
-  private final Map<String, String> headers = new HashMap<>();
+  private final ListMultimap<String, RequestContextHeader> headers =
+      MultimapBuilder.linkedHashKeys().linkedListValues().build();
   private final JwtParser jwtParser = new JwtParser();
 
   /** Reads tenant id from this RequestContext based on the tenant id http header and returns it. */
   public Optional<String> getTenantId() {
-    return get(RequestContextConstants.TENANT_ID_HEADER_KEY);
+    return getHeaderValue(RequestContextConstants.TENANT_ID_HEADER_KEY);
   }
 
   public Optional<String> getUserId() {
@@ -96,28 +102,110 @@ public Optional<JwtClaim> getClaim(String claimName) {
   }
 
   public Optional<String> getRequestId() {
-    return this.get(RequestContextConstants.REQUEST_ID_HEADER_KEY);
+    return this.getHeaderValue(RequestContextConstants.REQUEST_ID_HEADER_KEY);
   }
 
   private Optional<Jwt> getJwt() {
-    return get(RequestContextConstants.AUTHORIZATION_HEADER).flatMap(jwtParser::fromAuthHeader);
+    return this.getHeaderValue(RequestContextConstants.AUTHORIZATION_HEADER)
+        .flatMap(jwtParser::fromAuthHeader);
   }
 
-  /** Method to read all GRPC request headers from this RequestContext. */
+  /**
+   * This is retained for backwards compatibility, but is based on the incorrect assumption that a
+   * header only can have one value. For the updated API, please use {@link #getAllHeaders()}}
+   */
+  @Deprecated
   public Map<String, String> getRequestHeaders() {
     return getAll();
   }
 
-  public void add(String headerKey, String headerValue) {
-    this.headers.put(headerKey, headerValue);
+  /**
+   * This is retained for backwards compatibility. It previously was implemented with the assumption
+   * of a single value per header, so overwrote on addition. This is maintained by first removing
+   * any values for the given header name, then adding
+   *
+   * @param headerName
+   * @param headerValue
+   */
+  @Deprecated
+  public void add(String headerName, String headerValue) {
+    this.removeHeader(headerName);
+    this.put(headerName, headerValue);
   }
 
-  public Optional<String> get(String headerKey) {
-    return Optional.ofNullable(this.headers.get(headerKey));
+  /** Prefer {@link #getHeaderValue} */
+  @Deprecated
+  public Optional<String> get(String headerName) {
+    return this.getHeaderValue(headerName);
   }
 
+  /**
+   * This is retained for backwards compatibility, but is based on the incorrect assumption that a
+   * header only can have one value. For the updated API, please use {@link #getAllHeaders()} ()}
+   */
+  @Deprecated
   public Map<String, String> getAll() {
-    return Map.copyOf(headers);
+    return this.getHeaderNames().stream()
+        .flatMap(key -> this.headers.get(key).stream().findFirst().stream())
+        .collect(
+            Collectors.toUnmodifiableMap(
+                RequestContextHeader::getName, RequestContextHeader::getValue));
+  }
+
+  /**
+   * Adds the provided header name and header value. Duplicates are allowed. For fluency, the
+   * current instance is returned.
+   */
+  public RequestContext put(String headerName, String headerValue) {
+    this.headers.put(
+        this.normalizeHeaderName(requireNonNull(headerName)),
+        new RequestContextHeader(headerName, headerValue));
+    return this;
+  }
+
+  /** Returns all header names in normalized form (case not preserved) */
+  @Nonnull
+  public Set<String> getHeaderNames() {
+    return Set.copyOf(this.headers.keySet());
+  }
+
+  /**
+   * Removes and returns all headers matching the provided name. Header names are case insensitive.
+   * Returns an empty list if no headers have been removed.
+   */
+  @Nonnull
+  public List<String> removeHeader(String name) {
+    return this.headers.removeAll(this.normalizeHeaderName(name)).stream()
+        .map(RequestContextHeader::getValue)
+        .collect(Collectors.toUnmodifiableList());
+  }
+
+  /** Returns all header values matching the provided header name, case insensitively. */
+  @Nonnull
+  public List<String> getAllHeaderValues(String key) {
+    return this.headers.get(this.normalizeHeaderName(key)).stream()
+        .map(RequestContextHeader::getValue)
+        .collect(Collectors.toUnmodifiableList());
+  }
+
+  /**
+   * Returns all header name-value pairs, with the original case preserved. Multiple headers with
+   * the same name, and potentially same value may be returned.
+   */
+  @Nonnull
+  public List<RequestContextHeader> getAllHeaders() {
+    return List.copyOf(this.headers.values());
+  }
+
+  /**
+   * Gets the first header value specified for the provided name (case insensitive), or empty if no
+   * match.
+   */
+  @Nonnull
+  public Optional<String> getHeaderValue(String key) {
+    return this.headers.get(this.normalizeHeaderName(key)).stream()
+        .map(RequestContextHeader::getValue)
+        .findFirst();
   }
 
   public <V> V call(@Nonnull Callable<V> callable) {
@@ -198,8 +286,12 @@ public Metadata buildTrailers() {
     return trailers;
   }
 
-  private Map<String, String> getHeadersOtherThanAuth() {
-    return Maps.filterKeys(
+  private String normalizeHeaderName(@Nonnull String headerName) {
+    return headerName.toLowerCase();
+  }
+
+  private Multimap<String, RequestContextHeader> getHeadersOtherThanAuth() {
+    return Multimaps.filterKeys(
         headers, key -> !key.equals(RequestContextConstants.AUTHORIZATION_HEADER));
   }
 
@@ -213,4 +305,10 @@ public String toString() {
         + getJwt().map(Jwt::toString).orElse(emptyValue)
         + '}';
   }
+
+  @Value
+  public static class RequestContextHeader {
+    String name;
+    String value;
+  }
 }