diff --git a/agent/src/main/java/reactor/blockhound/AllowAdvice.java b/agent/src/main/java/reactor/blockhound/AllowAdvice.java deleted file mode 100644 index 59ede09..0000000 --- a/agent/src/main/java/reactor/blockhound/AllowAdvice.java +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright (c) 2019-Present Pivotal Software Inc, All Rights Reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package reactor.blockhound; - -import net.bytebuddy.asm.Advice; - -import java.lang.annotation.Documented; -import java.lang.annotation.ElementType; -import java.lang.annotation.Retention; -import java.lang.annotation.RetentionPolicy; - -class AllowAdvice { - - @Documented - @Retention(RetentionPolicy.RUNTIME) - @java.lang.annotation.Target(ElementType.PARAMETER) - @interface AllowedArgument { - } - - @Advice.OnMethodEnter - static boolean onEnter(@AllowedArgument boolean allowed) { - Boolean previous = BlockHoundRuntime.IS_ALLOWED.get(); - if (previous == null || previous == allowed) { - return allowed; - } - BlockHoundRuntime.IS_ALLOWED.set(allowed); - return previous; - } - - @Advice.OnMethodExit(onThrowable = Throwable.class) - static void onExit(@Advice.Enter boolean wasAllowed, @AllowedArgument boolean allowed) { - if (wasAllowed != allowed) { - BlockHoundRuntime.IS_ALLOWED.set(wasAllowed); - } - } -} diff --git a/agent/src/main/java/reactor/blockhound/AllowancesByteBuddyTransformer.java b/agent/src/main/java/reactor/blockhound/AllowancesByteBuddyTransformer.java new file mode 100644 index 0000000..f80ea6e --- /dev/null +++ b/agent/src/main/java/reactor/blockhound/AllowancesByteBuddyTransformer.java @@ -0,0 +1,128 @@ +/* + * Copyright (c) 2019-Present Pivotal Software Inc, All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package reactor.blockhound; + +import net.bytebuddy.agent.builder.AgentBuilder; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.asm.AsmVisitorWrapper; +import net.bytebuddy.description.annotation.AnnotationDescription; +import net.bytebuddy.description.method.ParameterDescription; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.dynamic.DynamicType; +import net.bytebuddy.utility.JavaModule; + +import java.lang.annotation.Documented; +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.util.Map; + +/** + * This transformer applies {@link AllowAdvice} to every method + * registered with {@link BlockHound.Builder#allowBlockingCallsInside(String, String)}. + */ +class AllowancesByteBuddyTransformer implements AgentBuilder.Transformer { + + private Map> allowances; + + AllowancesByteBuddyTransformer(Map> allowances) { + this.allowances = allowances; + } + + @Override + public DynamicType.Builder transform( + DynamicType.Builder builder, + TypeDescription typeDescription, + ClassLoader classLoader, + JavaModule module + ) { + Map methods = allowances.get(typeDescription.getName()); + + if (methods == null) { + return builder; + } + + AsmVisitorWrapper advice = Advice + .withCustomMapping() + .bind(new AllowedArgument.Factory(methods)) + .to(AllowAdvice.class) + .on(method -> methods.containsKey(method.getName())); + + return builder.visit(advice); + } + + @Documented + @Retention(RetentionPolicy.RUNTIME) + @java.lang.annotation.Target(ElementType.PARAMETER) + @interface AllowedArgument { + + /** + * Binds advice method's argument annotated with {@link AllowedArgument} + * to boolean where `true` means "allowed" and `false" means "disallowed" + */ + class Factory implements Advice.OffsetMapping.Factory { + + final Map methods; + + Factory(Map methods) { + this.methods = methods; + } + + @Override + public Class getAnnotationType() { + return AllowedArgument.class; + } + + @Override + public Advice.OffsetMapping make( + ParameterDescription.InDefinedShape target, + AnnotationDescription.Loadable annotation, + AdviceType adviceType + ) { + return (instrumentedType, instrumentedMethod, assigner, argumentHandler, sort) -> { + boolean allowed = methods.get(instrumentedMethod.getName()); + return Advice.OffsetMapping.Target.ForStackManipulation.of(allowed); + }; + } + } + } + + static class AllowAdvice { + + @Advice.OnMethodEnter + static boolean onEnter( + @AllowancesByteBuddyTransformer.AllowedArgument boolean allowed + ) { + Boolean previous = BlockHoundRuntime.IS_ALLOWED.get(); + if (previous == null || previous == allowed) { + return allowed; + } + BlockHoundRuntime.IS_ALLOWED.set(allowed); + return previous; + } + + @Advice.OnMethodExit(onThrowable = Throwable.class) + static void onExit( + @Advice.Enter boolean wasAllowed, + @AllowancesByteBuddyTransformer.AllowedArgument boolean allowed + ) { + if (wasAllowed != allowed) { + BlockHoundRuntime.IS_ALLOWED.set(wasAllowed); + } + } + } +} diff --git a/agent/src/main/java/reactor/blockhound/BlockHound.java b/agent/src/main/java/reactor/blockhound/BlockHound.java index a3a7978..c289a05 100644 --- a/agent/src/main/java/reactor/blockhound/BlockHound.java +++ b/agent/src/main/java/reactor/blockhound/BlockHound.java @@ -25,13 +25,9 @@ import net.bytebuddy.agent.builder.AgentBuilder.RedefinitionStrategy.DiscoveryStrategy; import net.bytebuddy.agent.builder.AgentBuilder.TypeStrategy; import net.bytebuddy.asm.Advice; -import net.bytebuddy.asm.AsmVisitorWrapper; -import net.bytebuddy.description.type.TypeDescription; -import net.bytebuddy.dynamic.DynamicType; import net.bytebuddy.matcher.ElementMatchers; import net.bytebuddy.pool.TypePool; import net.bytebuddy.pool.TypePool.CacheProvider; -import net.bytebuddy.utility.JavaModule; import reactor.blockhound.integration.BlockHoundIntegration; import java.lang.instrument.ClassFileTransformer; @@ -221,6 +217,8 @@ public void install() { Instrumentation instrumentation = ByteBuddyAgent.install(); InstrumentationUtils.injectBootstrapClasses(instrumentation, BLOCK_HOUND_RUNTIME_TYPE.getInternalName()); + // Since BlockHoundRuntime is injected into the bootstrap classloader, + // we use raw Object[] here instead of `BlockingMethod` to avoid classloading issues BlockHoundRuntime.blockingMethodConsumer = args -> { String className = (String) args[0]; String methodName = (String) args[1]; @@ -239,13 +237,15 @@ public void install() { } } - private void instrument(Instrumentation instrumentation) throws Exception { + private void instrument(Instrumentation instrumentation) { ClassFileTransformer transformer = new NativeWrappingClassFileTransformer(blockingMethods); instrumentation.addTransformer(transformer, true); instrumentation.setNativeMethodPrefix(transformer, PREFIX); new AgentBuilder.Default() .with(RedefinitionStrategy.RETRANSFORMATION) + // Explicit strategy is almost 2 times faster than SinglePass + // TODO https://github.com/raphw/byte-buddy/issues/715 .with(new DiscoveryStrategy.Explicit( Stream .of(instrumentation.getAllLoadedClasses()) @@ -266,83 +266,33 @@ private void instrument(Instrumentation instrumentation) throws Exception { )) .with(TypeStrategy.Default.DECORATE) .with(InitializationStrategy.NoOp.INSTANCE) + // this DescriptionStrategy is required to force ByteBuddy to parse the bytes + // and not cache them, since we run another transformer (see NativeWrappingClassFileTransformer) + // before ByteBuddy .with(DescriptionStrategy.Default.POOL_FIRST) + // Override PoolStrategy because the default one will cache java.lang.Object, + // and we need to instrument it. .with((PoolStrategy) (classFileLocator, classLoader) -> new TypePool.Default( new CacheProvider.Simple(), classFileLocator, TypePool.Default.ReaderMode.FAST )) .with(AgentBuilder.Listener.StreamWriting.toSystemError().withErrorsOnly()) + // Do not ignore JDK classes .ignore(ElementMatchers.none()) - .type((typeDescription, classLoader, module, classBeingRedefined, protectionDomain) -> { - if (blockingMethods.containsKey(typeDescription.getInternalName())) { - return true; - } - if (allowances.containsKey(typeDescription.getName())) { - return true; - } - return false; - }) - .transform(new BlockingCallsTransformer()) - .transform(new AllowancesTransformer()) - .installOn(instrumentation); - } - - private class BlockingCallsTransformer implements AgentBuilder.Transformer { - @Override - public DynamicType.Builder transform( - DynamicType.Builder builder, - TypeDescription typeDescription, - ClassLoader classLoader, - JavaModule module - ) { - Map> methods = blockingMethods.get(typeDescription.getInternalName()); + // Instrument blocking calls + .type(it -> blockingMethods.containsKey(it.getInternalName())) + .transform(new BlockingCallsByteBuddyTransformer(blockingMethods)) + .asTerminalTransformation() - if (methods == null) { - return builder; - } - - AsmVisitorWrapper advice = Advice.withCustomMapping() - .bind(BlockingCallAdvice.ModifiersArgument.class, (type, method, assigner, argumentHandler, sort) -> { - return Advice.OffsetMapping.Target.ForStackManipulation.of(method.getModifiers()); - }) - .to(BlockingCallAdvice.class) - .on(method -> { - Set descriptors = methods.get(method.getName()); - return descriptors != null && descriptors.contains(method.getDescriptor()); - }); - - return builder.visit(advice); - } - } + // Instrument allowed/disallowed methods + .type(it -> allowances.containsKey(it.getName())) + .transform(new AllowancesByteBuddyTransformer(allowances)) + .asTerminalTransformation() - private class AllowancesTransformer implements AgentBuilder.Transformer { - - @Override - public DynamicType.Builder transform( - DynamicType.Builder builder, - TypeDescription typeDescription, - ClassLoader classLoader, - JavaModule module - ) { - Map methods = allowances.get(typeDescription.getName()); - - if (methods == null) { - return builder; - } - - AsmVisitorWrapper advice = Advice - .withCustomMapping() - .bind(AllowAdvice.AllowedArgument.class, (type, method, assigner, argumentHandler, sort) -> { - return Advice.OffsetMapping.Target.ForStackManipulation.of(methods.get(method.getName())); - }) - .to(AllowAdvice.class) - .on(method -> methods.containsKey(method.getName())); - - return builder.visit(advice); - } + .installOn(instrumentation); } } } diff --git a/agent/src/main/java/reactor/blockhound/BlockingCallAdvice.java b/agent/src/main/java/reactor/blockhound/BlockingCallAdvice.java deleted file mode 100644 index cfc9ba0..0000000 --- a/agent/src/main/java/reactor/blockhound/BlockingCallAdvice.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * Copyright (c) 2019-Present Pivotal Software Inc, All Rights Reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package reactor.blockhound; - -import net.bytebuddy.asm.Advice; - -import java.lang.annotation.Documented; -import java.lang.annotation.ElementType; -import java.lang.annotation.Retention; -import java.lang.annotation.RetentionPolicy; - -class BlockingCallAdvice { - - @Documented - @Retention(RetentionPolicy.RUNTIME) - @java.lang.annotation.Target(ElementType.PARAMETER) - @interface ModifiersArgument { - } - - @Advice.OnMethodEnter - static void onEnter( - @Advice.Origin("#t") String declaringType, - @Advice.Origin("#m") String methodName, - @BlockingCallAdvice.ModifiersArgument int modifiers - ) { - BlockHoundRuntime.checkBlocking(declaringType, methodName, modifiers); - } -} diff --git a/agent/src/main/java/reactor/blockhound/BlockingCallsByteBuddyTransformer.java b/agent/src/main/java/reactor/blockhound/BlockingCallsByteBuddyTransformer.java new file mode 100644 index 0000000..e1579d7 --- /dev/null +++ b/agent/src/main/java/reactor/blockhound/BlockingCallsByteBuddyTransformer.java @@ -0,0 +1,113 @@ +/* + * Copyright (c) 2019-Present Pivotal Software Inc, All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package reactor.blockhound; + +import net.bytebuddy.agent.builder.AgentBuilder; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.asm.AsmVisitorWrapper; +import net.bytebuddy.description.annotation.AnnotationDescription; +import net.bytebuddy.description.method.ParameterDescription; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.dynamic.DynamicType; +import net.bytebuddy.utility.JavaModule; + +import java.lang.annotation.Documented; +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.util.Map; +import java.util.Set; + +/** + * This transformer applies {@link BlockingCallAdvice} to every method + * registered with {@link BlockHound.Builder#markAsBlocking(Class, String, String)}. + */ +class BlockingCallsByteBuddyTransformer implements AgentBuilder.Transformer { + + private Map>> blockingMethods; + + BlockingCallsByteBuddyTransformer(Map>> blockingMethods) { + this.blockingMethods = blockingMethods; + } + + @Override + public DynamicType.Builder transform( + DynamicType.Builder builder, + TypeDescription typeDescription, + ClassLoader classLoader, + JavaModule module + ) { + Map> methods = blockingMethods.get(typeDescription.getInternalName()); + + if (methods == null) { + return builder; + } + + AsmVisitorWrapper advice = Advice.withCustomMapping() + .bind(ModifiersArgument.Factory.INSTANCE) + .to(BlockingCallAdvice.class) + .on(method -> { + Set descriptors = methods.get(method.getName()); + return descriptors != null && descriptors.contains(method.getDescriptor()); + }); + + return builder.visit(advice); + } + + @Documented + @Retention(RetentionPolicy.RUNTIME) + @java.lang.annotation.Target(ElementType.PARAMETER) + @interface ModifiersArgument { + + /** + * Binds advice method's argument annotated with {@link ModifiersArgument} + * to method's modifiers (static, final, private, etc) + */ + enum Factory implements Advice.OffsetMapping.Factory { + INSTANCE; + + @Override + public Class getAnnotationType() { + return ModifiersArgument.class; + } + + @Override + public Advice.OffsetMapping make( + ParameterDescription.InDefinedShape target, + AnnotationDescription.Loadable annotation, + AdviceType adviceType + ) { + return (instrumentedType, instrumentedMethod, assigner, argumentHandler, sort) -> { + int modifiers = instrumentedMethod.getModifiers(); + return Advice.OffsetMapping.Target.ForStackManipulation.of(modifiers); + }; + } + } + } + + static class BlockingCallAdvice { + + @Advice.OnMethodEnter + static void onEnter( + @Advice.Origin("#t") String declaringType, + @Advice.Origin("#m") String methodName, + @BlockingCallsByteBuddyTransformer.ModifiersArgument int modifiers + ) { + BlockHoundRuntime.checkBlocking(declaringType, methodName, modifiers); + } + } +}