Add Okhttp body instrumentation and test server (#92)

* Add Okhttp body instrumentation and test server

Signed-off-by: Pavol Loffay <[email protected]>

* Fix muzzle

Signed-off-by: Pavol Loffay <[email protected]>

* remove junk

Signed-off-by: Pavol Loffay <[email protected]>

Author: pavolloffay

README.md

@@ -16,6 +16,7 @@ List of supported frameworks with additional capabilities:
 | [Servlet](https://javaee.github.io/javaee-spec/javadocs/javax/servlet/package-summary.html)            | 2.3+            |
 | [Spark Web Framework](https://github.com/perwendel/spark)                                              | 2.3+            |
 | [gRPC](https://github.com/grpc/grpc-java)                                                              | 1.5+            |
+| [OkHttp](https://github.com/square/okhttp/)                                                            | 3.0+            |
 
 
 ## Build

instrumentation/build.gradle.kts

@@ -9,7 +9,9 @@ subprojects {
         implementation("com.google.auto.service:auto-service:1.0-rc7")
         annotationProcessor("com.google.auto.service:auto-service:1.0-rc7")
 
+        implementation("io.opentelemetry:opentelemetry-api:0.9.1")
         implementation("io.opentelemetry.javaagent:opentelemetry-javaagent-tooling:0.9.0")
+        implementation("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api:0.9.0")
         implementation(project(":javaagent-core"))
         implementation(project(":blocking"))
     }
@@ -28,6 +30,8 @@ dependencies{
     implementation(project(":instrumentation:servlet:servlet-3.0"))
     implementation(project(":instrumentation:servlet:servlet-3.1"))
     implementation(project(":instrumentation:spark-web-framework-2.3"))
+    implementation(project(":instrumentation:grpc-1.5"))
+    implementation(project(":instrumentation:okhttp:okhttp-3.0"))
 }
 
 tasks {

instrumentation/okhttp/okhttp-3.0/build.gradle.kts

@@ -0,0 +1,34 @@
+plugins {
+    `java-library`
+    id("net.bytebuddy.byte-buddy-gradle-plugin") version "1.10.10"
+    id("io.opentelemetry.instrumentation.auto-instrumentation")
+    muzzle
+}
+
+muzzle {
+    pass {
+        group = "com.squareup.okhttp3"
+        module = "okhttp"
+        versions = "[3.0,)"
+        assertInverse = true
+    }
+}
+
+afterEvaluate{
+    byteBuddy {
+        transformation(closureOf<net.bytebuddy.build.gradle.Transformation> {
+            setTasks(kotlin.collections.setOf("compileJava", "compileScala", "compileKotlin"))
+            plugin = "io.opentelemetry.javaagent.tooling.muzzle.collector.MuzzleCodeGenerationPlugin"
+            setClassPath(project(":javaagent-tooling").configurations["instrumentationMuzzle"] + configurations.runtimeClasspath + sourceSets["main"].output)
+        })
+    }
+}
+
+dependencies {
+    api("io.opentelemetry.javaagent.instrumentation:opentelemetry-javaagent-okhttp-3.0:0.9.0")
+
+    compileOnly("com.squareup.okhttp3:okhttp:3.0.0")
+    implementation("net.bytebuddy:byte-buddy:1.10.10")
+
+    testImplementation(project(":testing-common"))
+}

instrumentation/okhttp/okhttp-3.0/src/main/java/io/opentelemetry/instrumentation/hypertrace/okhttp/v3_0/OkHttp3BodyInstrumentation.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.instrumentation.hypertrace.okhttp.v3_0;
+
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.tooling.Instrumenter;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import okhttp3.Interceptor;
+import okhttp3.OkHttpClient;
+
+@AutoService(Instrumenter.class)
+public class OkHttp3BodyInstrumentation extends Instrumenter.Default {
+
+  public OkHttp3BodyInstrumentation() {
+    super("okhttp", "okhttp-3");
+  }
+
+  @Override
+  public int getOrder() {
+    return 1;
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("okhttp3.OkHttpClient");
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isConstructor().and(takesArgument(0, named("okhttp3.OkHttpClient$Builder"))),
+        OkHttp3BodyInstrumentation.class.getName() + "$OkHttp3Advice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "org.hypertrace.agent.core.ContentTypeUtils",
+      "org.hypertrace.agent.core.HypertraceSemanticAttributes",
+      packageName + ".OkHttpTracingInterceptor",
+    };
+  }
+
+  public static class OkHttp3Advice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void addTracingInterceptor(@Advice.Argument(0) OkHttpClient.Builder builder) {
+      for (Interceptor interceptor : builder.interceptors()) {
+        if (interceptor instanceof OkHttpTracingInterceptor) {
+          return;
+        }
+      }
+      OkHttpTracingInterceptor interceptor = new OkHttpTracingInterceptor();
+      builder.addInterceptor(interceptor);
+    }
+  }
+}

instrumentation/okhttp/okhttp-3.0/src/main/java/io/opentelemetry/instrumentation/hypertrace/okhttp/v3_0/OkHttpTracingInterceptor.java

@@ -0,0 +1,103 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.instrumentation.hypertrace.okhttp.v3_0;
+
+import io.opentelemetry.common.AttributeKey;
+import io.opentelemetry.javaagent.instrumentation.okhttp.v3_0.OkHttpClientTracer;
+import io.opentelemetry.trace.Span;
+import java.io.IOException;
+import java.util.function.Function;
+import okhttp3.Headers;
+import okhttp3.Interceptor;
+import okhttp3.MediaType;
+import okhttp3.Request;
+import okhttp3.RequestBody;
+import okhttp3.Response;
+import okhttp3.ResponseBody;
+import okio.Buffer;
+import org.hypertrace.agent.core.ContentTypeUtils;
+import org.hypertrace.agent.core.HypertraceSemanticAttributes;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class OkHttpTracingInterceptor implements Interceptor {
+  private static final Logger log = LoggerFactory.getLogger(OkHttpTracingInterceptor.class);
+
+  private static final OkHttpClientTracer TRACER = OkHttpClientTracer.TRACER;
+
+  @Override
+  public Response intercept(Chain chain) throws IOException {
+    Span span = TRACER.getCurrentSpan();
+
+    Request request = chain.request();
+    captureHeaders(span, request.headers(), HypertraceSemanticAttributes::httpRequestHeader);
+    captureRequestBody(span, request.body());
+
+    Response response = chain.proceed(request);
+    captureHeaders(span, response.headers(), HypertraceSemanticAttributes::httpResponseHeader);
+    return captureResponseBody(span, response);
+  }
+
+  private static void captureRequestBody(Span span, RequestBody requestBody) {
+    if (requestBody == null) {
+      return;
+    }
+    MediaType mediaType = requestBody.contentType();
+    if (mediaType == null || !ContentTypeUtils.shouldCapture(mediaType.toString())) {
+      return;
+    }
+    try {
+      Buffer buffer = new Buffer();
+      requestBody.writeTo(buffer);
+      span.setAttribute(HypertraceSemanticAttributes.HTTP_REQUEST_BODY, buffer.readUtf8());
+    } catch (IOException e) {
+      log.error("Could not read request requestBody", e);
+    }
+  }
+
+  private static Response captureResponseBody(Span span, final Response response) {
+    if (response.body() == null) {
+      return response;
+    }
+    ResponseBody responseBody = response.body();
+    MediaType mediaType = responseBody.contentType();
+    if (mediaType == null || !ContentTypeUtils.shouldCapture(mediaType.toString())) {
+      return response;
+    }
+
+    try {
+      String body = responseBody.string();
+      span.setAttribute(HypertraceSemanticAttributes.HTTP_RESPONSE_BODY, body);
+      return response
+          .newBuilder()
+          .body(ResponseBody.create(responseBody.contentType(), body))
+          .build();
+    } catch (IOException e) {
+      log.error("Could not read response body", e);
+    }
+    return response;
+  }
+
+  private static void captureHeaders(
+      Span span, Headers headers, Function<String, AttributeKey<String>> headerNameProvider) {
+    for (String name : headers.names()) {
+      for (String value : headers.values(name)) {
+        span.setAttribute(headerNameProvider.apply(name), value);
+      }
+    }
+  }
+}