Merge pull request #14 from izzytwosheds/payload_alignment

Payload alignment

dexlib2/src/main/java/org/jf/dexlib2/writer/util/InstructionWriteUtil.java

@@ -32,6 +32,7 @@
 package org.jf.dexlib2.writer.util;
 
 import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import org.jf.dexlib2.Format;
 import org.jf.dexlib2.Opcode;
 import org.jf.dexlib2.ReferenceType;
@@ -148,17 +149,14 @@ public class InstructionWriteUtil<Insn extends Instruction, StringRef extends St
     private void findCodeOffsetShifts() {
         // first, process const-string to const-string/jumbo conversions
         int currentCodeOffset = 0;
+        codeOffsetShifts = Lists.newArrayList();
+        offsetToNewInstructionMap = Maps.newHashMap();
+
         for (Instruction instruction: originalInstructions) {
             if (instruction.getOpcode().equals(Opcode.CONST_STRING)) {
                 ReferenceInstruction refInstr = (ReferenceInstruction) instruction;
                 int referenceIndex = stringIndexProvider.getItemIndex((StringRef)refInstr.getReference());
                 if (referenceIndex > 0xFFFF) {
-                    if (codeOffsetShifts == null) {
-                        codeOffsetShifts = new ArrayList<Integer>();
-                    }
-                    if (offsetToNewInstructionMap == null) {
-                        offsetToNewInstructionMap = new HashMap<Integer,Format>();
-                    }
                     codeOffsetShifts.add(currentCodeOffset+instruction.getCodeUnits());
                     offsetToNewInstructionMap.put(currentCodeOffset, Opcode.CONST_STRING_JUMBO.format);
                 }
@@ -166,10 +164,6 @@ public class InstructionWriteUtil<Insn extends Instruction, StringRef extends St
             currentCodeOffset += instruction.getCodeUnits();
         }
 
-        if (codeOffsetShifts == null) {
-            return;
-        }
-
         // next, let's check if this caused any conversions in goto instructions due to changes in offset values
         // since code offset delta is equivalent to the position of instruction's code offset in the shift list,
         // we use it as a position here

dexlib2/src/test/java/org/jf/dexlib2/writer/JumboStringConversionTest.java

@@ -272,6 +272,7 @@ public class JumboStringConversionTest {
         for (Instruction instr: writeUtil.getInstructions()) {
             if (codeOffset == 21) {
                 Assert.assertEquals("array payload was not aligned properly", instr.getOpcode(), Opcode.NOP);
+                break;
             }
             codeOffset += instr.getCodeUnits();
         }
@@ -289,6 +290,7 @@ public class JumboStringConversionTest {
         for (Instruction instr: writeUtil.getInstructions()) {
             if (codeOffset == 7) {
                 Assert.assertEquals("packed switch payload was not aligned properly", instr.getOpcode(), Opcode.NOP);
+                break;
             }
             codeOffset += instr.getCodeUnits();
         }
@@ -307,6 +309,7 @@ public class JumboStringConversionTest {
         for (Instruction instr: writeUtil.getInstructions()) {
             if (codeOffset == 15) {
                 Assert.assertEquals("packed switch payload was not aligned properly", instr.getOpcode(), Opcode.NOP);
+                break;
             }
             codeOffset += instr.getCodeUnits();
         }

dexlib2/src/test/java/org/jf/dexlib2/writer/PayloadAlignmentTest.java

@@ -0,0 +1,131 @@
+/*
+ * Copyright 2012, Google Inc.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google Inc. nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package org.jf.dexlib2.writer;
+
+import com.google.common.collect.Lists;
+import junit.framework.Assert;
+import org.jf.dexlib2.Opcode;
+import org.jf.dexlib2.iface.MethodImplementation;
+import org.jf.dexlib2.iface.instruction.Instruction;
+import org.jf.dexlib2.iface.instruction.SwitchElement;
+import org.jf.dexlib2.iface.reference.Reference;
+import org.jf.dexlib2.iface.reference.StringReference;
+import org.jf.dexlib2.immutable.ImmutableMethodImplementation;
+import org.jf.dexlib2.immutable.instruction.*;
+import org.jf.dexlib2.writer.util.InstructionWriteUtil;
+import org.junit.Before;
+import org.junit.Test;
+
+import javax.annotation.Nonnull;
+import java.util.ArrayList;
+
+public class PayloadAlignmentTest {
+    private MockStringIndexProvider mockStringIndexProvider;
+
+    private class InsnWriteUtil extends InstructionWriteUtil<Instruction, StringReference, Reference> {
+        public InsnWriteUtil(@Nonnull MethodImplementation implementation) {
+            super(implementation.getInstructions(), mockStringIndexProvider, ImmutableInstructionFactory.INSTANCE);
+        }
+    }
+
+    @Before
+    public void setup() {
+        mockStringIndexProvider = new MockStringIndexProvider();
+    }
+
+    @Test
+    public void testArrayPayloadAlignment() {
+        ArrayList<ImmutableInstruction> instructions = Lists.newArrayList();
+
+        // add misaligned array payload
+        instructions.add(new ImmutableInstruction10x(Opcode.NOP));
+        instructions.add(new ImmutableArrayPayload(4, null));
+
+        ImmutableMethodImplementation methodImplementation = new ImmutableMethodImplementation(1, instructions, null, null);
+        InsnWriteUtil writeUtil = new InsnWriteUtil(methodImplementation);
+
+        int codeOffset = 0;
+        for (Instruction instr: writeUtil.getInstructions()) {
+            if (instr.getOpcode().equals(Opcode.ARRAY_PAYLOAD)) {
+                Assert.assertEquals("array payload was not aligned properly", codeOffset%2, 0);
+                break;
+            }
+            codeOffset += instr.getCodeUnits();
+        }
+    }
+
+    @Test
+    public void testPackedSwitchAlignment() {
+        ArrayList<ImmutableInstruction> instructions = Lists.newArrayList();
+        // add misaligned packed switch payload
+        ArrayList<SwitchElement> switchElements = Lists.newArrayList();
+        switchElements.add(new ImmutableSwitchElement(0, 5));
+        instructions.add(new ImmutableInstruction10x(Opcode.NOP));
+        instructions.add(new ImmutablePackedSwitchPayload(switchElements));
+
+        ImmutableMethodImplementation methodImplementation = new ImmutableMethodImplementation(1, instructions, null, null);
+        InsnWriteUtil writeUtil = new InsnWriteUtil(methodImplementation);
+
+        int codeOffset = 0;
+        for (Instruction instr: writeUtil.getInstructions()) {
+            if (instr.getOpcode().equals(Opcode.PACKED_SWITCH_PAYLOAD)) {
+                Assert.assertEquals("packed switch payload was not aligned properly", codeOffset%2, 0);
+                break;
+            }
+            codeOffset += instr.getCodeUnits();
+        }
+    }
+
+    @Test
+    public void testSparseSwitchAlignment() {
+        ArrayList<ImmutableInstruction> instructions = Lists.newArrayList();
+
+        // add misaligned sparse switch payload
+        ArrayList<SwitchElement> switchElements = Lists.newArrayList();
+        switchElements.add(new ImmutableSwitchElement(0, 5));
+
+        instructions.add(new ImmutableInstruction10x(Opcode.NOP));
+        instructions.add(new ImmutableSparseSwitchPayload(switchElements));
+
+        ImmutableMethodImplementation methodImplementation = new ImmutableMethodImplementation(1, instructions, null, null);
+        InsnWriteUtil writeUtil = new InsnWriteUtil(methodImplementation);
+
+        int codeOffset = 0;
+        for (Instruction instr: writeUtil.getInstructions()) {
+            if (instr.getOpcode().equals(Opcode.SPARSE_SWITCH_PAYLOAD)) {
+                Assert.assertEquals("packed switch payload was not aligned properly", codeOffset%2, 0);
+                break;
+            }
+            codeOffset += instr.getCodeUnits();
+        }
+    }
+}