Added exception catches for the OkHTTPClient header vulnerability

communication/src/main/java/datadog/communication/http/OkHttpUtils.java

@@ -153,9 +153,8 @@ private static OkHttpClient buildHttpClient(
               public Request authenticate(final Route route, final Response response) {
                 final String credential =
                     Credentials.basic(proxyUsername, proxyPassword == null ? "" : proxyPassword);
-                return response
-                    .request()
-                    .newBuilder()
+
+                return new SafeRequestBuilder.Builder(response.request().newBuilder())
                     .header("Proxy-Authorization", credential)
                     .build();
               }
@@ -175,8 +174,9 @@ public Request authenticate(final Route route, final Response response) {
   }
 
   public static Request.Builder prepareRequest(final HttpUrl url, Map<String, String> headers) {
-    final Request.Builder builder =
-        new Request.Builder()
+
+    final SafeRequestBuilder.Builder builder =
+        new SafeRequestBuilder.Builder()
             .url(url)
             .addHeader(DATADOG_META_LANG, "java")
             .addHeader(DATADOG_META_LANG_VERSION, JAVA_VERSION)
@@ -192,7 +192,7 @@ public static Request.Builder prepareRequest(final HttpUrl url, Map<String, Stri
       builder.addHeader(e.getKey(), e.getValue());
     }
 
-    return builder;
+    return builder.getBuilder();
   }
 
   public static Request.Builder prepareRequest(
@@ -206,7 +206,7 @@ public static Request.Builder prepareRequest(
     if (agentless && apiKey != null) {
       // we only add the api key header if we know we're doing agentless. No point in adding it to
       // other agent-based requests since we know the datadog-agent isn't going to make use of it.
-      builder.addHeader(DD_API_KEY, apiKey);
+      builder = new SafeRequestBuilder.Builder(builder).addHeader(DD_API_KEY, apiKey).getBuilder();
     }
 
     return builder;

communication/src/main/java/datadog/communication/http/SafeRequestBuilder.java

@@ -0,0 +1,148 @@
+package datadog.communication.http;
+
+import java.net.URL;
+import javax.annotation.Nullable;
+import okhttp3.CacheControl;
+import okhttp3.Headers;
+import okhttp3.HttpUrl;
+import okhttp3.Request;
+import okhttp3.RequestBody;
+
+/*
+The purpose of this class is to wrap okhttp3.Request.Builder.
+This adapter is used to resolve a vulnerability in Request.builder where
+the .header() and .addHeader() methods can cause header secrets to be
+printed/logged when invalid characters are parsed.
+ */
+public final class SafeRequestBuilder {
+  public static class Builder {
+    private final Request.Builder requestBuilder;
+
+    public Builder() {
+      this.requestBuilder = new Request.Builder();
+    }
+
+    // Constructs a SafeRequestBuilder from an existing request
+    public Builder(Request.Builder request) {
+      this.requestBuilder = request;
+    }
+
+    public SafeRequestBuilder.Builder url(HttpUrl url) {
+      this.requestBuilder.url(url);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder url(String url) {
+      this.requestBuilder.url(url);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder url(URL url) {
+      this.requestBuilder.url(url);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder header(String name, String value) {
+      // This try/catch block prevents header secrets from being outputted to
+      // console or logs.
+      try {
+        this.requestBuilder.header(name, value);
+        return this;
+      } catch (IllegalArgumentException e) {
+        throw new IllegalArgumentException(
+            new StringBuilder()
+                .append("InvalidArgumentException at header() for header: ")
+                .append(name)
+                .toString());
+      }
+    }
+
+    public SafeRequestBuilder.Builder addHeader(String name, String value) {
+      // This try/catch block prevents header secrets from being outputted to
+      // console or logs.
+      try {
+        this.requestBuilder.addHeader(name, value);
+        return this;
+      } catch (IllegalArgumentException e) {
+        throw new IllegalArgumentException(
+            new StringBuilder()
+                .append("InvalidArgumentException at addHeader() for header: ")
+                .append(name)
+                .toString());
+      }
+    }
+
+    public SafeRequestBuilder.Builder removeHeader(String name) {
+      this.requestBuilder.removeHeader(name);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder headers(Headers headers) {
+      this.requestBuilder.headers(headers);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder cacheControl(CacheControl cacheControl) {
+      this.requestBuilder.cacheControl(cacheControl);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder get() {
+      this.requestBuilder.get();
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder head() {
+      this.requestBuilder.head();
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder post(RequestBody body) {
+      this.requestBuilder.post(body);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder delete(@Nullable RequestBody body) {
+      this.requestBuilder.delete(body);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder delete() {
+      this.requestBuilder.delete();
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder put(RequestBody body) {
+      this.requestBuilder.put(body);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder patch(RequestBody body) {
+      this.requestBuilder.patch(body);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder method(String method, @Nullable RequestBody body) {
+      this.requestBuilder.method(method, body);
+      return this;
+    }
+
+    public SafeRequestBuilder.Builder tag(@Nullable Object tag) {
+      this.requestBuilder.tag(tag);
+      return this;
+    }
+
+    public <T> SafeRequestBuilder.Builder tag(Class<? super T> type, @Nullable T tag) {
+      this.requestBuilder.tag(type, tag);
+      return this;
+    }
+
+    public Request build() {
+      return this.requestBuilder.build();
+    }
+
+    public Request.Builder getBuilder() {
+      return this.requestBuilder;
+    }
+  }
+}

communication/src/test/groovy/datadog/communication/http/SafeRequestBuilderTest.groovy

@@ -0,0 +1,46 @@
+package datadog.communication.http
+
+import datadog.communication.http.SafeRequestBuilder
+import okhttp3.Headers
+import okhttp3.Request
+import okhttp3.RequestBody
+import org.junit.Test
+import org.junit.Assert
+
+class SafeRequestBuilderTest {
+  SafeRequestBuilder.Builder testBuilder = new SafeRequestBuilder.Builder()
+
+  @Test
+  void "test add header"(){
+    testBuilder.url("http:localhost").addHeader("test","test")
+    Assert.assertEquals(testBuilder.build().headers().get("test"),"test")
+  }
+  @Test(expected = IllegalArgumentException.class)
+  void "test adding bad header"(){
+    testBuilder.addHeader("\n\n","\n\n")
+  }
+  @Test (expected = IllegalArgumentException.class)
+  void "test adding bad header2"(){
+    testBuilder.url("localhost").header("\n\n\n","\n\n\n")
+  }
+  @Test
+  void "test building result"(){
+    Request testRequest = new SafeRequestBuilder.Builder().url("http://localhost")
+      .header("key","value").build()
+    Assert.assertEquals(Request.class,testRequest.getClass())
+  }
+  @Test
+  void "test getBuilder"(){
+    Request.Builder originalBuilder = new Request.Builder().url("http://localhost")
+      .addHeader("test","test")
+    testBuilder = new SafeRequestBuilder.Builder(originalBuilder)
+    Assert.assertEquals(originalBuilder,testBuilder.getBuilder())
+  }
+  @Test
+  void "test get method"(){
+    testBuilder = new SafeRequestBuilder.Builder().url("http://localhost")
+    testBuilder.get()
+    testBuilder.headers(new Headers())
+    Assert.assertEquals(testBuilder.getClass(),SafeRequestBuilder.Builder.class)
+  }
+}

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/uploader/BatchUploader.java

@@ -2,6 +2,7 @@
 
 import com.datadog.debugger.util.DebuggerMetrics;
 import datadog.common.container.ContainerInfo;
+import datadog.communication.http.SafeRequestBuilder;
 import datadog.trace.api.Config;
 import datadog.trace.relocate.api.RatelimitedLogger;
 import java.io.IOException;
@@ -151,7 +152,8 @@ private void makeUploadRequest(byte[] json, String tags) throws IOException {
     if (!tags.isEmpty()) {
       builder.addQueryParameter("ddtags", tags);
     }
-    Request.Builder requestBuilder = new Request.Builder().url(builder.build()).post(body);
+    SafeRequestBuilder.Builder requestBuilder =
+        new SafeRequestBuilder.Builder().url(builder.build()).post(body);
     if (apiKey != null) {
       if (apiKey.isEmpty()) {
         log.debug("API key is empty");

dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java

@@ -5,13 +5,13 @@
 
 import com.squareup.moshi.JsonAdapter;
 import com.squareup.moshi.Moshi;
+import datadog.communication.http.SafeRequestBuilder;
 import datadog.trace.api.DDId;
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.core.propagation.ExtractedContext;
 import okhttp3.MediaType;
 import okhttp3.OkHttpClient;
-import okhttp3.Request;
 import okhttp3.RequestBody;
 import okhttp3.Response;
 import org.slf4j.Logger;
@@ -65,7 +65,7 @@ public static AgentSpan.Context notifyStartInvocation(Object event) {
     try (Response response =
         HTTP_CLIENT
             .newCall(
-                new Request.Builder()
+                new SafeRequestBuilder.Builder()
                     .url(EXTENSION_BASE_URL + START_INVOCATION)
                     .addHeader(DATADOG_META_LANG, "java")
                     .post(body)
@@ -99,14 +99,15 @@ public static AgentSpan.Context notifyStartInvocation(Object event) {
   }
 
   public static boolean notifyEndInvocation(AgentSpan span, boolean isError) {
+
     if (null == span || null == span.getSamplingPriority()) {
       log.error(
           "could not notify the extension as the lambda span is null or no sampling priority has been found");
       return false;
     }
     RequestBody body = RequestBody.create(jsonMediaType, "{}");
-    Request.Builder builder =
-        new Request.Builder()
+    SafeRequestBuilder.Builder builder =
+        new SafeRequestBuilder.Builder()
             .url(EXTENSION_BASE_URL + END_INVOCATION)
             .addHeader(DATADOG_META_LANG, "java")
             .addHeader(DATADOG_TRACE_ID, span.getTraceId().toString())
@@ -117,6 +118,7 @@ public static boolean notifyEndInvocation(AgentSpan span, boolean isError) {
     if (isError) {
       builder.addHeader(DATADOG_INVOCATION_ERROR, "true");
     }
+
     try (Response response = HTTP_CLIENT.newCall(builder.build()).execute()) {
       if (response.isSuccessful()) {
         log.debug("notifyEndInvocation success");

telemetry/src/main/java/datadog/telemetry/RequestBuilder.java

@@ -3,6 +3,7 @@
 import com.squareup.moshi.JsonAdapter;
 import com.squareup.moshi.Moshi;
 import datadog.common.container.ContainerInfo;
+import datadog.communication.http.SafeRequestBuilder;
 import datadog.telemetry.api.ApiVersion;
 import datadog.telemetry.api.Application;
 import datadog.telemetry.api.Host;
@@ -105,7 +106,7 @@ public Request build(RequestType requestType, Payload payload) {
     String json = JSON_ADAPTER.toJson(telemetry);
     RequestBody body = RequestBody.create(JSON, json);
 
-    return new Request.Builder()
+    return new SafeRequestBuilder.Builder()
         .url(httpUrl)
         .addHeader("Content-Type", JSON.toString())
         .addHeader("DD-Telemetry-API-Version", API_VERSION.toString())