Add support for Akka-HTTP

Author: pimterry

README.md

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

@@ -40,6 +40,8 @@ dependencies {
     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'
+    compileOnly group: 'com.typesafe.akka', name: 'akka-http-core_2.13', version: '10.2.4'
+    compileOnly group: 'com.typesafe.akka', name: 'akka-actor_2.13', version: '2.6.13'
 
     // Test deps:
     testImplementation group: 'io.kotest', name: 'kotest-runner-junit5-jvm', version: '4.4.0'

src/main/java/tech/httptoolkit/javaagent/advice/akka/OverrideHttpSettingsAdvice.java

@@ -0,0 +1,28 @@
+package tech.httptoolkit.javaagent.advice.akka;
+
+import akka.http.scaladsl.ClientTransport;
+import akka.http.scaladsl.ConnectionContext;
+import akka.http.scaladsl.settings.ClientConnectionSettings;
+import net.bytebuddy.asm.Advice;
+import tech.httptoolkit.javaagent.HttpProxyAgent;
+
+import java.net.InetSocketAddress;
+
+public class OverrideHttpSettingsAdvice {
+    @Advice.OnMethodEnter
+    public static void beforeOutgoingConnection(
+        @Advice.Argument(value = 2, readOnly = false) ClientConnectionSettings clientSettings,
+        @Advice.Argument(value = 3, readOnly = false) ConnectionContext connectionContext
+    ) {
+        // Change all new outgoing connections to use the proxy:
+        clientSettings = clientSettings.withTransport(
+            ClientTransport.httpsProxy(new InetSocketAddress(
+                    HttpProxyAgent.getAgentProxyHost(),
+                    HttpProxyAgent.getAgentProxyPort()
+            ))
+        );
+
+        // Change all new outgoing connections to trust our certificate:
+        connectionContext = ConnectionContext.httpsClient(HttpProxyAgent.getInterceptedSslContext());
+    }
+}

src/main/java/tech/httptoolkit/javaagent/advice/akka/ResetAllConnectionPoolsAdvice.java

@@ -0,0 +1,26 @@
+package tech.httptoolkit.javaagent.advice.akka;
+
+import net.bytebuddy.asm.Advice;
+
+import java.lang.reflect.Method;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public class ResetAllConnectionPoolsAdvice {
+
+    public static AtomicBoolean connectionPoolsReset = new AtomicBoolean(false);
+
+    @Advice.OnMethodEnter
+    public static void resetConnectionsBeforeDispatch(
+        @Advice.This Object thisPoolMaster
+    ) throws Exception {
+        boolean resetNeeded = ResetAllConnectionPoolsAdvice
+                .connectionPoolsReset
+                .compareAndSet(false, true);
+
+        if (!resetNeeded) return;
+
+        // Just once, when we first try to dispatch a request, reset all existing connection pools.
+        Method shutdownMethod = thisPoolMaster.getClass().getDeclaredMethod("shutdownAll");
+        shutdownMethod.invoke(thisPoolMaster);
+    }
+}

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

@@ -103,7 +103,9 @@ fun interceptAllHttps(config: Config, instrumentation: Instrumentation) {
         ReactorNettyHttpClientSecureTransformer(logger),
         KtorClientEngineConfigTransformer(logger),
         KtorCioEngineTransformer(logger),
-        KtorClientTlsTransformer(logger)
+        KtorClientTlsTransformer(logger),
+        AkkaHttpTransformer(logger),
+        AkkaPoolTransformer(logger)
     ).forEach { matchingAgentTransformer ->
         agentBuilder = matchingAgentTransformer.register(agentBuilder)
     }

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

@@ -0,0 +1,45 @@
+package tech.httptoolkit.javaagent
+
+import net.bytebuddy.agent.builder.AgentBuilder
+import net.bytebuddy.asm.Advice
+import net.bytebuddy.dynamic.DynamicType
+import net.bytebuddy.matcher.ElementMatchers.hasMethodName
+import net.bytebuddy.matcher.ElementMatchers.named
+import tech.httptoolkit.javaagent.advice.akka.OverrideHttpSettingsAdvice
+import tech.httptoolkit.javaagent.advice.akka.ResetAllConnectionPoolsAdvice
+
+// First, we hook outgoing connection creation, and ensure that new connections always go via the proxy & trust us:
+class AkkaHttpTransformer(logger: TransformationLogger) : MatchingAgentTransformer(logger) {
+    override fun register(builder: AgentBuilder): AgentBuilder {
+        return builder
+            .type(named("akka.http.scaladsl.HttpExt")) // Scala compiles Http()s methods here for some reason
+            .transform(this)
+    }
+
+    override fun transform(builder: DynamicType.Builder<*>): DynamicType.Builder<*> {
+        return builder
+            .visit(
+                Advice.to(OverrideHttpSettingsAdvice::class.java)
+                    .on(hasMethodName("_outgoingConnection"))
+            )
+    }
+}
+
+// Then, to ensure that any existing connections trust us too, we do a one-off connection pool reset,
+// triggered by the new PoolMaster.dispatchRequest (seems to happen for every request). This seems to
+// affects shared pools and individual pools too.
+class AkkaPoolTransformer(logger: TransformationLogger) : MatchingAgentTransformer(logger) {
+    override fun register(builder: AgentBuilder): AgentBuilder {
+        return builder
+            .type(named("akka.http.impl.engine.client.PoolMaster"))
+            .transform(this)
+    }
+
+    override fun transform(builder: DynamicType.Builder<*>): DynamicType.Builder<*> {
+        return builder
+            .visit(
+                Advice.to(ResetAllConnectionPoolsAdvice::class.java)
+                    .on(hasMethodName("dispatchRequest"))
+            )
+    }
+}

test-app/build.gradle

@@ -28,6 +28,10 @@ dependencies {
 
     implementation group: 'io.ktor', name: 'ktor-client-core', version: '1.5.2'
     implementation group: 'io.ktor', name: 'ktor-client-cio', version: '1.5.2'
+
+    implementation group: 'com.typesafe.akka', name: 'akka-actor_2.13', version: '2.6.13'
+    implementation group: 'com.typesafe.akka', name: 'akka-http_2.13', version: '10.2.4'
+    implementation group: 'com.typesafe.akka', name: 'akka-stream_2.13', version: '2.6.13'
 }
 
 test {
@@ -43,4 +47,13 @@ jar {
 compileJava {
     sourceCompatibility = '11'
     targetCompatibility = '11'
+}
+
+import com.github.jengelman.gradle.plugins.shadow.transformers.AppendingTransformer
+
+shadowJar {
+    transform(AppendingTransformer) {
+        resource = 'reference.conf' // Required for akka
+    }
+    with jar
 }

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

@@ -28,7 +28,9 @@ public class Main {
         entry("jetty-client", new JettyClientCase()),
         entry("async-http-client", new AsyncHttpClientCase()),
         entry("spring-web", new SpringWebClientCase()),
-        entry("ktor-cio", new KtorCioCase())
+        entry("ktor-cio", new KtorCioCase()),
+        entry("akka-req-http", new AkkaRequestClientCase()),
+        entry("akka-host-http", new AkkaHostClientCase())
     );
 
     public static void main(String[] args) throws Exception {

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

@@ -0,0 +1,54 @@
+package tech.httptoolkit.testapp.cases;
+
+import akka.NotUsed;
+import akka.actor.ActorSystem;
+import akka.actor.ExtendedActorSystem;
+import akka.http.javadsl.HostConnectionPool;
+import akka.http.javadsl.Http;
+import akka.http.javadsl.model.HttpRequest;
+import akka.http.javadsl.model.HttpResponse;
+import akka.japi.Pair;
+import akka.stream.javadsl.Flow;
+import akka.stream.javadsl.Sink;
+import akka.stream.javadsl.Source;
+import scala.util.Try;
+
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+
+// Here we use a host-based persistent Akka connection pool, just to make sure that that case is covered too, in
+// addition to the request client case.
+public class AkkaHostClientCase extends ClientCase<
+    Flow<Pair<HttpRequest, NotUsed>, Pair<Try<HttpResponse>, NotUsed>, HostConnectionPool>
+> {
+
+    private static final ActorSystem system = ExtendedActorSystem.create();
+
+    @Override
+    public Flow<Pair<HttpRequest, NotUsed>, Pair<Try<HttpResponse>, NotUsed>, HostConnectionPool> newClient(String url) throws Exception {
+        URI uri = new URI(url);
+        return new Http((ExtendedActorSystem) system)
+            .cachedHostConnectionPool(uri.getHost());
+    }
+
+    @Override
+    public int test(
+            String url,
+            Flow<Pair<HttpRequest, NotUsed>, Pair<Try<HttpResponse>, NotUsed>, HostConnectionPool> clientFlow
+    ) throws URISyntaxException, ExecutionException, InterruptedException {
+        Source<Pair<HttpRequest, NotUsed>, NotUsed> requestSource = Source.single(
+            new Pair<>(HttpRequest.create(url), null)
+        );
+
+        CompletableFuture<Pair<Try<HttpResponse>, NotUsed>> responseFuture = requestSource
+            .via(clientFlow)
+            .runWith(Sink.head(), system)
+            .toCompletableFuture();
+
+        HttpResponse response = responseFuture.get().first().get();
+        response.discardEntityBytes(system);
+        return response.status().intValue();
+    }
+}

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

@@ -0,0 +1,29 @@
+package tech.httptoolkit.testapp.cases;
+
+import akka.actor.ActorSystem;
+import akka.http.javadsl.Http;
+import akka.http.javadsl.model.HttpRequest;
+import akka.http.javadsl.model.HttpResponse;
+
+import java.net.URISyntaxException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+
+public class AkkaRequestClientCase extends ClientCase<Http> {
+
+    private static final ActorSystem system = ActorSystem.create();
+
+    @Override
+    public Http newClient(String url) throws Exception {
+        return Http.get(system);
+    }
+
+    @Override
+    public int test(String url, Http client) throws URISyntaxException, ExecutionException, InterruptedException {
+        CompletableFuture<HttpResponse> responseFuture = client
+                .singleRequest(HttpRequest.create(url))
+                .toCompletableFuture();
+        HttpResponse response = responseFuture.get();
+        return response.status().intValue();
+    }
+}