Is it possible to encapsulate a native method in a thread using ByteBuddy Advice mechanism?

Viewed 145

I am trying to delay writing to a socket in a separate thread to simulate latency in the network stack using an javaagent. Using ByteBuddy, I managed to decorate the native java.net.SocketOutputStream::writeSocket0. However, doing so, I pause the current thread which is not what I want to achieve. To avoid this issue, I thought of creating a thread from which the native method can be called. Unfortunately, my knowledge of ByteBuddy remains rather limited and I didn't manage to find a way to achieve this. My best attempts looks like the code below:

//premain function
public static void premain(final String agentArgs, final Instrumentation instrumentation) {
    try{
        BootLoaderInjector.inject(DelayedExecutor.class);
        long duration = 100;

        new AgentBuilder.Default()
            .with(new ByteBuddy().with(Implementation.Context.Disabled.Factory.INSTANCE))
            .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
            .with(AgentBuilder.TypeStrategy.Default.REBASE)
            .enableNativeMethodPrefix("somePrefix")
            .with(AgentBuilder.Listener.StreamWriting.toSystemError().withErrorsOnly())
            .ignore(nameStartsWith("net.bytebuddy."))
            .type(named("java.net.SocketOutputStream"))
            .transform((transformer, typeDescription, classLoader, javaModule) ->
                transformer.method(named("socketWrite0"))
                    .intercept(
                        Advice.withCustomMapping().bind(Duration.class, duration)
                        .to(DelayInterceptor.class)
                )
            )
            .installOn(instrumentation);
    }
    catch (Throwable t){
        System.err.printf(t.getMessage());
    }
}

// interceptor class
public class DelayInterceptor {
    @Advice.OnMethodEnter
    public static void intercept(@SuperCall Callable<?> caller, @Duration long duration) {
        DelayedExecutor.execute(callable, duration);
    }
}

// execute the intercepted method in a new thread with a delay
public class DelayedExecutor implements Runnable {
    private final Callable<?> callable;
    private final long delay;

    private DelayedExecutor(Callable<?> callable, long delay) {
        this.callable = callable;
        this.delay = delay;
    }

    public static void execute(Callable<?> callable, long delay){
        new Thread(new DelayedExecutor(callable, delay)).start();
    }

    @Override
    public void run() {
        try {
            Thread.sleep(this.delay);
            this.callable.call();
        } catch (Exception ignore) {}
    }
}

However, this fails because @Callable is not part of the Advice API.

Not using the advice seems to pose serious limitations. Indeed using MethodDelegation raises the following error:

Registration of auxiliary types was disabled: net.bytebuddy.implementation.auxiliary.MethodCallProxy@3206bfce

Is there a way to wrap a function call inside a block using the Advice API of ByteBuddy?

Edit

To add more clarity and to address the answer suggestion, here is the implementation of BootLoaderInjector.inject which shows how types are added to the bootloader classloader relying on ClassInjector.UsingUnsafe.ofBootLoader():

public class BootLoaderInjector {
    public static void inject(Class<?> targetClass) {
        try {
            ClassInjector.UsingUnsafe.ofBootLoader().inject(singletonMap(
                    new TypeDescription.ForLoadedType(targetClass),
                    ClassFileLocator.ForClassLoader.read(targetClass)
            ));
        } catch (Throwable e){
            System.out.println("Something went terribly wrong: " + e.getMessage());
            e.printStackTrace();
        }
    }
}
1 Answers

You would need to register an InjectionStrategy to inject these required types into the targeted class loader. For the boot loader, this could for example be one UsingInstrumentation. More efficiently, but relying on JVM internal API, you could also use UsingUnsafe.

Related