Add support for async-http-client

pimterry · pimterry · commit e06a0041a8b4 · 2021-02-24T17:35:33.000+01:00
diff --git a/README.md b/README.md
@@ -16,6 +16,7 @@ Traffic can be captured from at least:
 - [x] OkHttp v2, v3 & v4
 - [x] Retrofit
 - [x] Jetty-Client v9, 10 & 11
+- [x] Async-Http-Client
 
 This will also capture HTTP(S) from any downstream libraries based on each of these clients, and many other untested clients sharing similar implementations, and so should cover a very large percentage of HTTP client usage.
 
diff --git a/build.gradle b/build.gradle
@@ -29,6 +29,7 @@ dependencies {
     compileOnly group: 'org.apache.httpcomponents', name: 'httpclient', version: '4.5'
     compileOnly group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '5.0.3'
     compileOnly group: 'org.eclipse.jetty', name: 'jetty-client', version: '11.0.1'
+    compileOnly group: 'org.asynchttpclient', name: 'async-http-client', version: '2.12.2'
 
     // Test deps:
     testImplementation group: 'io.kotest', name: 'kotest-runner-junit5-jvm', version: '4.4.0'
diff --git a/src/main/java/tech/httptoolkit/javaagent/asynchttpclient/AsyncHttpClientReturnProxySelectorAdvice.java b/src/main/java/tech/httptoolkit/javaagent/asynchttpclient/AsyncHttpClientReturnProxySelectorAdvice.java
@@ -0,0 +1,22 @@
+package tech.httptoolkit.javaagent.asynchttpclient;
+
+import net.bytebuddy.asm.Advice;
+import org.asynchttpclient.proxy.ProxyServer;
+import org.asynchttpclient.proxy.ProxyServerSelector;
+import tech.httptoolkit.javaagent.HttpProxyAgent;
+
+import java.net.ProxySelector;
+import java.util.Optional;
+
+public class AsyncHttpClientReturnProxySelectorAdvice {
+
+    public static ProxyServerSelector proxyServerSelector = uri -> new ProxyServer.Builder(
+        HttpProxyAgent.getAgentProxyHost(),
+        HttpProxyAgent.getAgentProxyPort()
+    ).build();
+
+    @Advice.OnMethodExit
+    public static void getProxyServerSelector(@Advice.Return(readOnly = false) ProxyServerSelector returnValue) {
+        returnValue = proxyServerSelector;
+    }
+}
diff --git a/src/main/java/tech/httptoolkit/javaagent/asynchttpclient/AsyncHttpClientReturnSslContextAdvice.java b/src/main/java/tech/httptoolkit/javaagent/asynchttpclient/AsyncHttpClientReturnSslContextAdvice.java
@@ -0,0 +1,22 @@
+package tech.httptoolkit.javaagent.asynchttpclient;
+
+import io.netty.handler.ssl.SslContext;
+import io.netty.handler.ssl.SslContextBuilder;
+import net.bytebuddy.asm.Advice;
+import tech.httptoolkit.javaagent.HttpProxyAgent;
+
+import javax.net.ssl.SSLException;
+
+public class AsyncHttpClientReturnSslContextAdvice {
+    @Advice.OnMethodExit
+    public static void getSslContext(@Advice.Return(readOnly = false) SslContext returnValue) {
+        try {
+            returnValue = SslContextBuilder
+                    .forClient()
+                    .trustManager(HttpProxyAgent.getInterceptedTrustManagerFactory())
+                    .build();
+        } catch (SSLException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}
diff --git a/src/main/java/tech/httptoolkit/javaagent/asynchttpclient/AsyncHttpResetSslEngineFactoryAdvice.java b/src/main/java/tech/httptoolkit/javaagent/asynchttpclient/AsyncHttpResetSslEngineFactoryAdvice.java
@@ -0,0 +1,43 @@
+package tech.httptoolkit.javaagent.asynchttpclient;
+
+import net.bytebuddy.asm.Advice;
+import org.asynchttpclient.AsyncHttpClientConfig;
+import org.asynchttpclient.SslEngineFactory;
+
+import java.util.Collections;
+import java.util.Set;
+import java.util.WeakHashMap;
+
+public class AsyncHttpResetSslEngineFactoryAdvice {
+
+    // Track each ChannelManager with a weak ref, to avoid unnecessary reflection overhead by only
+    // initializing them once, instead of every request
+    public static Set<Object> patchedChannelManagers = Collections.newSetFromMap(new WeakHashMap());
+
+    @Advice.OnMethodEnter
+    public static void createSslHandler(
+        @Advice.This Object thisChannelManager
+    ) {
+        if (patchedChannelManagers.contains(thisChannelManager)) return;
+
+        try {
+            Class<?> ChannelManager = thisChannelManager.getClass();
+
+            SslEngineFactory sslEngineFactory = (SslEngineFactory) ChannelManager
+                    .getDeclaredField("sslEngineFactory")
+                    .get(thisChannelManager);
+
+            AsyncHttpClientConfig config = (AsyncHttpClientConfig) ChannelManager
+                    .getDeclaredField("config")
+                    .get(thisChannelManager);
+
+            // Reinitialize the SSL Engine from the config (which uses our new cert)
+            // before building the SSL handler.
+            sslEngineFactory.init(config);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+        patchedChannelManagers.add(thisChannelManager);
+    }
+}
diff --git a/src/main/kotlin/tech/httptoolkit/javaagent/AgentMain.kt b/src/main/kotlin/tech/httptoolkit/javaagent/AgentMain.kt
@@ -10,11 +10,15 @@ import java.lang.instrument.Instrumentation
 import javax.net.ssl.SSLContext
 import java.net.*
 import javax.net.ssl.HttpsURLConnection
+import javax.net.ssl.TrustManagerFactory
 
 
 lateinit var InterceptedSslContext: SSLContext
     private set
 
+lateinit var InterceptedTrustManagerFactory: TrustManagerFactory
+    private set
+
 lateinit var AgentProxyHost: String
     private set
 
@@ -48,7 +52,8 @@ fun agentmain(arguments: String?, instrumentation: Instrumentation) {
 fun interceptAllHttps(config: Config, instrumentation: Instrumentation) {
     val (certPath, proxyHost, proxyPort) = config
 
-    InterceptedSslContext = buildSslContextForCertificate(certPath)
+    InterceptedTrustManagerFactory = buildTrustManagerFactoryForCertificate(certPath)
+    InterceptedSslContext = buildSslContextForCertificate(InterceptedTrustManagerFactory)
     AgentProxyHost = proxyHost
     AgentProxyPort = proxyPort
 
@@ -74,7 +79,9 @@ fun interceptAllHttps(config: Config, instrumentation: Instrumentation) {
         ApacheClientRoutingV5Transformer(),
         ApacheSslSocketFactoryTransformer(),
         JavaClientTransformer(),
-        JettyClientTransformer()
+        JettyClientTransformer(),
+        AsyncHttpClientConfigTransformer(),
+        AsyncHttpChannelManagerTransformer()
     ).forEach { matchingAgentTransformer ->
         agentBuilder = matchingAgentTransformer.register(agentBuilder)
     }
diff --git a/src/main/kotlin/tech/httptoolkit/javaagent/AsyncHttpClientConfigTransformers.kt b/src/main/kotlin/tech/httptoolkit/javaagent/AsyncHttpClientConfigTransformers.kt
@@ -0,0 +1,59 @@
+package tech.httptoolkit.javaagent
+
+import net.bytebuddy.agent.builder.AgentBuilder
+import net.bytebuddy.asm.Advice
+import net.bytebuddy.description.type.TypeDescription
+import net.bytebuddy.dynamic.DynamicType
+import net.bytebuddy.utility.JavaModule
+import net.bytebuddy.matcher.ElementMatchers.*
+import tech.httptoolkit.javaagent.asynchttpclient.AsyncHttpClientReturnProxySelectorAdvice
+import tech.httptoolkit.javaagent.asynchttpclient.AsyncHttpClientReturnSslContextAdvice
+import tech.httptoolkit.javaagent.asynchttpclient.AsyncHttpResetSslEngineFactoryAdvice
+
+// For new clients, we just need to override the properties on the convenient config
+// class that contains both proxy & SSL configuration.
+class AsyncHttpClientConfigTransformer : MatchingAgentTransformer {
+    override fun register(builder: AgentBuilder): AgentBuilder {
+        return builder
+            .type(
+                hasSuperType(named("org.asynchttpclient.AsyncHttpClientConfig"))
+            ).and(
+                not(isInterface())
+            ).transform(this)
+    }
+
+    override fun transform(
+        builder: DynamicType.Builder<*>,
+        typeDescription: TypeDescription,
+        classLoader: ClassLoader?,
+        module: JavaModule?
+    ): DynamicType.Builder<*>? {
+        return builder
+            .visit(Advice.to(AsyncHttpClientReturnSslContextAdvice::class.java)
+                .on(hasMethodName("getSslContext")))
+            .visit(Advice.to(AsyncHttpClientReturnProxySelectorAdvice::class.java)
+                .on(hasMethodName("getProxyServerSelector")))
+    }
+}
+
+// For existing classes, we need to hook SSL Handler creation, called for
+// every new connection
+class AsyncHttpChannelManagerTransformer : MatchingAgentTransformer {
+    override fun register(builder: AgentBuilder): AgentBuilder {
+        return builder
+            .type(
+                named("org.asynchttpclient.netty.channel.ChannelManager")
+            ).transform(this)
+    }
+
+    override fun transform(
+        builder: DynamicType.Builder<*>,
+        typeDescription: TypeDescription,
+        classLoader: ClassLoader?,
+        module: JavaModule?
+    ): DynamicType.Builder<*>? {
+        return builder
+            .visit(Advice.to(AsyncHttpResetSslEngineFactoryAdvice::class.java)
+                .on(hasMethodName("createSslHandler")))
+    }
+}
diff --git a/src/main/kotlin/tech/httptoolkit/javaagent/CustomSslContext.kt b/src/main/kotlin/tech/httptoolkit/javaagent/CustomSslContext.kt
@@ -8,7 +8,7 @@ import java.security.cert.CertificateFactory
 import javax.net.ssl.SSLContext
 import javax.net.ssl.TrustManagerFactory
 
-fun buildSslContextForCertificate(certPath: String): SSLContext {
+fun buildTrustManagerFactoryForCertificate(certPath: String): TrustManagerFactory {
     val certFile = File(certPath)
     val certificate: Certificate = CertificateFactory.getInstance("X.509")
         .generateCertificate(FileInputStream(certFile))
@@ -19,9 +19,11 @@ fun buildSslContextForCertificate(certPath: String): SSLContext {
 
     val trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm())
     trustManagerFactory.init(keyStore)
+    return trustManagerFactory
+}
 
+fun buildSslContextForCertificate(trustManagerFactory: TrustManagerFactory): SSLContext {
     val sslContext = SSLContext.getInstance("TLS")
     sslContext.init(null, trustManagerFactory.trustManagers, null)
-
     return sslContext
 }
diff --git a/src/test/kotlin/IntegrationTests.kt b/src/test/kotlin/IntegrationTests.kt
@@ -74,8 +74,8 @@ class IntegrationTests : StringSpec({
         agentAttachProc.isAlive.shouldBe(false)
         agentAttachProc.exitValue().shouldBe(0)
 
-        // Target should pick up proxy details & quit happily too
-        targetProc.waitFor(10, TimeUnit.SECONDS)
+        // Target should pick up proxy details & quit happily, eventually
+        targetProc.waitFor(15, TimeUnit.SECONDS)
         targetProc.isAlive.shouldBe(false)
         targetProc.exitValue().shouldBe(0)
     }
diff --git a/test-app/build.gradle b/test-app/build.gradle
@@ -16,7 +16,8 @@ dependencies {
     implementation group: 'com.squareup.okhttp3', name: 'okhttp', version: '4.9.1'
     implementation group: 'com.squareup.okhttp', name: 'okhttp', version: '2.7.5'
     implementation group: 'com.squareup.retrofit2', name: 'retrofit', version: '2.9.0'
-    implementation group: 'org.eclipse.jetty', name: 'jetty-client', version: '11.0.1'
+    implementation group: 'org.eclipse.jetty', name: 'jetty-client', version: '10.0.0'
+    implementation group: 'org.asynchttpclient', name: 'async-http-client', version: '2.12.2'
 }
 
 test {
diff --git a/test-app/src/main/java/tech/httptoolkit/testapp/Main.java b/test-app/src/main/java/tech/httptoolkit/testapp/Main.java
@@ -18,7 +18,8 @@ public class Main {
             "okhttp-v2", new OkHttpV2Case(),
             "okhttp-v4", new OkHttpV4Case(),
             "retrofit", new RetrofitCase(),
-            "jetty-client", new JettyClientCase()
+            "jetty-client", new JettyClientCase(),
+            "async-http-client", new AsyncHttpClientCase()
     );
 
     public static void main(String[] args) throws Exception {
diff --git a/test-app/src/main/java/tech/httptoolkit/testapp/cases/AsyncHttpClientCase.java b/test-app/src/main/java/tech/httptoolkit/testapp/cases/AsyncHttpClientCase.java
@@ -0,0 +1,25 @@
+package tech.httptoolkit.testapp.cases;
+
+import org.asynchttpclient.AsyncHttpClient;
+import org.asynchttpclient.Request;
+import org.asynchttpclient.Response;
+
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+
+import static org.asynchttpclient.Dsl.*;
+
+public class AsyncHttpClientCase extends ClientCase<AsyncHttpClient> {
+
+    @Override
+    public AsyncHttpClient newClient(String url) {
+        return asyncHttpClient();
+    }
+
+    @Override
+    public int test(String url, AsyncHttpClient client) throws ExecutionException, InterruptedException {
+        Request request = get(url).build();
+        Future<Response> whenResponse = client.executeRequest(request);
+        return whenResponse.get().getStatusCode();
+    }
+}
diff --git a/test-app/src/main/java/tech/httptoolkit/testapp/cases/OkHttpV4Case.java b/test-app/src/main/java/tech/httptoolkit/testapp/cases/OkHttpV4Case.java
@@ -20,6 +20,7 @@ public int test(String url, OkHttpClient client) throws IOException {
                 .build();
 
         Response response = client.newCall(request).execute();
+        response.body().close();
         return response.code();
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,8 @@ dependencies {`
`16`	`16`	`implementation group: 'com.squareup.okhttp3', name: 'okhttp', version: '4.9.1'`
`17`	`17`	`implementation group: 'com.squareup.okhttp', name: 'okhttp', version: '2.7.5'`
`18`	`18`	`implementation group: 'com.squareup.retrofit2', name: 'retrofit', version: '2.9.0'`
`19`		`- implementation group: 'org.eclipse.jetty', name: 'jetty-client', version: '11.0.1'`
	`19`	`+ implementation group: 'org.eclipse.jetty', name: 'jetty-client', version: '10.0.0'`
	`20`	`+ implementation group: 'org.asynchttpclient', name: 'async-http-client', version: '2.12.2'`
`20`	`21`	`}`
`21`	`22`
`22`	`23`	`test {`