Make Scala interception more reliable for existing connection pools

Previously the tests would frequently fail because the pool wasn't
properly reset when interception began. This seems to be a race
condition if there's a request currently in progress, and it's something
that would likely resolve in practice _eventually_ but it's better to
detect cases and fix it properly per pool, rather than relying on
shutdownAll working reliably with concurrent requests.

Still more work to do here to make Akka reliable with older versions,
and check a couple of other edge cases.

Author: pimterry

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

@@ -0,0 +1,52 @@
+package tech.httptoolkit.javaagent.advice.akka;
+
+import akka.http.impl.engine.client.PoolId;
+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;
+
+public class ResetOldPoolsAdvice {
+
+    public static Set<PoolId> resetPoolIds = Collections.newSetFromMap(new WeakHashMap<>());
+
+    @Advice.OnMethodEnter
+    public static void beforeDispatchRequest(
+        @Advice.This Object thisPoolMaster,
+        @Advice.Argument(value = 0) PoolId poolId
+    ) throws Exception {
+        // If a pool config has been changed to use our proxy already, then we're perfect
+        ClientTransport transport = poolId.hcps().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 = resetPoolIds.contains(poolId);
+
+        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 = thisPoolMaster.getClass()
+                .getDeclaredMethod("shutdown", PoolId.class);
+
+        Future<?> shutdownFuture = (Future<?>) shutdownMethod.invoke(thisPoolMaster, poolId);
+
+        // 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 id, so that we don't unnecessarily reset it again in future:
+        resetPoolIds.add(poolId);
+    }
+
+}

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

@@ -0,0 +1,33 @@
+package tech.httptoolkit.javaagent.advice.akka;
+
+import akka.http.scaladsl.ClientTransport;
+import akka.http.scaladsl.ConnectionContext;
+import akka.http.scaladsl.settings.ConnectionPoolSettings;
+import net.bytebuddy.asm.Advice;
+import tech.httptoolkit.javaagent.HttpProxyAgent;
+
+import java.net.InetSocketAddress;
+
+public class ResetPoolSetupAdvice {
+
+    // We use this to avoid re-instantiating the proxy endlessly, but also to recognize intercepted
+    // and pre-existing settings configurations when they're used.
+    public static ClientTransport interceptedProxyTransport = ClientTransport.httpsProxy(
+        new InetSocketAddress(
+            HttpProxyAgent.getAgentProxyHost(),
+            HttpProxyAgent.getAgentProxyPort()
+        )
+    );
+
+    @Advice.OnMethodExit
+    public static void afterConstructor(
+        @Advice.FieldValue(value = "settings", readOnly = false) ConnectionPoolSettings settings,
+        @Advice.FieldValue(value = "connectionContext", readOnly = false) 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());
+    }
+}

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

@@ -105,6 +105,7 @@ fun interceptAllHttps(config: Config, instrumentation: Instrumentation) {
         KtorCioEngineTransformer(logger),
         KtorClientTlsTransformer(logger),
         AkkaHttpTransformer(logger),
+        AkkaPoolSettingsTransformer(logger),
         AkkaPoolTransformer(logger)
     ).forEach { matchingAgentTransformer ->
         agentBuilder = matchingAgentTransformer.register(agentBuilder)

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

@@ -3,16 +3,14 @@ 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
+import net.bytebuddy.matcher.ElementMatchers.*
+import tech.httptoolkit.javaagent.advice.akka.*
 
 // 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
+            .type(named("akka.http.scaladsl.HttpExt")) // Scala compiles Http()s methods as 'Ext' for some reason
             .transform(this)
     }
 
@@ -25,9 +23,27 @@ class AkkaHttpTransformer(logger: TransformationLogger) : MatchingAgentTransform
     }
 }
 
-// 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.
+// Second, when a connection pool setup is created (part of creating any connection pool, but also for
+// sending any individual request) we change its configuration. This isn't strictly necessary given the above,
+// but helps generally, and makes the 3rd step possible.
+class AkkaPoolSettingsTransformer(logger: TransformationLogger) : MatchingAgentTransformer(logger) {
+    override fun register(builder: AgentBuilder): AgentBuilder {
+        return builder
+            .type(named("akka.http.impl.settings.ConnectionPoolSetup"))
+            .transform(this)
+    }
+
+    override fun transform(builder: DynamicType.Builder<*>): DynamicType.Builder<*> {
+        return builder
+            .visit(
+                Advice.to(ResetPoolSetupAdvice::class.java)
+                    .on(isConstructor())
+            )
+    }
+}
+
+// Then, to ensure that any existing connections trust us too, we monitor all calls to dispatchRequest, and reset
+// any pools that don't have intercepted configuration (so preare -existing), just once per pool id.
 class AkkaPoolTransformer(logger: TransformationLogger) : MatchingAgentTransformer(logger) {
     override fun register(builder: AgentBuilder): AgentBuilder {
         return builder
@@ -38,7 +54,7 @@ class AkkaPoolTransformer(logger: TransformationLogger) : MatchingAgentTransform
     override fun transform(builder: DynamicType.Builder<*>): DynamicType.Builder<*> {
         return builder
             .visit(
-                Advice.to(ResetAllConnectionPoolsAdvice::class.java)
+                Advice.to(ResetOldPoolsAdvice::class.java)
                     .on(hasMethodName("dispatchRequest"))
             )
     }

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

@@ -62,15 +62,4 @@ class KtorCioEngineTransformer(logger: TransformationLogger) : MatchingAgentTran
             .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()
- *
-*/
+}

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

@@ -24,6 +24,7 @@ public int test(String url, Http client) throws URISyntaxException, ExecutionExc
                 .singleRequest(HttpRequest.create(url))
                 .toCompletableFuture();
         HttpResponse response = responseFuture.get();
+        response.discardEntityBytes(system);
         return response.status().intValue();
     }
 }