Avoid cross-contaminating injector-added locals by when capturing locals

Fixes SpongePowered#493

Adds a preinject pass to local-capturing injectors so they see only
valid locals. Also decorate mixin-added locals to safeguard against
accidental leakage.

src/main/java/org/spongepowered/asm/mixin/injection/callback/CallbackInjector.java

@@ -326,6 +326,11 @@ int marshalVar() {
 
     }
 
+    /**
+     * Decorator key for local variables decoration
+     */
+    private static final String LOCALS_KEY = "locals";
+
     /**
      * True if cancellable 
      */
@@ -417,25 +422,28 @@ protected void addTargetNode(Target target, List<InjectionNode> myNodes, Abstrac
         myNodes.add(injectionNode);
         this.totalInjections++;
     }
-
-    /* (non-Javadoc)
-     * @see org.spongepowered.asm.mixin.injection.callback.BytecodeInjector
-     *      #inject(org.spongepowered.asm.mixin.injection.callback.Target,
-     *      org.objectweb.asm.tree.AbstractInsnNode)
-     */
+
     @Override
-    protected void inject(Target target, InjectionNode node) {
-        LocalVariableNode[] locals = null;
-
+    protected void preInject(Target target, InjectionNode node) {
         if (this.localCapture.isCaptureLocals() || this.localCapture.isPrintLocals()) {
-            locals = Locals.getLocalsAt(this.classNode, target.method, node.getCurrentTarget());
+            LocalVariableNode[] locals = Locals.getLocalsAt(this.classNode, target.method, node.getCurrentTarget());
             for (int j = 0; j < locals.length; j++) {
                 if (locals[j] != null && locals[j].desc != null && locals[j].desc.startsWith("Lorg/spongepowered/asm/mixin/injection/callback/")) {
                     locals[j] = null;
                 }
             }
+            node.<LocalVariableNode[]>decorate(CallbackInjector.LOCALS_KEY, locals);
         }
+    }
 
+    /* (non-Javadoc)
+     * @see org.spongepowered.asm.mixin.injection.callback.BytecodeInjector
+     *      #inject(org.spongepowered.asm.mixin.injection.callback.Target,
+     *      org.objectweb.asm.tree.AbstractInsnNode)
+     */
+    @Override
+    protected void inject(Target target, InjectionNode node) {
+        LocalVariableNode[] locals = node.<LocalVariableNode[]>getDecoration(CallbackInjector.LOCALS_KEY);
         this.inject(new Callback(this.methodNode, target, node, locals, this.localCapture.isCaptureLocals()));
     }
 

src/main/java/org/spongepowered/asm/mixin/injection/code/Injector.java

@@ -247,6 +247,18 @@ protected void addTargetNode(Target target, List<InjectionNode> myNodes, Abstrac
         myNodes.add(target.addInjectionNode(node));
     }
 
+    /**
+     * Performs pre-injection checks and tasks on the specified target
+     * 
+     * @param target target potentially being injected into
+     * @param nodes selected target nodes
+     */
+    public final void preInject(Target target, List<InjectionNode> nodes) {
+        for (InjectionNode node : nodes) {
+            this.preInject(target, node);
+        }
+    }
+
     /**
      * Performs the injection on the specified target
      * 
@@ -381,6 +393,10 @@ protected void checkTargetForNode(Target target, InjectionNode node, RestrictTar
         this.checkTargetModifiers(target, true);
     }
 
+    protected void preInject(Target target, InjectionNode node) {
+        // stub
+    }
+
     protected abstract void inject(Target target, InjectionNode node);
 
     protected void postInject(Target target, InjectionNode node) {
@@ -439,7 +455,7 @@ protected AbstractInsnNode invokeHandlerWithArgs(Type[] args, InsnList insns, in
         this.pushArgs(args, insns, argMap, startArg, endArg);
         return this.invokeHandler(insns);
     }
-
+    
     /**
      * Store args on the stack starting at the end and working back to position
      * specified by start, return the generated argMap
@@ -451,8 +467,24 @@ protected AbstractInsnNode invokeHandlerWithArgs(Type[] args, InsnList insns, in
      * @return the generated argmap
      */
     protected int[] storeArgs(Target target, Type[] args, InsnList insns, int start) {
+        return this.storeArgs(target, args, insns, start, null, null);
+    }
+
+    /**
+     * Store args on the stack starting at the end and working back to position
+     * specified by start, return the generated argMap
+     * 
+     * @param target target method
+     * @param args argument types
+     * @param insns instruction list to generate insns into
+     * @param start Starting index
+     * @param from The label marking the start of the region for the locals
+     * @param to The label marking the end of the region for the stored locals
+     * @return the generated argmap
+     */
+    protected int[] storeArgs(Target target, Type[] args, InsnList insns, int start, LabelNode from, LabelNode to) {
         int[] argMap = target.generateArgMap(args, start);
-        this.storeArgs(args, insns, argMap, start, args.length);
+        this.storeArgs(target, args, insns, argMap, start, args.length, from, to);
         return argMap;
     }
 
@@ -465,9 +497,25 @@ protected int[] storeArgs(Target target, Type[] args, InsnList insns, int start)
      * @param start Starting index
      * @param end Ending index
      */
-    protected void storeArgs(Type[] args, InsnList insns, int[] argMap, int start, int end) {
+    protected void storeArgs(Target target, Type[] args, InsnList insns, int[] argMap, int start, int end) {
+        this.storeArgs(target, args, insns, argMap, start, end, null, null);
+    }
+
+    /**
+     * Store args on the stack to their positions allocated based on argMap
+     * 
+     * @param args argument types
+     * @param insns instruction list to generate insns into
+     * @param argMap generated argmap containing local indices for all args
+     * @param start Starting index
+     * @param end Ending index
+     * @param from The label marking the start of the region for the locals
+     * @param to The label marking the end of the region for the stored locals
+     */
+    protected void storeArgs(Target target, Type[] args, InsnList insns, int[] argMap, int start, int end, LabelNode from, LabelNode to) {
         for (int arg = end - 1; arg >= start; arg--) {
             insns.add(new VarInsnNode(args[arg].getOpcode(Opcodes.ISTORE), argMap[arg]));
+            target.addLocalVariable(argMap[arg], String.format("injectorAllocatedLocal%d", argMap[arg]), args[arg].getDescriptor(), from, to);
         }
     }
 

src/main/java/org/spongepowered/asm/mixin/injection/modify/ModifyVariableInjector.java

@@ -43,6 +43,7 @@
 import org.spongepowered.asm.mixin.injection.struct.InjectionPointData;
 import org.spongepowered.asm.mixin.injection.struct.Target;
 import org.spongepowered.asm.mixin.injection.struct.Target.Extension;
+import org.spongepowered.asm.mixin.injection.throwables.InjectionError;
 import org.spongepowered.asm.mixin.injection.throwables.InvalidInjectionException;
 import org.spongepowered.asm.mixin.refmap.IMixinContext;
 import org.spongepowered.asm.util.Bytecode;
@@ -55,7 +56,7 @@
  * {@link ModifyVariable}.
  */
 public class ModifyVariableInjector extends Injector {
-        
+
     /**
      * Target context information
      */
@@ -93,11 +94,16 @@ public boolean find(String desc, InsnList insns, Collection<AbstractInsnNode> no
 
     }
 
+    /**
+     * Decorator key for Context decoration
+     */
+    private static final String LOCALS_CONTEXT_KEY = "localcontext";
+
     /**
      * True to consider only method args
      */
     private final LocalVariableDiscriminator discriminator;
-
+    
     /**
      * @param info Injection info
      * @param discriminator discriminator
@@ -136,6 +142,12 @@ protected void sanityCheck(Target target, List<InjectionPoint> injectionPoints)
         }
     }
 
+    @Override
+    protected void preInject(Target target, InjectionNode node) {
+        Context context = new Context(this.returnType, this.discriminator.isArgsOnly(), target, node.getCurrentTarget());
+        node.<Context>decorate(ModifyVariableInjector.LOCALS_CONTEXT_KEY, context);
+    }
+
     /**
      * Do the injection
      */
@@ -145,15 +157,27 @@ protected void inject(Target target, InjectionNode node) {
             throw new InvalidInjectionException(this.info, "Variable modifier target for " + this + " was removed by another injector");
         }
 
-        Context context = new Context(this.returnType, this.discriminator.isArgsOnly(), target, node.getCurrentTarget());
+        Context context = node.<Context>getDecoration(ModifyVariableInjector.LOCALS_CONTEXT_KEY);
+        if (context == null) {
+            throw new InjectionError(String.format(
+                    "%s injector target is missing CONTEXT decoration for %s. PreInjection failure or illegal internal state change",
+                    this.annotationType, this.info));
+        }
 
         if (this.discriminator.printLVT()) {
             this.printLocals(target, context);
         }
-
+        
         this.checkTargetForNode(target, node, RestrictTargetLevel.ALLOW_ALL);
 
         InjectorData handler = new InjectorData(target, "handler", false);
+
+        if (this.returnType == Type.VOID_TYPE) {
+            throw new InvalidInjectionException(this.info, String.format(
+                    "%s %s method %s from %s has an invalid signature, cannot return a VOID type.",
+                    this.annotationType, handler, this, this.info.getMixin()));
+        }
+
         this.validateParams(handler, this.returnType, this.returnType);
 
         Extension extraStack = target.extendStack();

src/main/java/org/spongepowered/asm/mixin/injection/struct/InjectionInfo.java

@@ -428,6 +428,15 @@ public void prepare() {
         }
     }
 
+    /**
+     * Perform pre-injection checks and tasks
+     */
+    public void preInject() {
+        for (Entry<Target, List<InjectionNode>> entry : this.targetNodes.entrySet()) {
+            this.injector.preInject(entry.getKey(), entry.getValue());
+        }
+    }
+
     /**
      * Perform injections
      */

src/main/java/org/spongepowered/asm/mixin/injection/struct/Target.java

@@ -28,7 +28,6 @@
 import java.util.Iterator;
 import java.util.List;
 
-import org.objectweb.asm.Label;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.AbstractInsnNode;
@@ -45,6 +44,7 @@
 import org.spongepowered.asm.util.Bytecode;
 import org.spongepowered.asm.util.Bytecode.DelegateInitialiser;
 import org.spongepowered.asm.util.Constants;
+import org.spongepowered.asm.util.Locals.SyntheticLocalVariableNode;
 
 /**
  * Information about the current injection target, mainly just convenience
@@ -751,13 +751,7 @@ public void removeNode(AbstractInsnNode insn) {
      * @param desc local variable type
      */
     public void addLocalVariable(int index, String name, String desc) {
-        if (this.start == null) {
-            this.start = new LabelNode(new Label());
-            this.end = new LabelNode(new Label());
-            this.insns.insert(this.start);
-            this.insns.add(this.end);
-        }
-        this.addLocalVariable(index, name, desc, this.start, this.end);
+        this.addLocalVariable(index, name, desc, null, null);
     }
 
     /**
@@ -766,14 +760,50 @@ public void addLocalVariable(int index, String name, String desc) {
      * @param index local variable index
      * @param name local variable name
      * @param desc local variable type
-     * @param start start of range
-     * @param end end of range
+     * @param from start of range
+     * @param to end of range
      */
-    private void addLocalVariable(int index, String name, String desc, LabelNode start, LabelNode end) {
+    public void addLocalVariable(int index, String name, String desc, LabelNode from, LabelNode to) {
+        if (from == null) {
+            from = this.getStartLabel();
+        }
+
+        if (to == null) {
+            to = this.getEndLabel();
+        }
+
         if (this.method.localVariables == null) {
             this.method.localVariables = new ArrayList<LocalVariableNode>();
         }
-        this.method.localVariables.add(new LocalVariableNode(name, desc, null, start, end, index));
+
+        for (Iterator<LocalVariableNode> iter = this.method.localVariables.iterator(); iter.hasNext();) {
+            LocalVariableNode local = iter.next();
+            if (local != null && local.index == index && from == local.start && to == local.end) {
+                iter.remove();
+            }
+        }
+
+        this.method.localVariables.add(new SyntheticLocalVariableNode(name, desc, null, from, to, index));
     }
-
+
+    /**
+     * Get a label which marks the very start of the method
+     */
+    private LabelNode getStartLabel() {
+        if (this.start == null) {
+            this.insns.insert(this.start = new LabelNode());
+        }
+        return this.start;
+    }
+
+    /**
+     * Get a label which marks the very end of the method
+     */
+    private LabelNode getEndLabel() {
+        if (this.end == null) {
+            this.insns.add(this.end = new LabelNode());
+        }
+        return this.end;
+    }
+
 }

src/main/java/org/spongepowered/asm/mixin/transformer/MixinTargetContext.java

@@ -1277,7 +1277,14 @@ void applyInjections() {
         this.activities.clear();
 
         try {
-            Activity applyActivity = this.activities.begin("Inject");
+            Activity applyActivity = this.activities.begin("PreInject");
+            Activity preInjectActivity = this.activities.begin("?");
+            for (InjectionInfo injectInfo : this.injectors) {
+                preInjectActivity.next(injectInfo.toString());
+                injectInfo.preInject();
+            }
+
+            applyActivity.next("Inject");
             Activity injectActivity = this.activities.begin("?");
             for (InjectionInfo injectInfo : this.injectors) {
                 injectActivity.next(injectInfo.toString());

src/main/java/org/spongepowered/asm/util/Bytecode.java

@@ -659,6 +659,58 @@ public static void loadArgs(Type[] args, InsnList insns, int start, int end, Typ
         }
     }
 
+//    /**
+//     * Generate required APPEND frame nodes for the listed types and append them
+//     * to the supplied instruction list
+//     * 
+//     * @param insns Insnlist to append to
+//     * @param types frame types to append
+//     */
+//    public static void extendFrame(InsnList insns, Type[] types) {
+//        int offset = 0;
+//        Object[] locals = new Object[3];
+//        
+//        for (int index = 0; index < types.length; index++) {
+//            
+//            Type type = types[index];
+//            int size = type.getSize();
+//            if (offset + size <= 3) {
+//                if (type.getSort() < Type.ARRAY) {
+//                    locals[offset] = type.getSort();
+//                } else {
+//                    locals[offset] = type.getInternalName();
+//                }
+//                if (size > 1) {
+//                    locals[offset + 1] = null;
+//                }
+//                offset += size;
+//            } else {
+//                insns.add(new FrameNode(Opcodes.F_APPEND, offset, locals, 0, null));
+//                offset = 0;
+//            }
+//        }
+//        
+//        if (offset > 0) {
+//            insns.add(new FrameNode(Opcodes.F_APPEND, offset, locals, 0, null));
+//        }
+//    }
+//    
+//    /**
+//     * Generate required CHOP frame nodes for the listed types and append them
+//     * to the supplied insn list
+//     * 
+//     * @param insns Insnlist to append generated nodes to
+//     * @param types types to chop
+//     */
+//    public static void chopFrame(InsnList insns, Type[] types) {
+//        int size = Bytecode.getArgsSize(types);
+//        while (size > 0) {
+//            int chop = Math.min(size, 3);
+//            insns.add(new FrameNode(Opcodes.F_CHOP, chop, null, 0, null));
+//            size -= chop;
+//        }
+//    }
+
     /**
      * Get an array of Types from an array of classes.
      *