Skip to content

Commit

Permalink
Comment and clean-up some code (#47)
Browse files Browse the repository at this point in the history
1. Add `BlockingCallsByteBuddyTransformer` and move the advice in it
2. Add `AllowancesByteBuddyTransformer` and move the advice
3. Document ByteBuddy customizations
  • Loading branch information
bsideup committed Aug 29, 2019
1 parent eb337c4 commit 312fc5f
Show file tree
Hide file tree
Showing 5 changed files with 261 additions and 162 deletions.
50 changes: 0 additions & 50 deletions agent/src/main/java/reactor/blockhound/AllowAdvice.java

This file was deleted.

Original file line number Diff line number Diff line change
@@ -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<String, Map<String, Boolean>> allowances;

AllowancesByteBuddyTransformer(Map<String, Map<String, Boolean>> allowances) {
this.allowances = allowances;
}

@Override
public DynamicType.Builder<?> transform(
DynamicType.Builder<?> builder,
TypeDescription typeDescription,
ClassLoader classLoader,
JavaModule module
) {
Map<String, Boolean> 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<AllowedArgument> {

final Map<String, Boolean> methods;

Factory(Map<String, Boolean> methods) {
this.methods = methods;
}

@Override
public Class<AllowedArgument> getAnnotationType() {
return AllowedArgument.class;
}

@Override
public Advice.OffsetMapping make(
ParameterDescription.InDefinedShape target,
AnnotationDescription.Loadable<AllowedArgument> 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);
}
}
}
}
90 changes: 20 additions & 70 deletions agent/src/main/java/reactor/blockhound/BlockHound.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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];
Expand All @@ -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())
Expand All @@ -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<String, Set<String>> 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<String> 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<String, Boolean> 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);
}
}
}
42 changes: 0 additions & 42 deletions agent/src/main/java/reactor/blockhound/BlockingCallAdvice.java

This file was deleted.

Loading

0 comments on commit 312fc5f

Please sign in to comment.