How to share some data between agent and app with bytebuddy

[white-blade]

Hi,

I am working on a project to intercept class java.net.Socket method connect(SocketAddress endpoint, int timeout) using byte buddy.

I create a class SocketWriterStorage to store the data. The goal is to store data in SocketWriterStorage from app and write the data to socket output steam every time after jave.net.Socket.connect is called.

public class SocketWriterStorage {
   private static final InheritableThreadLocal<SocketWriterStorage> holder = new InheritableThreadLocal<>();
   public static void set(SocketWriterStorage storage) {
        holder.set(storage);
    }
    public static SocketWriterStorage get() {
        return holder.get();
    }
}

• In the agent, I inject SocketWriterStorage into bootstrap class loader so that jave.net.Socket native methods could have access to the class.

File temp = Files.createTempDirectory("tmp").toFile();
Map<TypeDescription.ForLoadedType, byte[]> injectMap = new HashMap<>();
injectMap.put(new TypeDescription.ForLoadedType(SocketWriterStorage.class),
                    ClassFileLocator.ForClassLoader.read(SocketWriterStorage.class));

ClassInjector.UsingInstrumentation.of(temp, ClassInjector.UsingInstrumentation.Target.BOOTSTRAP, inst).inject(injectMap);

Then I intercept java.net.Socket.connect method

@Advice.OnMethodExit
public static void intercept(SocketAddress endpoint, int timeout) {
    SocketWriterStorage storage = SocketWriterStorage.get();
    System.out.println("Agent SocketWriterStorage classloader: " + SocketWriterStorage.class.getClassLoader());
}   

• In the app, it sets the storage data

public static void main(String[] args) {
    SocketWriterStorage.set(xxx);
    System.out.println("App SocketWriterStorage classloader: " + SocketWriterStorage.class.getClassLoader());
}

I realized this will not work as SocketWriterStorage in the agent is loaded by bootstrap loader (null) while SocketWriterStorage in the app is loaded by app loader (jdk.internal.loader.ClassLoaders$AppClassLoader@)

I tried -Xbootclasspath/a: when starting jvm wth jar containing need classes. I also tried defining a thread-safe context in agent, intercept app set storage method and set the context data then pass to jave.net.Socket.connect interceptor but neither worked.

May I know whether there is any possible way I can pass the data from client to agent and used by interceptor of java.net.Socket? Apologies if it is not the right place for this question. Appreciate any help here!

Thanks

[white-blade]

I found there is a way to write the data in a text file and read it in interceptor which is pretty hacky. Hope there will be a better solution.

[raphw]

Ideally, you create a different jar and add it as a resource to your agent jar. Upon startup of the agent, you extract the jar to a temporary location and then attach the jar to the boot loader. This way, you avoid that it gets loaded twice!

[white-blade]

Thanks @raphw for the quick response! You suggestion works perfectly for me! Also I realized I had the test app in the same build environment as the agent so some classes will be loaded by app class loader. Now I separate the test app code in another repo and force loading the classes from bootstrap using Java reflection (e.g. Class.forName(xxxx, true, null)) to make sure both agent and app are referring to the same class from bootstrap only.