net.bytebuddy.agent.builder.AgentBuilder Java Examples

/**
 * Add {@link ElementMatcher} for the hooks.
 */
private static AgentBuilder applyHooks(AgentBuilder agentBuilder, SortedSet<HookMetadata> hookMetadata, ClassLoaderCache classLoaderCache) {
    Map<String, SortedSet<MethodSignature>> instruments = getInstruments(hookMetadata);
    for (Map.Entry<String, SortedSet<MethodSignature>> entry : instruments.entrySet()) {
        String instrumentedClassName = entry.getKey();
        Set<MethodSignature> instrumentedMethods = entry.getValue();
        agentBuilder = agentBuilder
                .type(ElementMatchers.hasSuperType(named(instrumentedClassName)))
                .transform(new AgentBuilder.Transformer.ForAdvice()
                        .include(classLoaderCache.currentClassLoader()) // must be able to load PromagentAdvice
                        .advice(matchAnyMethodIn(instrumentedMethods), PromagentAdvice.class.getName())
                );
    }
    return agentBuilder;
}
 

private static AgentBuilder addDriverCommandInstrumentation(final AgentBuilder agentBuilder)
{
    if (CmdInterceptor.EVENTS.stream().noneMatch(DRIVER_EVENT_CODES::contains))
    {
        return agentBuilder;
    }

    return agentBuilder
        .type(nameEndsWith("ClientCommandAdapter"))
        .transform((builder, typeDescription, classLoader, javaModule) -> builder
            .visit(to(CmdInterceptor.class)
                .on(named("onMessage"))))
        .type(nameEndsWith("ClientProxy"))
        .transform((builder, typeDescription, classLoader, javaModule) -> builder
            .visit(to(CmdInterceptor.class)
                .on(named("transmit"))));
}
 

private static AgentBuilder initAgentBuilder(ElasticApmTracer tracer, Instrumentation instrumentation,
                                             Iterable<ElasticApmInstrumentation> instrumentations, Logger logger,
                                             AgentBuilder.DescriptionStrategy descriptionStrategy, boolean premain) {
    final CoreConfiguration coreConfiguration = tracer.getConfig(CoreConfiguration.class);
    ElasticApmAgent.instrumentation = instrumentation;
    final ByteBuddy byteBuddy = new ByteBuddy()
        .with(TypeValidation.of(logger.isDebugEnabled()))
        .with(FailSafeDeclaredMethodsCompiler.INSTANCE);
    AgentBuilder agentBuilder = getAgentBuilder(byteBuddy, coreConfiguration, logger, descriptionStrategy, premain);
    int numberOfAdvices = 0;
    for (final ElasticApmInstrumentation advice : instrumentations) {
        if (isIncluded(advice, coreConfiguration)) {
            numberOfAdvices++;
            agentBuilder = applyAdvice(tracer, agentBuilder, advice, new ElementMatcher.Junction.Conjunction<>(advice.getTypeMatcher(), not(isInterface())));
        }
    }
    logger.debug("Applied {} advices", numberOfAdvices);
    return agentBuilder;
}
 

private static synchronized void agent(final boolean shouldRedefine, final Instrumentation instrumentation)
{
    BufferAlignmentAgent.instrumentation = instrumentation;
    // all Int methods, and all String method other than
    // XXXStringWithoutLengthXXX or getStringXXX(int, int)
    final Junction<MethodDescription> intVerifierMatcher = nameContains("Int")
        .or(nameMatches(".*String[^W].*").and(not(ElementMatchers.takesArguments(int.class, int.class))));

    alignmentTransformer = new AgentBuilder.Default(new ByteBuddy().with(TypeValidation.DISABLED))
        .with(new AgentBuilderListener())
        .disableClassFormatChanges()
        .with(shouldRedefine ?
            AgentBuilder.RedefinitionStrategy.RETRANSFORMATION : AgentBuilder.RedefinitionStrategy.DISABLED)
        .type(isSubTypeOf(DirectBuffer.class).and(not(isInterface())))
        .transform((builder, typeDescription, classLoader, module) -> builder
            .visit(to(LongVerifier.class).on(nameContains("Long")))
            .visit(to(DoubleVerifier.class).on(nameContains("Double")))
            .visit(to(IntVerifier.class).on(intVerifierMatcher))
            .visit(to(FloatVerifier.class).on(nameContains("Float")))
            .visit(to(ShortVerifier.class).on(nameContains("Short")))
            .visit(to(CharVerifier.class).on(nameContains("Char"))))
        .installOn(instrumentation);
}
 

private AgentBuilder newAgentBuilder() {
        val byteBuddy = new ByteBuddy()
            .with(TypeValidation.of(config.isDebugMode()))
            .with(MethodGraph.Compiler.ForDeclaredMethods.INSTANCE);

        AgentBuilder agentBuilder = new AgentBuilder.Default(byteBuddy)
                .with(poolStrategyCache);


        agentBuilder = withRetransformationForRuntime(agentBuilder);
        agentBuilder = withBootstrapAttaching(agentBuilder);
        agentBuilder = withIgnore(agentBuilder);

        return agentBuilder
                .with(DefaultInstrumentationListener.instance())
                .with(additionalListeners());
}
 

private AgentBuilder newBuilder(final Instrumentation inst, final PluginManifest pluginManifest, final Event[] events) {
  // Prepare the builder to be used to implement transformations in AgentRule(s)
  AgentBuilder agentBuilder = new Default(byteBuddy);
  if (Adapter.tracerClassLoader != null)
    agentBuilder = agentBuilder.ignore(any(), is(Adapter.tracerClassLoader));

  agentBuilder = agentBuilder
    .ignore(nameStartsWith("net.bytebuddy.").or(nameStartsWith("sun.reflect.")).or(isSynthetic()), any(), any())
    .disableClassFormatChanges()
    .with(RedefinitionStrategy.RETRANSFORMATION)
    .with(InitializationStrategy.NoOp.INSTANCE)
    .with(TypeStrategy.Default.REDEFINE)
    .with(bootFallbackLocationStrategy);

  if (inst == null)
    return agentBuilder;

  if (pluginManifest != null)
    return agentBuilder.with(new TransformationListener(inst, pluginManifest, events));

  if (transformationListener == null)
    transformationListener = new TransformationListener(inst, null, events);

  return agentBuilder.with(transformationListener);
}
 

private static AgentBuilder addClusterElectionInstrumentation(final AgentBuilder agentBuilder)
{
    if (!CLUSTER_EVENT_CODES.contains(ClusterEventCode.ELECTION_STATE_CHANGE))
    {
        return agentBuilder;
    }

    return agentBuilder
        .type(nameEndsWith("Election"))
        .transform(((builder, typeDescription, classLoader, module) -> builder
            .visit(to(ClusterInterceptor.ElectionStateChange.class)
                .on(named("stateChange")))));
}
 

private static AgentBuilder addArchiveControlSessionInstrumentation(final AgentBuilder agentBuilder)
{
    if (!ARCHIVE_EVENT_CODES.contains(ArchiveEventCode.CONTROL_SESSION_STATE_CHANGE))
    {
        return agentBuilder;
    }

    return agentBuilder
        .type(nameEndsWith("ControlSession"))
        .transform(((builder, typeDescription, classLoader, module) -> builder
            .visit(to(ArchiveInterceptor.ControlSessionStateChange.class)
                .on(named("stateChange")))));
}
 

/**
 * @return
 */
public ClassFileTransformer getTransformer() {
    logger.info("Creating a class file transformer");
    AgentBuilder builder = new AgentBuilder.Default()
            .with(AgentBuilder.PoolStrategy.Default.FAST)
            .type(ElementMatchers.isSubTypeOf(Model.class).and(ElementMatchers.isAnnotatedWith(Entity.class)))
            .transform(this);
    return instrumentation != null ? builder.installOn(instrumentation) : builder.installOnByteBuddyAgent();
}
 

@Test
public void testTutorialGettingStartedJavaAgent() throws Exception {
    new AgentBuilder.Default().type(isAnnotatedWith(Rebase.class)).transform(new AgentBuilder.Transformer() {
        public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) {
            return builder.method(named("toString")).intercept(FixedValue.value("transformed"));
        }
    }).installOn(mock(Instrumentation.class));
}
 

/**
 * Remove the bytecode transformer and associated bytecode weaving so the alignment checks are not made.
 */
public static synchronized void removeTransformer()
{
    if (alignmentTransformer != null)
    {
        alignmentTransformer.reset(instrumentation, AgentBuilder.RedefinitionStrategy.RETRANSFORMATION);
        alignmentTransformer = null;
        BufferAlignmentAgent.instrumentation = null;
    }
}
 

private static AgentBuilder addClusterConsensusModuleAgentInstrumentation(final AgentBuilder agentBuilder)
{
    final boolean hasNewLeadershipTerm = CLUSTER_EVENT_CODES.contains(ClusterEventCode.NEW_LEADERSHIP_TERM);
    final boolean hasStateChange = CLUSTER_EVENT_CODES.contains(ClusterEventCode.STATE_CHANGE);
    final boolean hasRoleChange = CLUSTER_EVENT_CODES.contains(ClusterEventCode.ROLE_CHANGE);

    if (!hasNewLeadershipTerm && !hasStateChange && !hasRoleChange)
    {
        return agentBuilder;
    }

    return agentBuilder.type(nameEndsWith("ConsensusModuleAgent"))
        .transform(((builder, typeDescription, classLoader, module) ->
        {
            if (hasNewLeadershipTerm)
            {
                builder = builder
                    .visit(to(ClusterInterceptor.NewLeadershipTerm.class)
                        .on(named("onNewLeadershipTerm")));
            }
            if (hasStateChange)
            {
                builder = builder
                    .visit(to(ClusterInterceptor.ConsensusModuleStateChange.class)
                        .on(named("stateChange")));
            }
            if (hasRoleChange)
            {
                builder = builder
                    .visit(to(ClusterInterceptor.ConsensusModuleRoleChange.class)
                        .on(named("roleChange")));
            }

            return builder;
        }));
}
 

private static synchronized void agent(
    final AgentBuilder.RedefinitionStrategy redefinitionStrategy, final Instrumentation instrumentation)
{
    if (null != logTransformer)
    {
        throw new IllegalStateException("agent already instrumented");
    }

    EventConfiguration.init();

    if (DRIVER_EVENT_CODES.isEmpty() && ARCHIVE_EVENT_CODES.isEmpty() && CLUSTER_EVENT_CODES.isEmpty())
    {
        return;
    }

    EventLogAgent.instrumentation = instrumentation;

    readerAgentRunner = new AgentRunner(
        new SleepingMillisIdleStrategy(SLEEP_PERIOD_MS), Throwable::printStackTrace, null, newReaderAgent());

    AgentBuilder agentBuilder = new AgentBuilder.Default(new ByteBuddy()
        .with(TypeValidation.DISABLED))
        .disableClassFormatChanges()
        .with(new AgentBuilderListener())
        .with(redefinitionStrategy);

    agentBuilder = addDriverInstrumentation(agentBuilder);
    agentBuilder = addArchiveInstrumentation(agentBuilder);
    agentBuilder = addClusterInstrumentation(agentBuilder);

    logTransformer = agentBuilder.installOn(instrumentation);

    thread = new Thread(readerAgentRunner);
    thread.setName("event-log-reader");
    thread.setDaemon(true);
    thread.start();
}
 

private static AgentBuilder addClusterInstrumentation(final AgentBuilder agentBuilder)
{
    AgentBuilder tempBuilder = agentBuilder;
    tempBuilder = addClusterElectionInstrumentation(tempBuilder);
    tempBuilder = addClusterConsensusModuleAgentInstrumentation(tempBuilder);

    return tempBuilder;
}
 

private static AgentBuilder addArchiveReplaySessionInstrumentation(final AgentBuilder agentBuilder)
{
    if (!ARCHIVE_EVENT_CODES.contains(ArchiveEventCode.REPLAY_SESSION_ERROR))
    {
        return agentBuilder;
    }

    return agentBuilder
        .type(nameEndsWith("ReplaySession"))
        .transform(((builder, typeDescription, classLoader, module) -> builder
            .visit(to(ArchiveInterceptor.ReplaySession.class)
                .on(named("onPendingError")))));
}
 

private static AgentBuilder addDriverUdpChannelTransportInstrumentation(final AgentBuilder agentBuilder)
{
    final boolean hasFrameOut = DRIVER_EVENT_CODES.contains(DriverEventCode.FRAME_OUT);
    final boolean hasFrameIn = DRIVER_EVENT_CODES.contains(DriverEventCode.FRAME_IN);

    if (!hasFrameOut && !hasFrameIn)
    {
        return agentBuilder;
    }

    return agentBuilder
        .type(nameEndsWith("UdpChannelTransport"))
        .transform((builder, typeDescription, classLoader, javaModule) ->
        {
            if (hasFrameOut)
            {
                builder = builder
                    .visit(to(ChannelEndpointInterceptor.UdpChannelTransport.SendHook.class)
                        .on(named("sendHook")));
            }
            if (hasFrameIn)
            {
                builder = builder
                    .visit(to(ChannelEndpointInterceptor.UdpChannelTransport.ReceiveHook.class)
                        .on(named("receiveHook")));
            }

            return builder;
        });
}
 

private static AgentBuilder addDriverUntetheredSubscriptionInstrumentation(final AgentBuilder agentBuilder)
{
    if (!DRIVER_EVENT_CODES.contains(DriverEventCode.UNTETHERED_SUBSCRIPTION_STATE_CHANGE))
    {
        return agentBuilder;
    }

    return agentBuilder
        .type(nameEndsWith("UntetheredSubscription"))
        .transform((builder, typeDescription, classLoader, javaModule) ->
            builder
                .visit(to(DriverInterceptor.UntetheredSubscriptionStateChange.class)
                    .on(named("stateChange"))));
}
 

private static AgentBuilder addDriverConductorInstrumentation(final AgentBuilder agentBuilder)
{
    final boolean hasImageHook = DRIVER_EVENT_CODES.contains(DriverEventCode.REMOVE_IMAGE_CLEANUP);
    final boolean hasPublicationHook = DRIVER_EVENT_CODES.contains(DriverEventCode.REMOVE_PUBLICATION_CLEANUP);
    final boolean hasSubscriptionHook = DRIVER_EVENT_CODES.contains(DriverEventCode.REMOVE_SUBSCRIPTION_CLEANUP);

    if (!hasImageHook && !hasPublicationHook && !hasSubscriptionHook)
    {
        return agentBuilder;
    }

    return agentBuilder.type(nameEndsWith("DriverConductor"))
        .transform((builder, typeDescription, classLoader, javaModule) ->
        {
            if (hasImageHook)
            {
                builder = builder.visit(to(CleanupInterceptor.CleanupImage.class)
                    .on(named("cleanupImage")));
            }
            if (hasPublicationHook)
            {
                builder = builder.visit(to(CleanupInterceptor.CleanupPublication.class)
                    .on(named("cleanupPublication")));
            }
            if (hasSubscriptionHook)
            {
                builder = builder.visit(to(CleanupInterceptor.CleanupSubscriptionLink.class)
                    .on(named("cleanupSubscriptionLink")));
            }

            return builder;
        });
}
 

private static AgentBuilder addDriverInstrumentation(final AgentBuilder agentBuilder)
{
    AgentBuilder tempBuilder = agentBuilder;
    tempBuilder = addDriverConductorInstrumentation(tempBuilder);
    tempBuilder = addDriverCommandInstrumentation(tempBuilder);
    tempBuilder = addDriverSenderProxyInstrumentation(tempBuilder);
    tempBuilder = addDriverReceiverProxyInstrumentation(tempBuilder);
    tempBuilder = addDriverUdpChannelTransportInstrumentation(tempBuilder);
    tempBuilder = addDriverUntetheredSubscriptionInstrumentation(tempBuilder);

    return tempBuilder;
}
 

private static AgentBuilder addDriverReceiverProxyInstrumentation(final AgentBuilder agentBuilder)
{
    final boolean hasRegisterChannel = DRIVER_EVENT_CODES.contains(DriverEventCode.RECEIVE_CHANNEL_CREATION);
    final boolean hasCloseChannel = DRIVER_EVENT_CODES.contains(DriverEventCode.RECEIVE_CHANNEL_CLOSE);

    if (!hasRegisterChannel && !hasCloseChannel)
    {
        return agentBuilder;
    }

    return agentBuilder
        .type(nameEndsWith("ReceiverProxy"))
        .transform((builder, typeDescription, classLoader, javaModule) ->
        {
            if (hasRegisterChannel)
            {
                builder = builder
                    .visit(to(ChannelEndpointInterceptor.ReceiverProxy.RegisterReceiveChannelEndpoint.class)
                        .on(named("registerReceiveChannelEndpoint")));
            }
            if (hasCloseChannel)
            {
                builder = builder
                    .visit(to(ChannelEndpointInterceptor.ReceiverProxy.CloseReceiveChannelEndpoint.class)
                        .on(named("closeReceiveChannelEndpoint")));
            }

            return builder;
        });
}
 

public static synchronized void removeTransformer()
{
    if (logTransformer != null)
    {
        logTransformer.reset(instrumentation, AgentBuilder.RedefinitionStrategy.RETRANSFORMATION);
        instrumentation = null;
        logTransformer = null;
        thread = null;

        CloseHelper.close(readerAgentRunner);
        readerAgentRunner = null;
    }
}
 

@Subscribe
public void onStopModules(ReinstrumentationProtocol.StopModules stopEvent) {
    Logger.warn(() -> "Trying to stop modules.....");
    val stoppables = this.transformers.filter(KanelaFileTransformer::isStoppable)
                                      .map(KanelaFileTransformer::getClassFileTransformer)
                                      .map(transformer -> transformer.reset(this.instrumentation, AgentBuilder.RedefinitionStrategy.RETRANSFORMATION));

    if(stoppables.forAll(s -> s.equals(true))) Logger.warn(() -> "All modules are been stopped.");
    else Logger.warn(() -> "Error trying stop some modules.");
}
 

private AgentBuilder.Listener additionalListeners() {
    val listeners = new ArrayList<AgentBuilder.Listener>();
    if (config.getDump().isDumpEnabled()) listeners.add(ClassDumperListener.instance());
    if (config.getDebugMode()) listeners.add(DebugInstrumentationListener.instance());
    if (config.getInstrumentationRegistryConfig().isEnabled()) listeners.add(InstrumentationRegistryListener.instance());
    return new AgentBuilder.Listener.Compound(listeners);
}
 

private AgentBuilder withIgnore(AgentBuilder agentBuilder) {
    AgentBuilder.Ignored builder = agentBuilder.ignore(ignoreMatches())
            .or(moduleExcludes())
            .or(any(), isExtensionClassLoader())
            .or(any(), isKanelaClassLoader())
            .or(any(), isGroovyClassLoader())
            .or(any(), isSBTClassLoader())
            .or(any(), isSBTPluginClassLoader())
            .or(any(), isSBTCompilerClassLoader())
            .or(any(), isSBTCachedClassLoader())
            .or(any(), isReflectionClassLoader());

    if (moduleDescription.shouldInjectInBootstrap()) return builder;
    return builder.or(any(), isBootstrapClassLoader());
}
 

private AgentBuilder withBootstrapAttaching(AgentBuilder agentBuilder) {
    if(moduleDescription.shouldInjectInBootstrap()){
        Logger.info(() -> "Bootstrap Injection activated.");
        agentBuilder = agentBuilder.with(new InjectionStrategy.UsingUnsafe.OfFactory(ClassInjector.UsingUnsafe.Factory.resolve(instrumentation)));
    }
    return agentBuilder;
}
 

private AgentBuilder withRetransformationForRuntime(AgentBuilder agentBuilder) {
    if (config.isAttachedInRuntime() || moduleDescription.isStoppable() || moduleDescription.shouldInjectInBootstrap()) {
        Logger.info(() -> "Retransformation Strategy activated for: " + moduleDescription.getName());

        if(moduleDescription.isDisableClassFormatChanges())
            agentBuilder = agentBuilder.disableClassFormatChanges(); // enable restrictions imposed by most VMs and also HotSpot.

        agentBuilder = agentBuilder
            .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
            .withResubmission(PeriodicResubmitter.instance());
    }
    return agentBuilder;
}
 

AgentBuilder build() {
    return typeTransformations.build().foldLeft(newAgentBuilder(), (agent, typeTransformation) -> {
        val transformers = new ArrayList<AgentBuilder.Transformer>();
        transformers.addAll(typeTransformation.getBridges());
        transformers.addAll(typeTransformation.getMixins());
        transformers.addAll(typeTransformation.getTransformations());

        for (AgentBuilder.Transformer transformer : transformers) {
            agent  = agent
                    .type(typeTransformation.getElementMatcher().get(), RefinedClassLoaderMatcher.from(typeTransformation.getClassLoaderRefiner()))
                    .transform(transformer);
         }
         return agent;
    });
}
 

public AgentBuilder.Transformer makeTransformer() {
    val name = Option.of(advisorClassName).getOrElse(() -> advisorClass.getName());
    return new AgentBuilder.Transformer.ForAdvice()
            .advice(this.methodMatcher, name)
            .include(Thread.currentThread().getContextClassLoader())
            .withExceptionHandler(AdviceExceptionHandler.instance());
}
 

public AgentBuilder.Transformer makeTransformer() {
    return (builder, typeDescription, classLoader, module) -> {
        val interfaces = List.ofAll(Arrays.asList(mixinClass.getInterfaces()))
            .toSet()
            .map(TypeDescription.ForLoadedType::new)
            .toJavaList();

        return builder
            .implement(interfaces)
            .visit(MixinClassVisitorWrapper.of(this, typeDescription, classLoader));
    };

}
 

@SafeVarargs
static TypeTransformation of(String instrumentationName,
                             Option<ElementMatcher<? super TypeDescription>> elementMatcher,
                             Option<ClassLoaderRefiner> classLoaderRefiner,
                             List<AgentBuilder.Transformer> bridges,
                             List<AgentBuilder.Transformer> mixins,
                             List<AgentBuilder.Transformer>... transformers) {

    val transformations = Arrays.stream(transformers)
            .flatMap(Collection::stream)
            .collect(Collectors.toList());

    return new TypeTransformation(instrumentationName ,elementMatcher, classLoaderRefiner, bridges, mixins, transformations);
}