Fix file path formatting bug in Injector.java

README.md

@@ -31,16 +31,24 @@ methods and fields.
 
 * Inject fields
 
+* Handle exceptions
+
+* Inject try-catch-finally
+
+* Inject assertions
+
+*A caveat regarding assert-statements: the compiler synthesizes a static final field named `$assertionsDisabled`, so if a target
+class declares a static field with this name and does not declare any assertions in its code, loading of the field may already
+be done in static initialization or field declaration, preventing assertions from functioning as intended for injected code.*
+
 ### TODO
 
 * Better tests
 
-* Execution path optimization (e.g. remove unnecessary [GOTO](https://docs.oracle.com/javase/specs/jvms/se8/html/jvms-6.html#jvms-6.5.goto)s)
+* ~~Implement subroutines ([JSR](https://docs.oracle.com/javase/specs/jvms/se8/html/jvms-6.html#jvms-6.5.jsr) / [RET](https://docs.oracle.com/javase/specs/jvms/se8/html/jvms-6.html#jvms-6.5.ret))~~
 
-* Implement exceptions
-
-* Implement try-catch
-  * Implement subroutines ([JSR](https://docs.oracle.com/javase/specs/jvms/se8/html/jvms-6.html#jvms-6.5.jsr) / [RET](https://docs.oracle.com/javase/specs/jvms/se8/html/jvms-6.html#jvms-6.5.ret))
+*Note: Subroutines are prohibited as of class file version 51, so unless this project sees much use with projects targeting Java
+version 1.6 or lower on very specific compilers, this implementation can be set aside until everything else is polished.*
 
 ## How do I use this?
 
@@ -90,6 +98,10 @@ public class TweakFoo {
 }
 ```
 
+If this looks a bit confusing: don't worry, [it gets simpler](#a-simple-example) after we remove all the things that can
+be inferred automatically by the library. We start with a more complicated example to show explicitly what's happening
+under the hood in case it's needed in more complex use-cases.
+
 We specify using the `@Inject` annotation that we would like to inject the annotated method into the targeted class (
 we'll get to how we target a class in a sec), that we would like to append the code `AFTER` the existing code, that we
 are targeting a method named `addMyNumber` which accepts an int (specified by `(I)`) and returns an int (specified by
@@ -143,12 +155,11 @@ To target a class, simply create an instance of the `dev.w1zzrd.asm.Combine` cla
 annotated MethodNodes. For example, using the example code described earlier, we could do as follows:
 
 ```java
-import dev.w1zzrd.asm.Combine;
-import dev.w1zzrd.asm.GraftSource;
-import dev.w1zzrd.asm.Loader;
+import dev.w1zzrd.asm.*;
+import jdk.internal.org.objectweb.asm.tree.MethodNode;
 
 public class Run {
-  public static void main(String[] args) {
+  public static void main(String[] args) throws Exception {
     // This is the class we would like to inject code into
     Combine target = new Combine(Loader.getClassNode("Foo"));
 
@@ -184,7 +195,7 @@ available to the default ClassLoader's classpath. In this case, we can simply do
 import dev.w1zzrd.asm.Injector;
 
 public class Run {
-  public static void main(String[] args) {
+  public static void main(String[] args) throws Exception {
       // Locates all necessary tweaks and injects them into Foo
     Injector.injectAll("Foo").compile();
 
@@ -200,9 +211,7 @@ The only caveat is that classes injected this way must have an `@InjectClass` an
 In our example, this would mean that the `TweakFoo` class would look as follows:
 
 ```java
-import dev.w1zzrd.asm.InPlaceInjection;
-import dev.w1zzrd.asm.Inject;
-import dev.w1zzrd.asm.InjectClass;
+import dev.w1zzrd.asm.*;
 
 // This marks the class as targeting Foo
 @InjectClass(Foo.class)
@@ -221,11 +230,9 @@ method, the target should never be ambiguous. In fact, the resolution is intelli
 is unambiguous for the targeted method name in the targeted class, the `acceptOriginalReturn` value can even be omitted.
 This means that a minimal example implementation of the `TweakFoo` class could look as follows:
 
+#### A simple example:
 ```java
-import dev.w1zzrd.asm.InPlaceInjection;
-import dev.w1zzrd.asm.Inject;
-import dev.w1zzrd.asm.InjectClass;
-
+import dev.w1zzrd.asm.*;
 import static dev.w1zzrd.asm.InPlaceInjection.AFTER;
 
 @InjectClass(Foo.class)

src/dev/w1zzrd/asm/Combine.java

@@ -11,8 +11,6 @@ import jdk.internal.org.objectweb.asm.Handle;
 import jdk.internal.org.objectweb.asm.Opcodes;
 import jdk.internal.org.objectweb.asm.Type;
 import jdk.internal.org.objectweb.asm.tree.*;
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.*;
@@ -21,6 +19,10 @@ import java.util.stream.Collectors;
 import static jdk.internal.org.objectweb.asm.ClassWriter.COMPUTE_MAXS;
 
 public class Combine {
+    public static final String VAR_ASSERT_NAME = "$assertionsDisabled";
+    public static final int VAR_ASSERT_FLAGS = Opcodes.ACC_STATIC | Opcodes.ACC_SYNTHETIC | Opcodes.ACC_FINAL;
+
+
     private final ArrayList<DynamicSourceUnit> graftSources = new ArrayList<>();
 
     private final ClassNode target;
@@ -134,6 +136,10 @@ public class Combine {
         // Recompute maximum stack size
         resolution.node.maxStack = Math.max(resolution.node.maxStack, extension.maxStack);
 
+        // Merge try-catch blocks
+        resolution.node.tryCatchBlocks.addAll(extension.tryCatchBlocks);
+        // Exception list not merged to maintain original signature
+
         finishGrafting(extension, source);
     }
 
@@ -155,6 +161,9 @@ public class Combine {
         for (int i = 0; i < sig.getArgCount(); ++i)
             adjustArgument(target, getVarAt(target.localVariables, i), true, false);
 
+        target.tryCatchBlocks.addAll(extension.tryCatchBlocks);
+        // Exception list not merged to maintain original signature
+
         finishGrafting(extension, source);
     }
 
@@ -318,6 +327,88 @@ public class Combine {
         return target;
     }
 
+    protected void ensureLoadClassAssertionState() {
+        if (!hasDeclaredAssertionState())
+            target.fields.add(new FieldNode(
+                    VAR_ASSERT_FLAGS,
+                    VAR_ASSERT_NAME,
+                    "Z",
+                    null,
+                    null
+            ));
+
+        // Check if state is loaded
+        if (target.methods.stream().noneMatch(it -> it.name.equals("<clinit>"))) {
+            MethodNode mnode = new MethodNode(Opcodes.ACC_STATIC, "<clinit>", "()V", null, null);
+            injectAssertionLoad(mnode, true);
+            target.methods.add(mnode);
+        } else {
+            MethodNode clinit = target.methods.stream().filter(it -> it.name.equals("<clinit>")).findAny().get();
+
+            for (AbstractInsnNode node = clinit.instructions.getFirst(); node != null; node = node.getNext())
+                if (node instanceof FieldInsnNode &&
+                    node.getOpcode() == Opcodes.PUTSTATIC &&
+                    VAR_ASSERT_NAME.equals(((FieldInsnNode) node).name))
+                    return;
+
+            // Assertion state not loaded in the current clinit. Add it to the start of the clinit
+            injectAssertionLoad(clinit, false);
+        }
+    }
+
+    /**
+     * For internal use. Adds instructions to the given method to load assertion flag to a static final field
+     * @param node Method to inject instructions into
+     */
+    private void injectAssertionLoad(MethodNode node, boolean insertReturn) {
+        if (node.instructions == null) {
+            node.instructions = new InsnList();
+            insertReturn = true;
+        }
+
+        AbstractInsnNode current;
+        if (node.instructions.getFirst() == null || !(node.instructions.getFirst() instanceof LabelNode))
+            node.instructions.insert(current = new LabelNode());
+        else
+            current = node.instructions.getFirst();
+
+        node.instructions.insert(current, current = new LdcInsnNode(Type.getType("L"+target.name+";")));
+        node.instructions.insert(current, current = new MethodInsnNode(
+                Opcodes.INVOKEVIRTUAL,
+                "java/lang/Class",
+                "desiredAssertionStatus",
+                "()Z",
+                false
+        ));
+
+        final LabelNode jumpNE = new LabelNode();
+        final LabelNode jumpGOTO = new LabelNode();
+
+        node.instructions.insert(current, current = new JumpInsnNode(Opcodes.IFNE, jumpNE));
+        node.instructions.insert(current, current = new InsnNode(Opcodes.ICONST_1));
+        node.instructions.insert(current, current = new JumpInsnNode(Opcodes.GOTO, jumpGOTO));
+        node.instructions.insert(current, current = jumpNE);
+        node.instructions.insert(current, current = new FrameNode(Opcodes.F_SAME, 0, new Object[0], 0, new Object[0]));
+        node.instructions.insert(current, current = new InsnNode(Opcodes.ICONST_0));
+        node.instructions.insert(current, current = jumpGOTO);
+        node.instructions.insert(current, current = new FrameNode(
+                Opcodes.F_SAME1,
+                0,
+                new Object[0],
+                1,
+                new Object[]{ Opcodes.INTEGER }
+        ));
+
+        node.instructions.insert(current, current = new FieldInsnNode(Opcodes.PUTSTATIC, target.name, VAR_ASSERT_NAME, "Z"));
+
+        if (insertReturn)
+            node.instructions.insert(current, new InsnNode(Opcodes.RETURN));
+    }
+
+    protected boolean hasDeclaredAssertionState() {
+        return target.fields.stream().anyMatch(it -> VAR_ASSERT_NAME.equals(it.name));
+    }
+
     /**
      * Prepares a {@link MethodNode} for grafting on to a given method and into the targeted  {@link ClassNode}
      * @param node Node to adapt
@@ -329,8 +420,14 @@ public class Combine {
             if (insn instanceof MethodInsnNode) insn = adaptMethodInsn((MethodInsnNode) insn, source, node);
             else if (insn instanceof LdcInsnNode) adaptLdcInsn((LdcInsnNode) insn, source.getTypeName());
             else if (insn instanceof FrameNode) adaptFrameNode((FrameNode) insn, source);
-            else if (insn instanceof FieldInsnNode) adaptFieldInsn((FieldInsnNode) insn, source);
             else if (insn instanceof InvokeDynamicInsnNode) adaptInvokeDynamicInsn((InvokeDynamicInsnNode) insn, source);
+            else if (insn instanceof FieldInsnNode) {
+                adaptFieldInsn((FieldInsnNode) insn, source);
+
+                // If a method is trying to access the assertion state of the class, ensure the target class declares the state
+                if (VAR_ASSERT_NAME.equals(((FieldInsnNode) insn).name) && insn.getOpcode() == Opcodes.GETSTATIC)
+                    ensureLoadClassAssertionState();
+            }
         }
 
         // Adapt variable types
@@ -374,8 +471,13 @@ public class Combine {
     }
 
     private void storeAndGotoFromReturn(MethodNode source, InsnList nodes, int storeIndex, MethodSignature sig) {
+        // The last (injected) GOTO can always be removed
+        AbstractInsnNode lastGoto = null;
+        int jumpCount = 0;
+
         // If we already have a final frame, there's no need to add one
-        LabelNode endLabel = hasEndJumpFrame(nodes) ? findOrMakeEndLabel(nodes) : makeEndJumpFrame(nodes, sig, source);
+        boolean hadEJF = hasEndJumpFrame(nodes);
+        LabelNode endLabel = hadEJF ? findOrMakeEndLabel(nodes) : makeEndJumpFrame(nodes, sig, source);
 
         INSTRUCTION_LOOP:
         for (AbstractInsnNode current = nodes.getFirst(); current != null; current = current.getNext()) {
@@ -401,19 +503,33 @@ public class Combine {
                     break;
 
                 case Opcodes.RETURN:
-                    nodes.set(current, current = new JumpInsnNode(Opcodes.GOTO, endLabel));
+                    nodes.set(current, lastGoto = current = new JumpInsnNode(Opcodes.GOTO, endLabel));
+                    ++jumpCount;
                     // Fallthrough
 
                 default:
                     continue INSTRUCTION_LOOP;
             }
-            nodes.insert(current, current = new JumpInsnNode(Opcodes.GOTO, endLabel));
+            nodes.insert(current, lastGoto = current = new JumpInsnNode(Opcodes.GOTO, endLabel));
+            ++jumpCount;
+        }
+
+        if (lastGoto != null) {
+            nodes.remove(lastGoto);
+
+            // The final jump frame and label can be removed
+            if (jumpCount == 1 && !hadEJF)
+                nodes.remove(endLabel.getNext());
         }
     }
 
     private void popAndGotoFromReturn(MethodNode source, InsnList nodes, MethodSignature sig) {
+        AbstractInsnNode lastGoto = null;
+        int jumpCount = 0;
+
         // If we already have a final frame, there's no need to add one
-        LabelNode endLabel = hasEndJumpFrame(nodes) ? findOrMakeEndLabel(nodes) : makeEndJumpFrame(nodes, sig, source);
+        boolean hadEJF = hasEndJumpFrame(nodes);
+        LabelNode endLabel = hadEJF ? findOrMakeEndLabel(nodes) : makeEndJumpFrame(nodes, sig, source);
 
         INSTRUCTION_LOOP:
         for (AbstractInsnNode current = nodes.getFirst(); current != null; current = current.getNext()) {
@@ -430,14 +546,25 @@ public class Combine {
                     break;
 
                 case Opcodes.RETURN:
-                    nodes.set(current, current = new JumpInsnNode(Opcodes.GOTO, endLabel));
+                    nodes.set(current, lastGoto = current = new JumpInsnNode(Opcodes.GOTO, endLabel));
+                    ++jumpCount;
                     // Fallthrough
 
                 default:
                     continue INSTRUCTION_LOOP;
             }
 
-            nodes.insert(current, current = new JumpInsnNode(Opcodes.GOTO, endLabel));
+            nodes.insert(current, lastGoto = current = new JumpInsnNode(Opcodes.GOTO, endLabel));
+            ++jumpCount;
+        }
+
+
+        if (lastGoto != null) {
+            nodes.remove(lastGoto);
+
+            // The final jump frame and label can be removed
+            if (jumpCount == 1 && !hadEJF)
+                nodes.remove(endLabel.getNext());
         }
     }
 
@@ -750,7 +877,7 @@ public class Combine {
      * @return The located method node
      * @throws MethodNodeResolutionException If no method node matching the given description could be found
      */
-    protected final @NotNull MethodNode checkMethodExists(String name, MethodSignature descriptor) {
+    protected final MethodNode checkMethodExists(String name, MethodSignature descriptor) {
         final MethodNode target = findMethodNode(name, descriptor);
 
         if (target == null)
@@ -768,7 +895,7 @@ public class Combine {
      * @param desc Descriptor of the method node to find
      * @return A matching {@link MethodNode} if one exists, else null
      */
-    protected @Nullable MethodNode findMethodNode(String name, MethodSignature desc) {
+    protected MethodNode findMethodNode(String name, MethodSignature desc) {
         return target.methods
                 .stream()
                 .filter(it -> it.name.equals(name) && new MethodSignature(it.desc).equals(desc))
@@ -838,7 +965,7 @@ public class Combine {
     }
 
 
-    protected static @Nullable LabelNode findLabelBeforeReturn(AbstractInsnNode start, INodeTraversal traverse) {
+    protected static LabelNode findLabelBeforeReturn(AbstractInsnNode start, INodeTraversal traverse) {
         for (AbstractInsnNode cur = start; cur != null; cur = traverse.traverse(cur))
             if (cur instanceof LabelNode) // Traversal hit label
                 return (LabelNode) cur;

src/dev/w1zzrd/asm/GraftSource.java

@@ -6,8 +6,6 @@ import jdk.internal.org.objectweb.asm.tree.AnnotationNode;
 import jdk.internal.org.objectweb.asm.tree.ClassNode;
 import jdk.internal.org.objectweb.asm.tree.FieldNode;
 import jdk.internal.org.objectweb.asm.tree.MethodNode;
-import org.jetbrains.annotations.Nullable;
-
 import java.util.*;
 import java.util.stream.Collectors;
 
@@ -91,7 +89,7 @@ public final class GraftSource {
         return getInjectedMethod(name, desc) != null;
     }
 
-    public @Nullable MethodNode getInjectedMethod(String name, String desc) {
+    public MethodNode getInjectedMethod(String name, String desc) {
         return methodAnnotations
                 .entrySet()
                 .stream()

src/dev/w1zzrd/asm/Injector.java

@@ -9,6 +9,8 @@ import java.io.File;
 import java.io.IOException;
 import java.net.MalformedURLException;
 import java.net.URL;
+import java.net.URLDecoder;
+import java.nio.charset.StandardCharsets;
 import java.util.Enumeration;
 import java.util.Objects;
 
@@ -25,7 +27,10 @@ public class Injector {
     public static void injectAll(ClassLoader loader, Combine merger) throws IOException {
         Enumeration<URL> resources = loader.getResources("");
         while (resources.hasMoreElements())
-            injectDirectory(new File(resources.nextElement().getPath()), merger);
+            injectDirectory(new File(URLDecoder.decode(
+                    resources.nextElement().getFile(),
+                    StandardCharsets.UTF_8.name())
+            ), merger);
     }
 
     public static void injectAll(Combine merger) throws IOException {

src/dev/w1zzrd/asm/analysis/FrameState.java

@@ -7,9 +7,6 @@ import jdk.internal.org.objectweb.asm.Label;
 import jdk.internal.org.objectweb.asm.Opcodes;
 import jdk.internal.org.objectweb.asm.Type;
 import jdk.internal.org.objectweb.asm.tree.*;
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
-import sun.reflect.generics.reflectiveObjects.NotImplementedException;
 import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.List;
@@ -192,7 +189,7 @@ public class FrameState {
      * @param node {@link LabelNode} to find jumps referencing
      * @return A jump instruction node earlier in the instruction list or null if none could be found
      */
-    private static @Nullable JumpInsnNode findEarliestJump(LabelNode node) {
+    private static JumpInsnNode findEarliestJump(LabelNode node) {
         JumpInsnNode jump = null;
 
         // Traverse backward until we hit the beginning of the list
@@ -643,7 +640,7 @@ public class FrameState {
     private static List<String> getOpsByComplexity(
             boolean complexPush,
             boolean complexPop,
-            @Nullable Predicate<Integer> insnP
+            Predicate<Integer> insnP
     ) {
         ArrayList<Integer> opcodes = new ArrayList<>();
 
@@ -687,9 +684,9 @@ public class FrameState {
      * @return Negative values for instructions previous to the current instruction, positive values for instructions
      * after the current instruction. Null if instruction could not be found
      */
-    private static @Nullable Integer relativeIndexOf(
-            @NotNull AbstractInsnNode current,
-            @NotNull AbstractInsnNode find
+    private static Integer relativeIndexOf(
+            AbstractInsnNode current,
+            AbstractInsnNode find
     ) {
         // Check backward
         int idx = 0;

src/dev/w1zzrd/asm/signature/TypeSignature.java

@@ -2,8 +2,6 @@ package dev.w1zzrd.asm.signature;
 
 import dev.w1zzrd.asm.exception.SignatureInstanceMismatchException;
 import dev.w1zzrd.asm.exception.TypeSignatureParseException;
-import org.jetbrains.annotations.Nullable;
-
 import java.util.Arrays;
 import java.util.Objects;
 
@@ -86,7 +84,7 @@ public class TypeSignature {
      * @param primitive Primitive type internal name (V, J or D for Top types)
      * @param isTop Whether or not this is a Top type (only valid for 64-bit types J and D or as delimiter type V)
      */
-    public TypeSignature(@Nullable Character primitive, boolean isTop) {
+    public TypeSignature(Character primitive, boolean isTop) {
         if (primitive != null) {
             switch (Character.toUpperCase(primitive)) {
                 case 'J':

test/MergeInject.java

@@ -22,6 +22,8 @@ public class MergeInject extends MergeTest implements Runnable {
         Directives.callSuper();
         s = "Hello";
         number = 10;
+
+        assert false : "Test";
     }
 
 
@@ -45,8 +47,8 @@ public class MergeInject extends MergeTest implements Runnable {
     @Inject(AFTER)
     public int stackTest(int arg) {
         Runnable r = () -> {
-          System.out.println(arg / 15);
-          System.out.println("Heyo");
+            System.out.println(arg / 15);
+            System.out.println("Heyo");
         };
         r.run();
         return 69;
@@ -54,10 +56,20 @@ public class MergeInject extends MergeTest implements Runnable {
 
 
     @Inject(AFTER)
-    public String test(String retVal){
+    public String test(String retVal) throws Exception {
 
         System.out.println(retVal + "Cringe");
 
+        try {
+            if (ThreadLocalRandom.current().nextBoolean())
+                throw new Exception("Hello from exception");
+        }catch (Exception e) {
+            System.out.println("Hello from catch");
+            e.printStackTrace();
+        } finally {
+            System.out.println("Hello from finally");
+        }
+
         return "Modified";
     }
 
@@ -70,4 +82,4 @@ public class MergeInject extends MergeTest implements Runnable {
             System.out.println(test()+'\n');
         }
     }
-}
+}

test/MergeTest.java

@@ -14,7 +14,7 @@ public class MergeTest {
     }
 
     public String test(){
-        Class<?> c = Combine.class;
+        final Class<?> c = Combine.class;
         Runnable r = () -> {
             System.out.println("Sick");
             System.out.println(c.getName());
@@ -57,4 +57,4 @@ public class MergeTest {
         if (bool)
             System.out.println(k + a + b * getNumber() + i);
     }
-}
+}

test/Test.java

@@ -1,10 +1,16 @@
+import dev.w1zzrd.asm.Combine;
 import dev.w1zzrd.asm.Injector;
+
+import java.io.File;
+import java.io.FileOutputStream;
 import java.io.IOException;
 
 public class Test {
     public static void main(String... args) throws IOException {
+        ClassLoader.getSystemClassLoader().setDefaultAssertionStatus(false);
+
         // Load target class, inject all annotated classes and load compiled bytecode into JVM
-        Injector.injectAll("MergeTest").compile();
+        dumpFile(Injector.injectAll("MergeTest"), "MergeTest").compile();
 
         // Run simple injection tests
         new MergeTest().test();
@@ -17,4 +23,21 @@ public class Test {
         r.run();
     }
 
+    public static Combine dumpFile(Combine comb, String name) {
+        File f = new File(name + ".class");
+        try {
+            if ((f.isFile() && !f.delete()) || !f.createNewFile())
+                System.err.printf("Could not dump file %s.class%n", name);
+            else {
+                FileOutputStream fos = new FileOutputStream(f);
+                fos.write(comb.toByteArray());
+                fos.close(); // Implicit flush if underlying stream is buffered
+            }
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+
+        return comb;
+    }
+
 }