diff --git a/communication/src/main/java/datadog/communication/http/OkHttpUtils.java b/communication/src/main/java/datadog/communication/http/OkHttpUtils.java index ba57e004cd9..273cef30043 100644 --- a/communication/src/main/java/datadog/communication/http/OkHttpUtils.java +++ b/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(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 headers) { - final Request.Builder builder = - new Request.Builder() + + final SafeRequestBuilder builder = + new SafeRequestBuilder() .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 SafeRequestBuilder tag(Class 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; + } +} diff --git a/communication/src/test/groovy/datadog/communication/http/SafeRequestBuilderTest.groovy b/communication/src/test/groovy/datadog/communication/http/SafeRequestBuilderTest.groovy new file mode 100644 index 00000000000..b42b129ab3c --- /dev/null +++ b/communication/src/test/groovy/datadog/communication/http/SafeRequestBuilderTest.groovy @@ -0,0 +1,115 @@ +package datadog.communication.http + + +import okhttp3.Headers +import okhttp3.MediaType +import okhttp3.Request +import okhttp3.RequestBody +import org.junit.Test +import org.junit.Assert + +class SafeRequestBuilderTest { + SafeRequestBuilder testBuilder = new SafeRequestBuilder() + + @Test + void "test add header"(){ + testBuilder.url("http:localhost").addHeader("test","test") + Assert.assertEquals(testBuilder.build().headers().get("test"),"test") + } + @Test + void "test remove header"(){ + testBuilder.url("http:localhost").removeHeader("test") + Assert.assertEquals(testBuilder.build().headers().get("test"),null) + } + @Test + void "test static add header"(){ + Request.Builder builder = new Request.Builder().url("http://localhost") + builder = SafeRequestBuilder.addHeader(builder,"test","test") + Assert.assertEquals(builder.build().headers().get("test"),"test") + } + @Test (expected = IllegalArgumentException) + void "test bad static add header"(){ + Request.Builder builder = new Request.Builder().url("http://localhost") + builder = SafeRequestBuilder.addHeader(builder,"\n\n","\n\n") + } + @Test(expected = IllegalArgumentException) + void "test adding bad header"(){ + testBuilder.url("http:localhost").addHeader("\n\n","\n\n") + } + @Test (expected = IllegalArgumentException) + void "test adding bad header2"(){ + testBuilder.url("localhost").header("\u0019","\u0080") + } + @Test + void "test building result"(){ + Request testRequest = new SafeRequestBuilder().url("http://localhost") + .header("key","value").build() + Assert.assertEquals(Request,testRequest.getClass()) + } + @Test + void "test getBuilder"(){ + Request.Builder originalBuilder = new Request.Builder().url("http://localhost") + .addHeader("test","test") + testBuilder = new SafeRequestBuilder(originalBuilder) + Assert.assertEquals(originalBuilder,testBuilder.getBuilder()) + } + @Test + void "test get method"(){ + testBuilder = new SafeRequestBuilder().url("http://localhost") + testBuilder.get() + testBuilder.headers(new Headers()) + Assert.assertEquals(testBuilder.build().method(),"GET") + } + @Test + void "test post method"(){ + testBuilder = new SafeRequestBuilder().url("http://localhost") + RequestBody body = RequestBody.create(MediaType.parse("application/json"), "{}") + testBuilder.post(body) + Assert.assertEquals(testBuilder.build().method(),"POST") + } + @Test + void "test put method"(){ + testBuilder = new SafeRequestBuilder().url("http://localhost") + RequestBody body = RequestBody.create(MediaType.parse("application/json"), "{}") + testBuilder.put(body) + Assert.assertEquals(testBuilder.build().method(),"PUT") + } + @Test + void "test patch method"(){ + testBuilder = new SafeRequestBuilder().url("http://localhost") + RequestBody body = RequestBody.create(MediaType.parse("application/json"), "{}") + testBuilder.patch(body) + Assert.assertEquals(testBuilder.build().method(),"PATCH") + } + @Test + void "test delete method"(){ + testBuilder = new SafeRequestBuilder().url("http://localhost") + testBuilder.delete() + Assert.assertEquals(testBuilder.build().method(),"DELETE") + } + @Test + void "test method adder"(){ + testBuilder = new SafeRequestBuilder().url("http://localhost") + testBuilder.method("GET",null) + } + @Test + void "test tag"(){ + testBuilder = new SafeRequestBuilder() + URL url = new URL("http://localhost") + testBuilder.url(url).tag("test") + Assert.assertEquals(testBuilder.build().tag().toString(),"test") + } + @Test + void "test head"(){ + testBuilder = new SafeRequestBuilder().url("http://localhost") + testBuilder.head() + Assert.assertEquals(testBuilder.build().method(),"HEAD") + } + @Test + void "test delete with parameter method"(){ + testBuilder = new SafeRequestBuilder().url("http://localhost") + RequestBody body = RequestBody.create(MediaType.parse("application/json"), "{}") + testBuilder.delete(body) + Assert.assertEquals(testBuilder.build().method(),"DELETE") + } +} diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/uploader/BatchUploader.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/uploader/BatchUploader.java index 4bbf7c82fc9..07e819f847c 100644 --- a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/uploader/BatchUploader.java +++ b/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,7 @@ 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 requestBuilder = new SafeRequestBuilder().url(builder.build()).post(body); if (apiKey != null) { if (apiKey.isEmpty()) { log.debug("API key is empty"); diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java index 133a12a3689..0ea6757fa18 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java +++ b/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() .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 = + new SafeRequestBuilder() .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"); diff --git a/telemetry/src/main/java/datadog/telemetry/RequestBuilder.java b/telemetry/src/main/java/datadog/telemetry/RequestBuilder.java index 3ae2d701536..8ada6c91185 100644 --- a/telemetry/src/main/java/datadog/telemetry/RequestBuilder.java +++ b/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() .url(httpUrl) .addHeader("Content-Type", JSON.toString()) .addHeader("DD-Telemetry-API-Version", API_VERSION.toString())