BrokenAgent.java

import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.method.MethodList;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.dynamic.scaffold.InstrumentedType;
import net.bytebuddy.implementation.Implementation;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.implementation.bind.annotation.Origin;
import net.bytebuddy.implementation.bind.annotation.RuntimeType;
import net.bytebuddy.implementation.bind.annotation.SuperCall;
import net.bytebuddy.implementation.bytecode.ByteCodeAppender;
import net.bytebuddy.jar.asm.MethodVisitor;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.utility.JavaModule;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

public final class BrokenAgent
{
    static void load() throws IOException
    {
        performInstrumentation(ByteBuddyAgent.install());
    }

    private static void performInstrumentation(final Instrumentation instrumentation)
    {
        final ConcurrentMap<String, PerMethodInterceptor> interceptorByMethodMap = new ConcurrentHashMap<>();
        AgentBuilder agentBuilder = new AgentBuilder.Default().with(new AgentBuilder.Listener()
        {
            @Override
            public void onTransformation(final TypeDescription typeDescription, final ClassLoader classLoader, final JavaModule module, final DynamicType dynamicType)
            {
            }

            @Override
            public void onIgnored(final TypeDescription typeDescription, final ClassLoader classLoader, final JavaModule module)
            {
            }

            @Override
            public void onError(final String typeName, final ClassLoader classLoader, final JavaModule module, final Throwable throwable)
            {
                System.out.printf("onError: %s%n", typeName);
            }

            @Override
            public void onComplete(final String typeName, final ClassLoader classLoader, final JavaModule module)
            {
            }
        });

        agentBuilder = instrumentType(agentBuilder, isInstrumentationTarget(),
                ElementMatchers.declaresMethod(isInstrumentationTarget()),
                interceptorByMethodMap);
        final ClassFileTransformer classFileTransformer = agentBuilder.installOn(instrumentation);
    }

    private static final class Instrumenter implements Implementation
    {
        private final ConcurrentMap<String, PerMethodInterceptor> interceptorByMethodMap;

        public Instrumenter(final ConcurrentMap<String, PerMethodInterceptor> interceptorByMethodMap)
        {
            this.interceptorByMethodMap = interceptorByMethodMap;
        }

        @Override
        public ByteCodeAppender appender(final Target implementationTarget)
        {
            return new ByteCodeAppender()
            {
                @Override
                public Size apply(final MethodVisitor methodVisitor, final Context implementationContext,
                                  final MethodDescription instrumentedMethod)
                {
                    final PerMethodInterceptor delegate =
                            interceptorByMethodMap.get(instrumentedMethod.toGenericString() + "_fail_to_lookup");
                    return MethodDelegation.to(delegate).
                            appender(implementationTarget).apply(methodVisitor, implementationContext, instrumentedMethod);
                }
            };

        }

        @Override
        public InstrumentedType prepare(final InstrumentedType instrumentedType)
        {
            InstrumentedType updated = instrumentedType;
            final MethodList<MethodDescription.InDefinedShape> declaredMethods =
                    instrumentedType.getDeclaredMethods().filter(isInstrumentationTarget());
            for (MethodDescription.InDefinedShape declaredMethod : declaredMethods)
            {
                final PerMethodInterceptor perMethodInterceptor = interceptorByMethodMap.computeIfAbsent(
                        declaredMethod.toGenericString(), k ->
                {
                    return new PerMethodInterceptor();
                });
                updated = MethodDelegation.to(perMethodInterceptor).
                        filter(ElementMatchers.is(declaredMethod)).prepare(updated);
            }
            return updated;
        }
    }

    private static ElementMatcher.Junction<MethodDescription> isInstrumentationTarget()
    {
        return ElementMatchers.nameMatches("foo.*|bar");
    }

    private static AgentBuilder instrumentType(final AgentBuilder agentBuilder,
                                               final ElementMatcher.Junction<MethodDescription> methodMatcher,
                                               final ElementMatcher.Junction<TypeDescription> typeMatcher,
                                               final ConcurrentMap<String, PerMethodInterceptor> interceptorByMethodMap)
    {
        return agentBuilder.
                type(typeMatcher).
                transform(new AgentBuilder.Transformer()
        {
            @Override
            public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription arg1, ClassLoader arg2)
            {
                return builder.
                        method(methodMatcher).
                        intercept(new Instrumenter(interceptorByMethodMap));
            }

        });
    }

    private static final class PerMethodInterceptor
    {
        @RuntimeType
        public Object intercept(@Origin Method m, @SuperCall Callable<?> zuper) throws Exception
        {
            try
            {
                return zuper.call();
            }
            finally
            {
                System.out.println("method intercepted: " + m.getDeclaringClass().getSimpleName() + "." + m.getName());
                System.out.println("by interceptor: " + System.identityHashCode(this));
            }
        }
    }
}

BrokenAgentTest.java

import org.junit.Ignore;
import org.junit.Test;

public final class BrokenAgentTest
{
    @Test
    public void shouldName() throws Exception
    {
        BrokenAgent.load();

        new Foo().bar(42);
        new Bar().foo(17);
        new Bar().foo2(37);
        new Bar().foo2(42);
    }

    private static final class Foo
    {
        private String bar(final int number)
        {
            return Integer.toHexString(number);
        }
    }

    private static final class Bar
    {
        private String foo(final int number)
        {
            return Integer.toHexString(number);
        }
        private long foo2(final long number)
        {
            return number ^ 238764L << 34;
        }
    }
}