Update wrapping system: generalize logic for loading in agent_api.jar classes in a separate agent_bootstrap

Makefile

@@ -14,6 +14,10 @@ build: clean check_binaries
 	./gradlew agent_api:shadowJar
 	cp agent_api/build/libs/agent*-all.jar dist/agent_api.jar
 
+
+	./gradlew agent_bootstrap:shadowJar
+	cp agent_bootstrap/build/libs/agent*-all.jar dist/agent_bootstrap.jar
+
 mock_init:
 	docker kill mock_core && docker rm mock_core
 	cd end2end/server && docker build -t mock_core .

agent/build.gradle

@@ -5,6 +5,7 @@ plugins {
 
 dependencies {
     implementation project(':agent_api')
+    compileOnly project(':agent_bootstrap')
     implementation 'net.bytebuddy:byte-buddy:1.15.11'
     // Compile only for interface types :
     compileOnly 'jakarta.servlet:jakarta.servlet-api:6.1.0' // spring 3 -> jakarta

agent/src/main/java/dev/aikido/agent/Agent.java

@@ -6,11 +6,19 @@
 import dev.aikido.agent_api.helpers.logging.LogManager;
 import dev.aikido.agent_api.helpers.logging.Logger;
 import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
+import net.bytebuddy.jar.asm.MethodVisitor;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
 
 import java.io.File;
+import java.io.IOException;
 import java.lang.instrument.Instrumentation;
+import java.util.jar.JarFile;
 
 import static dev.aikido.agent.ByteBuddyInitializer.createAgentBuilder;
 import static dev.aikido.agent.DaemonStarter.startDaemon;
@@ -19,17 +27,31 @@
 
 public class Agent {
     private static final Logger logger = LogManager.getLogger(Agent.class);
+
     public static void premain(String agentArgs, Instrumentation inst) {
         // Check for 'AIKIDO_DISABLE' :
         if (new BooleanEnv("AIKIDO_DISABLE", /*default value*/ false).getValue()) {
             return; // AIKIDO_DISABLE is true, so we will not be wrapping anything.
         }
         logger.info("Zen by Aikido v%s starting.", Config.pkgVersion);
-        setAikidoSysProperties();
+        try {
+            setAikidoSysProperties(inst);
+        } catch (Exception e) {
+            logger.error(e.getMessage());
+        }
+
+        // Modify bootstrap class path (includes the core Java classes), so we can have some shared
+        // code for core java classes when wrapping.
+        try {
+            inst.appendToBootstrapClassLoaderSearch(new JarFile(getPathToAikidoDirectory() + "/agent_bootstrap.jar"));
+        } catch (IOException e) {
+            logger.error(e.getMessage());
+        }
 
         // Test loading of zen binaries :
         loadLibrary();
 
+
         ElementMatcher.Junction wrapperTypeDescriptors = ElementMatchers.none();
         for(Wrapper wrapper: WRAPPERS) {
             wrapperTypeDescriptors = wrapperTypeDescriptors.or(wrapper.getTypeMatcher());
@@ -56,11 +78,14 @@ public static AgentBuilder.Transformer get() {
             return adviceAgentBuilder;
         }
     }
-    private static void setAikidoSysProperties() {
+    private static String getPathToAikidoDirectory() {
         String pathToAgentJar = Agent.class.getProtectionDomain().getCodeSource().getLocation().getPath();
-        String pathToAikidoDirectory = new File(pathToAgentJar).getParent();
-        String jarPath = "file:" + pathToAikidoDirectory + "/agent_api.jar";
-        System.setProperty("AIK_agent_dir", pathToAikidoDirectory);
+        return new File(pathToAgentJar).getParent();
+    }
+    private static void setAikidoSysProperties(Instrumentation inst) throws IOException {
+
+        String jarPath = "file:" + getPathToAikidoDirectory() + "/agent_api.jar";
+        System.setProperty("AIK_agent_dir", getPathToAikidoDirectory());
         System.setProperty("AIK_agent_api_jar", jarPath);
     }
-}
+}

agent/src/main/java/dev/aikido/agent/wrappers/HttpClientSendWrapper.java

@@ -1,21 +1,16 @@
 package dev.aikido.agent.wrappers;
 
+import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
 
-import java.lang.reflect.Method;
-import java.net.HttpURLConnection;
-import java.net.MalformedURLException;
 import java.net.URL;
-import java.net.URLClassLoader;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 
-import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
-
 public class HttpClientSendWrapper implements Wrapper {
     public String getName() {
         // Wrap send(HttpRequest req, ...) and sendAsync(HttpRequest req, ...) on HttpClient instance
@@ -32,37 +27,15 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
         return ElementMatchers.isSubTypeOf(HttpClient.class);
     }
     public static class SendFunctionAdvice {
-        // Since we have to wrap a native Java Class stuff gets more complicated
-        // The classpath is not the same anymore, and we can't import our modules directly.
-        // To bypass this issue we load collectors from a .jar file
         @Advice.OnMethodEnter(suppress = Throwable.class)
         public static void before(
                 @Advice.Argument(0) HttpRequest httpRequest
-        ) throws Exception {
+        ) throws Throwable {
             if (httpRequest == null || httpRequest.uri() == null) {
                 return;
             }
-            String jarFilePath = System.getProperty("AIK_agent_api_jar");
-            URLClassLoader classLoader = null;
-            try {
-                URL[] urls = { new URL(jarFilePath) };
-                classLoader = new URLClassLoader(urls);
-            } catch (MalformedURLException ignored) {}
-            if (classLoader == null) {
-                return;
-            }
-
-            // Load the class from the JAR
-            Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.URLCollector");
-
-            // Run report with "argument"
-            for (Method method2: clazz.getMethods()) {
-                if(method2.getName().equals("report")) {
-                    method2.invoke(null, httpRequest.uri().toURL());
-                    break;
-                }
-            }
-            classLoader.close(); // Close the class loader
+            URL url = httpRequest.uri().toURL();
+            AikidoBootstrapClass.invoke(AikidoBootstrapClass.URL_COLLECTOR_REPORT, url);
         }
     }
 }

agent/src/main/java/dev/aikido/agent/wrappers/HttpClientWrapper.java

@@ -1,18 +1,16 @@
 package dev.aikido.agent.wrappers;
 
+import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
 
-import java.lang.reflect.Method;
 import java.net.*;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 
-import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
-
 public class HttpClientWrapper implements Wrapper {
     public String getName() {
         // Wrap getResponseCode function which executes HTTP requests
@@ -28,35 +26,14 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
         return ElementMatchers.isSubTypeOf(HttpClient.class);
     }
     public static class HttpClientAdvice {
-        // Since we have to wrap a native Java Class stuff gets more complicated
-        // The classpath is not the same anymore, and we can't import our modules directly.
-        // To bypass this issue we load collectors from a .jar file, specified with the AIKIDO_DIRECTORY env variable
         @Advice.OnMethodEnter(suppress = Throwable.class)
         public static void before(
-                @Advice.This(typing = DYNAMIC, optional = true) Object target,
                 @Advice.Argument(0) Object uriObject,
                 @Advice.Argument(2) Object httpRequestObject
-        ) throws Exception {
-            URI uri = (URI) uriObject;
-            HttpRequest httpRequest = (HttpRequest) httpRequestObject;
-
-            // Call to collector :
-            String jarFilePath = System.getProperty("AIK_agent_api_jar");
-            URL[] urls = { new URL(jarFilePath) };
-            URLClassLoader classLoader = new URLClassLoader(urls);
-
-            // Load the class from the JAR
-            Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.RedirectCollector");
-
-            // Run report func with its arguments
-            for (Method method2: clazz.getMethods()) {
-                if(method2.getName().equals("report")) {
-                    URL originUrl = httpRequest.uri().toURL();
-                    method2.invoke(null, originUrl, uri.toURL());
-                    break;
-                }
-            }
-            classLoader.close(); // Close the class loader
+        ) throws Throwable {
+            URL origin = ((HttpRequest) httpRequestObject).uri().toURL();
+            URL dest = ((URI) uriObject).toURL();
+            AikidoBootstrapClass.invoke(AikidoBootstrapClass.REDIRECT_COLLECTOR_REPORT, origin, dest);
         }
     }
 }

agent/src/main/java/dev/aikido/agent/wrappers/HttpConnectionRedirectWrapper.java

@@ -1,17 +1,13 @@
 package dev.aikido.agent.wrappers;
 
+import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
 
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
 import java.net.*;
-import java.net.http.HttpClient;
-import java.util.List;
-import java.util.Map;
 
 import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
 
@@ -32,38 +28,16 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
         return ElementMatchers.isSubTypeOf(HttpURLConnection.class);
     }
     public static class FollowRedirect0Advice {
-        // Since we have to wrap a native Java Class stuff gets more complicated
-        // The classpath is not the same anymore, and we can't import our modules directly.
-        // To bypass this issue we load collectors from a .jar file, specified with the AIKIDO_DIRECTORY env variable
         @Advice.OnMethodEnter(suppress = Throwable.class)
         public static void before(
                 @Advice.This(typing = DYNAMIC, optional = true) HttpURLConnection target,
                 @Advice.Argument(2) URL destUrl
-        ) throws Exception {
+        ) throws Throwable {
             URL origin = target.getURL();
             if (origin == null || destUrl == null) {
                 return;
             }
-            String jarFilePath = System.getProperty("AIK_agent_api_jar");
-            URLClassLoader classLoader = null;
-            try {
-                URL[] urls = { new URL(jarFilePath) };
-                classLoader = new URLClassLoader(urls);
-            } catch (MalformedURLException ignored) {}
-            if (classLoader == null) {
-                return;
-            }
-            // Load the class from the JAR
-            Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.RedirectCollector");
-
-            // Run report with "argument"
-            for (Method method2: clazz.getMethods()) {
-                if(method2.getName().equals("report")) {
-                    method2.invoke(null, origin, destUrl);
-                    break;
-                }
-            }
-            classLoader.close(); // Close the class loader
+            AikidoBootstrapClass.invoke(AikidoBootstrapClass.REDIRECT_COLLECTOR_REPORT, origin, destUrl);
         }
     }
 }

agent/src/main/java/dev/aikido/agent/wrappers/HttpURLConnectionWrapper.java

@@ -1,11 +1,11 @@
 package dev.aikido.agent.wrappers;
+
+import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import net.bytebuddy.matcher.ElementMatchers;
 
-import java.lang.reflect.Method;
 import java.net.*;
 
 import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
@@ -18,46 +18,20 @@ public String getName() {
         return ConstructorAdvice.class.getName();
     }
     public ElementMatcher<? super MethodDescription> getMatcher() {
-        return isConstructor().and(isDeclaredBy(is(HttpURLConnection.class)));
+        return isDeclaredBy(URL.class).and(named("openConnection"));
     }
 
     @Override
     public ElementMatcher<? super TypeDescription> getTypeMatcher() {
-        return ElementMatchers.is(HttpURLConnection.class);
+        return is(URL.class);
     }
 
     public static class ConstructorAdvice {
-        // Since we have to wrap a native Java Class stuff gets more complicated
-        // The classpath is not the same anymore, and we can't import our modules directly.
-        // To bypass this issue we load collectors from a .jar file
         @Advice.OnMethodExit(suppress = Throwable.class)
         public static void before(
-                @Advice.This(typing = DYNAMIC) Object target
-        ) throws Exception {
-            String jarFilePath = System.getProperty("AIK_agent_api_jar");
-            URLClassLoader classLoader = null;
-            try {
-                URL[] urls = { new URL(jarFilePath) };
-                classLoader = new URLClassLoader(urls);
-            } catch (MalformedURLException ignored) {}
-            if (classLoader == null) {
-                return;
-            }
-            URL url = ((HttpURLConnection) target).getURL();
-            if (target == null || url == null) {
-                return;
-            }
-            // Load the class from the JAR
-            Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.URLCollector");
-
-            // Run report with "argument"
-            for (Method method2: clazz.getMethods()) {
-                if(method2.getName().equals("report")) {
-                    method2.invoke(null, url);
-                    break;
-                }
-            }
-            classLoader.close(); // Close the class loader
+                @Advice.This(typing = DYNAMIC) URL url
+        ) throws Throwable {
+            AikidoBootstrapClass.invoke(AikidoBootstrapClass.URL_COLLECTOR_REPORT, url);
         }
     }
 }