Inferred Proxy Improvements

components/context/src/main/java/datadog/context/InferredProxyContext.java

@@ -0,0 +1,54 @@
+package datadog.context;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+public class InferredProxyContext implements ImplicitContextKeyed {
+  public static final ContextKey<InferredProxyContext> CONTEXT_KEY =
+      ContextKey.named("inferred-proxy-key");
+  private final Map<String, String> inferredProxy;
+
+  // at most 6 x-dd-proxy http headers to be extracted and stored into the Context hashmap,
+  // following API Gateway RFC
+  private final int DEFAULT_CAPACITY = 6;
+
+  public static InferredProxyContext fromContext(Context context) {
+    return context.get(CONTEXT_KEY);
+  }
+
+  public InferredProxyContext(Map<String, String> contextInfo) {
+    this.inferredProxy =
+        (contextInfo == null || contextInfo.isEmpty())
+            ? new HashMap<>(DEFAULT_CAPACITY)
+            : new HashMap<>(contextInfo);
+  }
+
+  public InferredProxyContext() {
+    this.inferredProxy = new HashMap<>(DEFAULT_CAPACITY);
+  }
+
+  public Map<String, String> getInferredProxyContext() {
+    return Collections.unmodifiableMap(inferredProxy);
+  }
+
+  public void putInferredProxyInfo(String key, String value) {
+    inferredProxy.put(key, value);
+  }
+
+  public void removeInferredProxyInfo(String key) {
+    inferredProxy.remove(key);
+  }
+
+  /**
+   * Creates a new context with this value under its chosen key.
+   *
+   * @param context the context to copy the original values from.
+   * @return the new context with the implicitly keyed value.
+   * @see Context#with(ImplicitContextKeyed)
+   */
+  @Override
+  public Context storeInto(Context context) {
+    return context.with(CONTEXT_KEY, this);
+  }
+}

components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java

@@ -0,0 +1,66 @@
+package datadog.context.propagation;
+
+import datadog.context.Context;
+import datadog.context.InferredProxyContext;
+import java.util.function.BiConsumer;
+
+public class InferredProxyPropagator implements Propagator {
+  public static final String INFERRED_PROXY_KEY = "x-dd-proxy";
+  /**
+   * METHOD STUB: InferredProxy is currently not meant to be injected to downstream services Injects
+   * a context into a downstream service using the given carrier.
+   *
+   * @param context the context containing the values to be injected.
+   * @param carrier the instance that will receive the key/value pairs to propagate.
+   * @param setter the callback to set key/value pairs into the carrier.
+   */
+  @Override
+  public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {}
+
+  /**
+   * Extracts a context from un upstream service.
+   *
+   * @param context the base context to store the extracted values on top, use {@link
+   *     Context#root()} for a default base context.
+   * @param carrier the instance to fetch the propagated key/value pairs from.
+   * @param visitor the callback to walk over the carrier and extract its key/value pais.
+   * @return A context with the extracted values on top of the given base context.
+   */
+  @Override
+  public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
+    if (context == null || carrier == null || visitor == null) {
+      return context;
+    }
+    InferredProxyContextExtractor extractor = new InferredProxyContextExtractor();
+    visitor.forEachKeyValue(carrier, extractor);
+
+    InferredProxyContext extractedContext = extractor.extractedContext;
+    if (extractedContext == null) {
+      return context;
+    }
+    return extractedContext.storeInto(context);
+  }
+
+  public static class InferredProxyContextExtractor implements BiConsumer<String, String> {
+    private InferredProxyContext extractedContext;
+
+    InferredProxyContextExtractor() {}
+
+    /**
+     * Performs this operation on the given arguments.
+     *
+     * @param key the first input argument from an http header
+     * @param value the second input argument from an http header
+     */
+    @Override
+    public void accept(String key, String value) {
+      if (key == null || key.isEmpty() || !key.startsWith(INFERRED_PROXY_KEY)) {
+        return;
+      }
+      if (extractedContext == null) {
+        extractedContext = new InferredProxyContext();
+      }
+      extractedContext.putInferredProxyInfo(key, value);
+    }
+  }
+}