Skip to content

Instantly share code, notes, and snippets.

@kriegaex
Forked from joshgord/BootstrapAgent.java
Last active January 19, 2024 03:52
Show Gist options
  • Save kriegaex/0f4dd4c9d05f3a68e3a8e1ed75359c3b to your computer and use it in GitHub Desktop.
Save kriegaex/0f4dd4c9d05f3a68e3a8e1ed75359c3b to your computer and use it in GitHub Desktop.
An example agent that intercepts a method of the bootstrap class loader. Tested with Byte Buddy 1.14.11 on JDKs 8 to 21.
package net.bytebuddy;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.loading.ClassInjector;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.implementation.bind.annotation.SuperCall;
import net.bytebuddy.matcher.ElementMatchers;
import java.lang.instrument.Instrumentation;
import java.net.URL;
import java.util.Collections;
import java.util.concurrent.Callable;
import static net.bytebuddy.matcher.ElementMatchers.none;
/**
* Inspired by <a href="https://github.com/apache/skywalking">Apache SkyWalking</a>, specifically
* <a href="https://github.com/apache/skywalking/blob/bc64c6a12770031478d29e2f19004796584374c9/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/plugin/bootstrap/BootstrapInstrumentBoost.java">
* this class</a>. Discussed in <a href="https://github.com/raphw/byte-buddy/issues/697">Byte Buddy issue #697</a>.
* <p>
* Successfully tested on JDKs 8 to 21. Should print:
* <pre>
* Intercepted!
* GET
* </pre>
*/
public class BootstrapAgent {
public static void main(String[] args) throws Exception {
premain(null, ByteBuddyAgent.install());
Object urlConnection = new URL("http://www.google.com").openConnection();
System.out.println(urlConnection.getClass().getMethod("getRequestMethod").invoke(urlConnection));
}
public static void premain(String arg, Instrumentation instrumentation) throws Exception {
ClassInjector.UsingUnsafe.Factory factory = ClassInjector.UsingUnsafe.Factory.resolve(instrumentation);
factory.make(null, null).injectRaw(
Collections.singletonMap(
MyInterceptor.class.getName(),
ClassFileLocator.ForClassLoader.read(MyInterceptor.class)
)
);
AgentBuilder agentBuilder = new AgentBuilder.Default();
agentBuilder = agentBuilder.with(new AgentBuilder.InjectionStrategy.UsingUnsafe.OfFactory(factory));
agentBuilder
.ignore(none())
.assureReadEdgeFromAndTo(instrumentation, Class.forName("java.net.HttpURLConnection"))
.assureReadEdgeFromAndTo(instrumentation, MyInterceptor.class)
.ignore(ElementMatchers.nameStartsWith("net.bytebuddy."))
.type(ElementMatchers.nameContains("HttpURLConnection"))
.transform((builder, typeDescription, classLoader, module, protectionDomain) -> builder
.method(ElementMatchers.named("getRequestMethod"))
.intercept(MethodDelegation.to(MyInterceptor.class))
)
.installOn(instrumentation);
}
public static class MyInterceptor {
public static String intercept(@SuperCall Callable<String> zuper) throws Exception {
System.out.println("Intercepted!");
return zuper.call();
}
}
}
@Ch35Tnut
Copy link

It works on me. Just want you to know that I'm not lying.And this strategy has been proven in Project.

That's so weird. I know the Class String was loaded before already. So it need to be retransformed, like this:

Class[] classes = instrumentation.getAllLoadedClasses();
for (Class clazz : classes) {
    if (instrumentation.isModifiableClass(clazz)) {
        instrumentation.retransformClasses(clazz);
    }
}

My sample program bellows , your sample program wokrs on my agent too.

public class Main {
    public static void main(String[] args) {
        while (true){
            String input = new Scanner(System.in).nextLine();
            System.out.println(input.concat(", hello!"));
        }
    }
}

Input:

ssss

Output

Enhance String::concat
ssss, hello!

Maybe my config is useful.
build agent with IDEA:
image
add VM option on sample project, -javaagent:C:\Users\A\Desktop\Code\Java\EnhanceStringAgent-main\target\EnhanceStringAgent-1.0-SNAPSHOT.jar
image
RUN sample project ,console log:
image

@kriegaex
Copy link
Author

kriegaex commented Jan 18, 2024

OK, I noticed that it started working for me locally when I downgraded my runtime from JDK 17 to 8, 11 or even 16. So I upgraded your code to use the latest ASM v9.6 and also bumped the ASM API version 8 in your project to Opcodes.ASM9. Now, it also works on JDKs 17 to 21. That is a starting point. I do not have any more time tonight to look into this any further, but tomorrow maybe I can spare a little bit of time.

I guess, I did something similar a few years ago in my testing tool Sarek, which I never got around to promote, even though it is feature-complete. (Only the API is not so easy to use. But I am degressing.) Basically, I can already say that if it works with ASM, there is no reason to believe that it will not work with Byte Buddy, too.

@kriegaex
Copy link
Author

kriegaex commented Jan 19, 2024

@Ch35Tnut, what you want is possible, like I said yesterday. But you need to avoid the interceptor approach in favour of an advice-based one and also make sure that you do not alter the original method signature. Finally, you need to issue retransformation, just like your custom ASM agent does.

See https://gist.github.com/kriegaex/8228c8ba664c730157c2fef3c5fd78e7.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment