hide BlockHoundRuntime from public API (#53)

`BlockHoundRuntime` must be public because it is injected into bootstrap ClassLoader and being accessed from the JDK classes later. But we don't want the users to call it, since it should only be used from the instrumentation. This change makes the original class internal and changes the visibility with ASM when we inject it into the bootstrap ClassLoader.
reactor · Oct 1, 2019 · b121197 · b121197
1 parent b2afb3b
commit b121197
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 13 deletions.
diff --git a/agent/src/main/java/reactor/blockhound/BlockHoundRuntime.java b/agent/src/main/java/reactor/blockhound/BlockHoundRuntime.java
@@ -22,15 +22,15 @@
 // Warning!!! This class MUST NOT be loaded by any classloader other than the bootstrap one.
 // Otherwise, non-bootstrap classes will be referring to it, but only the bootstrap one gets
 // initialized.
-public class BlockHoundRuntime {
+class BlockHoundRuntime {
 
     @SuppressWarnings("unused")
-    public static volatile Consumer<Object[]> blockingMethodConsumer;
+    static volatile Consumer<Object[]> blockingMethodConsumer;
 
     @SuppressWarnings("unused")
-    public static volatile Predicate<Thread> threadPredicate;
+    static volatile Predicate<Thread> threadPredicate;
 
-    public static final ThreadLocal<Boolean> IS_ALLOWED = ThreadLocal.withInitial(() -> {
+    static final ThreadLocal<Boolean> IS_ALLOWED = ThreadLocal.withInitial(() -> {
         if (threadPredicate.test(Thread.currentThread())) {
             return false;
         }
@@ -41,7 +41,7 @@ public class BlockHoundRuntime {
     });
 
     @SuppressWarnings("unused")
-    public static void checkBlocking(String internalClassName, String methodName, int modifiers) {
+    static void checkBlocking(String internalClassName, String methodName, int modifiers) {
         if (Boolean.FALSE == IS_ALLOWED.get()) {
             blockingMethodConsumer.accept(new Object[] {
                     internalClassName.replace("/", "."),

diff --git a/agent/src/main/java/reactor/blockhound/InstrumentationUtils.java b/agent/src/main/java/reactor/blockhound/InstrumentationUtils.java
@@ -16,6 +16,13 @@
 
 package reactor.blockhound;
 
+import net.bytebuddy.jar.asm.ClassReader;
+import net.bytebuddy.jar.asm.ClassVisitor;
+import net.bytebuddy.jar.asm.ClassWriter;
+import net.bytebuddy.jar.asm.FieldVisitor;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
+
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
@@ -35,19 +42,59 @@ static void injectBootstrapClasses(Instrumentation instrumentation, String... cl
         try (ZipOutputStream zipOutputStream = new ZipOutputStream(new FileOutputStream(tempJarFile))) {
             for (String className : classNames) {
                 String classFile = className.replace(".", "/") + ".class";
-                InputStream inputStream = classLoader.getResourceAsStream(classFile);
-                ZipEntry e = new ZipEntry(classFile);
-                zipOutputStream.putNextEntry(e);
-
-                byte[] buf = new byte[4096];
-                int n;
-                while ((n = inputStream.read(buf)) > 0) {
-                    zipOutputStream.write(buf, 0, n);
+                try (InputStream inputStream = classLoader.getResourceAsStream(classFile)) {
+                    ZipEntry entry = new ZipEntry(classFile);
+                    zipOutputStream.putNextEntry(entry);
+
+                    ClassReader cr = new ClassReader(inputStream);
+                    ClassWriter cw = new ClassWriter(cr, 0);
+
+                    cr.accept(new MakePublicClassVisitor(cw), 0);
+
+                    zipOutputStream.write(cw.toByteArray());
                 }
 
                 zipOutputStream.closeEntry();
             }
         }
         instrumentation.appendToBootstrapClassLoaderSearch(new JarFile(tempJarFile));
     }
+
+    /**
+     * Makes the class, fields and methods public
+     */
+    static class MakePublicClassVisitor extends ClassVisitor {
+
+        MakePublicClassVisitor(ClassWriter cw) {
+            super(Opcodes.ASM7, cw);
+        }
+
+        @Override
+        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
+            super.visit(version, access | Opcodes.ACC_PUBLIC, name, signature, superName, interfaces);
+        }
+
+        @Override
+        public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
+            switch (name) {
+                case "blockingMethodConsumer":
+                case "threadPredicate":
+                case "IS_ALLOWED":
+                    access = access | Opcodes.ACC_PUBLIC;
+                    break;
+            }
+            return super.visitField(access, name, descriptor, signature, value);
+        }
+
+        @Override
+        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+            switch (name) {
+                case "checkBlocking":
+                    access = access | Opcodes.ACC_PUBLIC;
+                    break;
+            }
+
+            return super.visitMethod(access, name, descriptor, signature, exceptions);
+        }
+    }
 }