Bytebuddy Not able to locate Custom Advice class file during tranformation

[prem0862]

Hi, I am trying to intercept the methods for ThreadPoolExecutor service and wrap them current thread context

Here is the code

public class Agent {
    public static void premain(String args, Instrumentation instrumentation) throws IOException, ClassNotFoundException {
        loadJar(instrumentation, configFilePath, parsedArgs);
    }

    public synchronized static void loadJar(Instrumentation instrumentation, String configFilePath, Map<String, String> parsedArgs) throws IOException {
        if (Boolean.getBoolean("Agent.attached")) {
            return;
        }

        String agentPath = "jar-path";
        try (JarFile jarFile = new JarFile(agentPath)) {
            instrumentation.appendToBootstrapClassLoaderSearch(jarFile);
        }

        try {
            // invoking via reflection to make sure the class is not loaded by the system classloader,
            // but only from the bootstrap classloader
            Class.forName("com.package.Agent", true, null)
                    .getMethod("init", String.class, Map.class, Instrumentation.class)
                    .invoke(null, configFilePath, parsedArgs, instrumentation);

            init(configFilePath, parsedArgs, instrumentation);
        } catch (Exception exception) {
            System.out.println("Method not found :" + exception);
        }

        System.setProperty("Agent.attached", Boolean.TRUE.toString());
    }

    public synchronized static void init(String configFilePath, Map<String, String> parsedArgs, Instrumentation instrumentation) {
        AgentConfig config;

        // Create AgentBuilder
        AgentBuilder agentBuilder = createAndInitialiseAgentBuilder();

        // Install the Agent on instrumentation
        agentBuilder.installOn(instrumentation);
    }

    public static AgentBuilder createAndInitialiseAgentBuilder() {

        AgentBuilder agentBuilder = new AgentBuilder.Default(
                new ByteBuddy().with(TypeValidation.DISABLED))
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                .with(AgentBuilder.RedefinitionStrategy.Listener.StreamWriting.toSystemError())
                .with(AgentBuilder.Listener.StreamWriting.toSystemError().withTransformationsOnly())
                .with(AgentBuilder.InstallationListener.StreamWriting.toSystemError());

        agentBuilder = agentBuilder.ignore(ElementMatchers.none());
        agentBuilder = agentBuilder.with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION);
        agentBuilder = agentBuilder.with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE);
        agentBuilder = agentBuilder.with(AgentBuilder.TypeStrategy.Default.REDEFINE);


        // Add Transformers to AgentBuilder
        for (AgentMatcherTransformer agentMatcherTransformer : getAgentMatcherTransformers()) {
            agentBuilder = agentBuilder.type(agentMatcherTransformer.getTypeMatcher()).transform(agentMatcherTransformer);
        }

        return agentBuilder;
    }

    public static ArrayList<AgentMatcherTransformer> getAgentMatcherTransformers() {
        ArrayList<AgentMatcherTransformer> agentMatcherTransformers = new ArrayList<>();
        agentMatcherTransformers.add(new ExecutorServiceTransformer());

        return agentMatcherTransformers;
    }
}

public class ExecutorServiceTransformer implements AgentMatcherTransformer {
    private final static String EXECUTOR_SERVICE_CLASS_NAME = "java.util.concurrent.ThreadPoolExecutor";

    @Override
    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
        // returns true when class matches and prefix start with
        Logger.info("Matching type for executor service...");
        return ElementMatchers.isSubTypeOf(Object.class)
                .and(ElementMatchers.named(EXECUTOR_SERVICE_CLASS_NAME));
    }


    @Override
    public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule javaModule, ProtectionDomain protectionDomain) {
        System.out.println(String.format("Starting transformation of ThreadPoolExecutor for class: %s", typeDescription.getName()));
        return builder
                // execute(Runnable)
                .visit(Advice.to(ThreadPoolExecutorAdvice.class)
                        .on(ElementMatchers.named("execute")
                                .and(ElementMatchers.takesArguments(Runnable.class))));
    }
}

public class ThreadPoolExecutorAdvice {

    @Advice.OnMethodEnter
    public static void onEnter(@Advice.Argument(value = 0, readOnly = false) Runnable task) {
        System.out.println(String.format("Intercepting execute(Runnable) method. Task type: %s", task.getClass().getName()));
        if (!(task instanceof ContextAwareRunnable)) {
            task = new ContextAwareRunnable(task);
        }
    }
}

public interface AgentMatcherTransformer extends AgentBuilder.Transformer {

    ElementMatcher<? super TypeDescription> getTypeMatcher();
}

But while running the app, advice class is not able to locate by bytebuddy

java.lang.IllegalStateException: Could not locate class file for com.package.advices.ThreadPoolExecutorAdvice
        at net.bytebuddy.dynamic.ClassFileLocator$Resolution$Illegal.resolve(ClassFileLocator.java:122)
        at net.bytebuddy.asm.Advice.to(Advice.java:406)
        at net.bytebuddy.asm.Advice.to(Advice.java:378)
        at net.bytebuddy.asm.Advice.to(Advice.java:355)
        at net.bytebuddy.asm.Advice.to(Advice.java:344)

How to resolve this?

[raphw]

When you appendToBootstrapClassLoaderSearch the class file's will not be added as resources. Advice does however allow you to provide a ClassFileLocator from which you can find the class files also at other locations.

[prem0862]

Added this class loader in Transformer and passed this as class loader in advice.

        ClassFileLocator adviceLocator = new ClassFileLocator.Compound(
                ClassFileLocator.ForClassLoader.ofSystemLoader(),
                ClassFileLocator.ForClassLoader.of(ThreadPoolExecutorAdvice.class.getClassLoader()),
                ClassFileLocator.ForClassLoader.of(getClass().getClassLoader())
        );

After passing it, I was able to resolve the issue with finding Advice classes. Thanks @raphw

However, I am facing an issue with Advice's dependencies not loading. for the below advice, okhttp3 is a dependency and facing. Since the provided fat JAR has all the dependencies packaged. It should be able to find the class from there.

One more thing, since we are loading the fat JAR into the boot loader it should be available by the system loader or any other loaders which has the Boot Loader as parent.

import okhttp3.Request;

public class Okhttp3Advice {

    @Advice.OnMethodEnter
    public static void intercept(@Advice.Argument(value = 0, readOnly = false) Request request ) {
        try{
            Request.Builder request = request.newBuilder();
        }
        catch (Exception e){
            Logger.error("Failed to intercept request: " + e.getMessage());
        }
    }
}

Stack trace are as below

[INFO] AgentRequestFilter registered successfully via ByteBuddy.
INFO  [2025-01-30 11:58:46,015] io.dropwizard.server.DefaultServerFactory: Registering jersey handler with root path prefix: /
INFO  [2025-01-30 11:58:46,017] io.dropwizard.server.DefaultServerFactory: Registering admin handler with root path prefix: /
[Byte Buddy] ERROR okhttp3.OkHttpClient [null, null, Thread[main,5,main], loaded=false]
net.bytebuddy.pool.TypePool$Resolution$NoSuchTypeException: Cannot resolve type description for okhttp3.Call$Factory
        at net.bytebuddy.pool.TypePool$Resolution$Illegal.resolve(TypePool.java:191)
        at net.bytebuddy.pool.TypePool$Default$WithLazyResolution$LazyTypeDescription.delegate(TypePool.java:1112)
        at net.bytebuddy.description.type.TypeDescription$AbstractBase$OfSimpleType$WithDelegation.getTypeVariables(TypeDescription.java:8528)