Add support for Ktor (testing the CIO engine specifically)

Author: pimterry

README.md

@@ -20,6 +20,7 @@ Traffic can be captured from at least:
 - [x] Async-Http-Client
 - [x] Reactor-Netty v0.9 & v1+
 - [x] Spring WebClient
+- [x] Ktor-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.
 

build.gradle

@@ -38,6 +38,8 @@ dependencies {
     compileOnly group: 'org.eclipse.jetty', name: 'jetty-client', version: '11.0.1'
     compileOnly group: 'org.asynchttpclient', name: 'async-http-client', version: '2.12.2'
     compileOnly group: 'io.projectreactor.netty', name: 'reactor-netty', version: '1.0.4'
+    compileOnly group: 'io.ktor', name: 'ktor-client-core', version: '1.5.2'
+    compileOnly group: 'io.ktor', name: 'ktor-client-cio', version: '1.5.2'
 
     // Test deps:
     testImplementation group: 'io.kotest', name: 'kotest-runner-junit5-jvm', version: '4.4.0'

src/main/java/tech/httptoolkit/javaagent/advice/ktor/KtorResetProxyFieldAdvice.java

@@ -0,0 +1,22 @@
+package tech.httptoolkit.javaagent.advice.ktor;
+
+import io.ktor.network.tls.TLSConfig;
+import net.bytebuddy.asm.Advice;
+import tech.httptoolkit.javaagent.HttpProxyAgent;
+
+import javax.net.ssl.X509TrustManager;
+import java.net.InetSocketAddress;
+import java.net.Proxy;
+
+public class KtorResetProxyFieldAdvice {
+
+    @Advice.OnMethodEnter
+    public static void beforeExecute(
+        @Advice.FieldValue(value = "proxy", readOnly = false) Proxy proxyField
+    ) throws Exception {
+        proxyField = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(
+            HttpProxyAgent.getAgentProxyHost(),
+            HttpProxyAgent.getAgentProxyPort()
+        ));
+    }
+}

src/main/java/tech/httptoolkit/javaagent/advice/ktor/KtorResetTlsClientTrustAdvice.java

@@ -0,0 +1,30 @@
+package tech.httptoolkit.javaagent.advice.ktor;
+
+import io.ktor.network.tls.TLSConfig;
+import net.bytebuddy.asm.Advice;
+import tech.httptoolkit.javaagent.HttpProxyAgent;
+
+import javax.net.ssl.X509TrustManager;
+
+public class KtorResetTlsClientTrustAdvice {
+
+    @Advice.OnMethodEnter
+    public static void beforeOpenTLSSession(
+        @Advice.Argument(value = 3, readOnly = false) TLSConfig tlsClientConfig
+    ) throws Exception {
+        // We're rewriting Kotlin bytecode from Java now, so some things get funky. Here it seems
+        // that coroutines result in a double call where the outer call has no args, so we need this:
+        if (tlsClientConfig == null) return;
+
+        // Clone the config, but replace the trust manager with one that trusts only our certificate:
+        tlsClientConfig = new TLSConfig(
+                tlsClientConfig.getRandom(),
+                tlsClientConfig.getCertificates(),
+                (X509TrustManager) HttpProxyAgent
+                        .getInterceptedTrustManagerFactory()
+                        .getTrustManagers()[0],
+                tlsClientConfig.getCipherSuites(),
+                tlsClientConfig.getServerName()
+        );
+    }
+}

src/main/kotlin/tech/httptoolkit/javaagent/AgentMain.kt

@@ -100,7 +100,10 @@ fun interceptAllHttps(config: Config, instrumentation: Instrumentation) {
         ReactorNettyClientConfigTransformer(logger),
         ReactorNettyProxyProviderTransformer(logger),
         ReactorNettyOverrideRequestAddressTransformer(logger),
-        ReactorNettyHttpClientSecureTransformer(logger)
+        ReactorNettyHttpClientSecureTransformer(logger),
+        KtorClientEngineConfigTransformer(logger),
+        KtorCioEngineTransformer(logger),
+        KtorClientTlsTransformer(logger)
     ).forEach { matchingAgentTransformer ->
         agentBuilder = matchingAgentTransformer.register(agentBuilder)
     }

src/main/kotlin/tech/httptoolkit/javaagent/KtorCioTransformers.kt

@@ -0,0 +1,76 @@
+package tech.httptoolkit.javaagent
+
+import net.bytebuddy.agent.builder.AgentBuilder
+import net.bytebuddy.asm.Advice
+import net.bytebuddy.description.method.MethodDescription
+import net.bytebuddy.dynamic.DynamicType
+import net.bytebuddy.matcher.ElementMatchers.*
+import tech.httptoolkit.javaagent.advice.ReturnProxyAdvice
+import tech.httptoolkit.javaagent.advice.ktor.KtorResetProxyFieldAdvice
+import tech.httptoolkit.javaagent.advice.ktor.KtorResetTlsClientTrustAdvice
+
+// To intercept HTTPS, we need to change the trust manager in TLSConfig instances. We don't want
+// to mess with server config though, so we clone the config argument and replace it with one that
+// uses our custom trust manager instead, every time a new TLS session is opened:
+class KtorClientTlsTransformer(logger: TransformationLogger) : MatchingAgentTransformer(logger) {
+    override fun register(builder: AgentBuilder): AgentBuilder {
+        return builder
+            .type(
+                named("io.ktor.network.tls.TLSClientSessionJvmKt")
+            ).transform(this)
+    }
+
+    override fun transform(builder: DynamicType.Builder<*>): DynamicType.Builder<*> {
+        return builder
+            .visit(Advice.to(KtorResetTlsClientTrustAdvice::class.java)
+                .on(hasMethodName<MethodDescription>("openTLSSession")
+                    .and(takesArgument(3, named("io.ktor.network.tls.TLSConfig")))))
+    }
+}
+
+// Proxy configuration for new clients is easy: we just hook getProxy() in the engine
+// configuration to return our proxy.
+class KtorClientEngineConfigTransformer(logger: TransformationLogger) : MatchingAgentTransformer(logger) {
+    override fun register(builder: AgentBuilder): AgentBuilder {
+        return builder
+            .type(
+                named("io.ktor.client.engine.HttpClientEngineConfig")
+            ).transform(this)
+    }
+
+    override fun transform(builder: DynamicType.Builder<*>): DynamicType.Builder<*> {
+        return builder
+            .visit(Advice.to(ReturnProxyAdvice::class.java)
+                .on(hasMethodName("getProxy")))
+        }
+}
+
+// Proxy configuration for existing clients is only mildly harder: we hook individual engines
+// elsewhere anyway, so it shouldn't matter much, but not CIO which is ktor specific. For CIO,
+// we just need one more hook that resets the proxy field to the value from config (hooked above)
+// before any requests are executed:
+class KtorCioEngineTransformer(logger: TransformationLogger) : MatchingAgentTransformer(logger) {
+    override fun register(builder: AgentBuilder): AgentBuilder {
+        return builder
+            .type(
+                named("io.ktor.client.engine.cio.CIOEngine")
+            ).transform(this)
+    }
+
+    override fun transform(builder: DynamicType.Builder<*>): DynamicType.Builder<*> {
+        return builder
+            .visit(Advice.to(KtorResetProxyFieldAdvice::class.java)
+                .on(hasMethodName("execute")))
+    }
+}
+
+/**
+ *
+ *
+ * Proxy settings live in HttpClientEngineConfig. Stored on CIOEngine at creation time,
+ * then used in execute() only. Reset then?
+ *
+ * Yes: change getProxy settings (changes all new engines immediately, though others shouldn't
+ * need it) and reset on execute()
+ *
+*/

test-app/build.gradle

@@ -1,6 +1,7 @@
 plugins {
     id 'java'
     id 'com.github.johnrengelman.shadow'
+    id 'org.jetbrains.kotlin.jvm'
 }
 
 group 'tech.httptoolkit'
@@ -24,6 +25,9 @@ dependencies {
     implementation group: 'org.asynchttpclient', name: 'async-http-client', version: '2.12.2'
     implementation group: 'org.springframework', name: 'spring-webflux', version: '5.3.4'
     implementation group: 'io.projectreactor.netty', name: 'reactor-netty', version: '1.0.4'
+
+    implementation group: 'io.ktor', name: 'ktor-client-core', version: '1.5.2'
+    implementation group: 'io.ktor', name: 'ktor-client-cio', version: '1.5.2'
 }
 
 test {

test-app/src/main/java/tech/httptoolkit/testapp/Main.java

@@ -27,7 +27,8 @@ public class Main {
         entry("retrofit", new RetrofitCase()),
         entry("jetty-client", new JettyClientCase()),
         entry("async-http-client", new AsyncHttpClientCase()),
-        entry("spring-web", new SpringWebClientCase())
+        entry("spring-web", new SpringWebClientCase()),
+        entry("ktor-cio", new KtorCioCase())
     );
 
     public static void main(String[] args) throws Exception {

test-app/src/main/java/tech/httptoolkit/testapp/cases/KtorCioCase.kt

@@ -0,0 +1,24 @@
+package tech.httptoolkit.testapp.cases
+
+import io.ktor.client.*
+import io.ktor.client.engine.cio.*
+import io.ktor.client.request.*
+import io.ktor.client.statement.*
+import kotlinx.coroutines.runBlocking
+
+class KtorCioCase : ClientCase<HttpClient>() {
+
+    override fun newClient(url: String): HttpClient {
+        return HttpClient(CIO)
+    }
+
+    override fun stopClient(client: HttpClient) {
+        client.close()
+    }
+
+    override fun test(url: String, client: HttpClient): Int = runBlocking {
+        val response = client.get<HttpResponse>(url)
+        response.status.value
+    }
+
+}