Add support for older Akka HTTP versions

Author: pimterry

README.md

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

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

@@ -1,18 +1,31 @@
 package tech.httptoolkit.javaagent.advice.akka;
 
-import akka.http.scaladsl.ClientTransport;
-import akka.http.scaladsl.ConnectionContext;
-import akka.http.scaladsl.settings.ClientConnectionSettings;
+import akka.http.javadsl.ClientTransport;
+import akka.http.javadsl.ConnectionContext;
+import akka.http.javadsl.settings.ClientConnectionSettings;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import tech.httptoolkit.javaagent.HttpProxyAgent;
 
 import java.net.InetSocketAddress;
+import java.util.Arrays;
 
 public class OverrideHttpSettingsAdvice {
+
+    public static final boolean hasHttpsSettingsMethod =
+            Arrays.stream(ConnectionContext.class.getDeclaredMethods())
+                    .anyMatch(method -> method.getName().equals("httpsClient"));
+
+    public static final ConnectionContext interceptedConnectionContext = hasHttpsSettingsMethod
+            // For 10.2+:
+            ? ConnectionContext.httpsClient(HttpProxyAgent.getInterceptedSslContext())
+            // For everything before then:
+            : akka.http.javadsl.ConnectionContext.https(HttpProxyAgent.getInterceptedSslContext());
+
     @Advice.OnMethodEnter
     public static void beforeOutgoingConnection(
-        @Advice.Argument(value = 2, readOnly = false) ClientConnectionSettings clientSettings,
-        @Advice.Argument(value = 3, readOnly = false) ConnectionContext connectionContext
+        @Advice.Argument(value = 2, readOnly = false, typing = Assigner.Typing.DYNAMIC) ClientConnectionSettings clientSettings,
+        @Advice.Argument(value = 3, readOnly = false, typing = Assigner.Typing.DYNAMIC) ConnectionContext connectionContext
     ) {
         // Change all new outgoing connections to use the proxy:
         clientSettings = clientSettings.withTransport(
@@ -23,6 +36,6 @@ public static void beforeOutgoingConnection(
         );
 
         // Change all new outgoing connections to trust our certificate:
-        connectionContext = ConnectionContext.httpsClient(HttpProxyAgent.getInterceptedSslContext());
+        connectionContext = interceptedConnectionContext;
     }
 }

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

@@ -0,0 +1,53 @@
+package tech.httptoolkit.javaagent.advice.akka;
+
+import akka.http.impl.settings.HostConnectionPoolSetup;
+import akka.http.scaladsl.ClientTransport;
+import net.bytebuddy.asm.Advice;
+import scala.concurrent.Await;
+import scala.concurrent.Future;
+import scala.concurrent.duration.Duration;
+
+import java.lang.reflect.Method;
+import java.util.Collections;
+import java.util.Set;
+import java.util.WeakHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+// This is very similar to ResetOldPoolsAdvice, but applies to older Akka setups, which use many PoolGateway instances,
+// one per config, rather than one PoolMaster instance. Otherwise the logic should be identical.
+public class ResetOldGatewaysAdvice {
+
+    public static Set<HostConnectionPoolSetup> resetPoolSetups = Collections.newSetFromMap(new WeakHashMap<>());
+
+    @Advice.OnMethodEnter
+    public static void beforeDispatchRequest(
+        @Advice.This Object thisPoolGateway,
+        @Advice.FieldValue(value = "hcps") HostConnectionPoolSetup poolSetup
+    ) throws Exception {
+        // If a pool config has been changed to use our proxy already, then we're perfect
+        ClientTransport transport = poolSetup.setup().settings().transport();
+        boolean alreadyIntercepted = transport == ResetPoolSetupAdvice.interceptedProxyTransport;
+        // If not, it's still OK, as long as we've previously reset the pool to ensure the connection was
+        // re-established (we hook connection setup too, so all new conns are intercepted, even with old config)
+        boolean alreadyReset = resetPoolSetups.contains(poolSetup);
+
+        if (alreadyIntercepted || alreadyReset) return;
+
+        // Otherwise this is a request to use a pre-existing connection pool which probably has connections open that
+        // aren't using our proxy. We shutdown the pool before the request. It'll be restarted automatically when
+        // the request does go through, but this ensures we re-establish connections (so it definitely gets intercepted)
+        Method shutdownMethod = thisPoolGateway.getClass().getDeclaredMethod("shutdown");
+
+        Future<?> shutdownFuture = (Future<?>) shutdownMethod.invoke(thisPoolGateway);
+
+        // We wait a little, just to ensure the shutdown is definitely started before this request is dispatched.
+        try {
+            Await.result(shutdownFuture, Duration.apply(10, TimeUnit.MILLISECONDS));
+        } catch (TimeoutException ignored) {}
+
+        // Lastly, we remember this pool setup, so that we don't unnecessarily reset it again in future:
+        resetPoolSetups.add(poolSetup);
+    }
+
+}

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

@@ -1,9 +1,10 @@
 package tech.httptoolkit.javaagent.advice.akka;
 
 import akka.http.scaladsl.ClientTransport;
-import akka.http.scaladsl.ConnectionContext;
 import akka.http.scaladsl.settings.ConnectionPoolSettings;
+import akka.http.javadsl.ConnectionContext;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import tech.httptoolkit.javaagent.HttpProxyAgent;
 
 import java.net.InetSocketAddress;
@@ -21,13 +22,13 @@ public class ResetPoolSetupAdvice {
 
     @Advice.OnMethodExit
     public static void afterConstructor(
-        @Advice.FieldValue(value = "settings", readOnly = false) ConnectionPoolSettings settings,
-        @Advice.FieldValue(value = "connectionContext", readOnly = false) ConnectionContext connContext
+        @Advice.FieldValue(value = "settings", readOnly = false, typing = Assigner.Typing.DYNAMIC) ConnectionPoolSettings settings,
+        @Advice.FieldValue(value = "connectionContext", readOnly = false, typing = Assigner.Typing.DYNAMIC) ConnectionContext connContext
     ) {
         // Change all new outgoing connections to use the proxy:
         settings = settings.withTransport(interceptedProxyTransport);
 
         // Change all new outgoing connections to trust our certificate:
-        connContext = ConnectionContext.httpsClient(HttpProxyAgent.getInterceptedSslContext());
+        connContext = OverrideHttpSettingsAdvice.interceptedConnectionContext;
     }
 }

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

@@ -106,7 +106,8 @@ fun interceptAllHttps(config: Config, instrumentation: Instrumentation) {
         KtorClientTlsTransformer(logger),
         AkkaHttpTransformer(logger),
         AkkaPoolSettingsTransformer(logger),
-        AkkaPoolTransformer(logger)
+        AkkaPoolTransformer(logger),
+        AkkaGatewayTransformer(logger)
     ).forEach { matchingAgentTransformer ->
         agentBuilder = matchingAgentTransformer.register(agentBuilder)
     }

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

@@ -58,4 +58,21 @@ class AkkaPoolTransformer(logger: TransformationLogger) : MatchingAgentTransform
                     .on(hasMethodName("dispatchRequest"))
             )
     }
-}
+}
+
+// The above works perfectly for new Akka, but as a last step we duplicate the 3rd step for slightly older versions:
+class AkkaGatewayTransformer(logger: TransformationLogger) : MatchingAgentTransformer(logger) {
+    override fun register(builder: AgentBuilder): AgentBuilder {
+        return builder
+            .type(named("akka.http.impl.engine.client.PoolGateway")) // Exists on <10.2.0
+            .transform(this)
+    }
+
+    override fun transform(builder: DynamicType.Builder<*>): DynamicType.Builder<*> {
+        return builder
+            .visit(
+                Advice.to(ResetOldGatewaysAdvice::class.java)
+                    .on(hasMethodName("apply"))
+            )
+    }
+}