[DEX] Parse smali code segments

REAndroid · May 8, 2024 · 6e891dd · 6e891dd
1 parent ca86436
commit 6e891dd
Show file tree

Hide file tree

Showing 20 changed files with 352 additions and 169 deletions.
diff --git a/src/main/java/com/reandroid/dex/data/CodeItem.java b/src/main/java/com/reandroid/dex/data/CodeItem.java
@@ -239,7 +239,7 @@ public void merge(CodeItem codeItem){
     public void fromSmali(SmaliMethod smaliMethod) throws IOException {
         setRegistersCount(smaliMethod.getRegistersCount());
         setParameterRegistersCount(smaliMethod.getParameterRegistersCount());
-        getInstructionList().fromSmali(smaliMethod);
+        getInstructionList().fromSmali(smaliMethod.getCodeSet());
         Iterator<SmaliCodeTryItem> iterator = smaliMethod.getTryItems();
         TryBlock tryBlock = null;
         if(iterator.hasNext()){
@@ -250,7 +250,7 @@ public void fromSmali(SmaliMethod smaliMethod) throws IOException {
         }
         if(smaliMethod.hasDebugs()){
             DebugInfo debugInfo = getOrCreateDebugInfo();
-            debugInfo.getDebugSequence().fromSmali(smaliMethod);
+            debugInfo.getDebugSequence().fromSmali(smaliMethod.getCodeSet());
         }
     }
     @Override

diff --git a/src/main/java/com/reandroid/dex/data/InstructionList.java b/src/main/java/com/reandroid/dex/data/InstructionList.java
@@ -34,8 +34,8 @@
 import com.reandroid.dex.sections.SectionType;
 import com.reandroid.dex.smali.SmaliFormat;
 import com.reandroid.dex.smali.SmaliWriter;
+import com.reandroid.dex.smali.model.SmaliCodeSet;
 import com.reandroid.dex.smali.model.SmaliInstruction;
-import com.reandroid.dex.smali.model.SmaliMethod;
 import com.reandroid.utils.CompareUtil;
 import com.reandroid.utils.collection.*;
 
@@ -319,6 +319,9 @@ public ConstString createStringAt(int index, int register, StringKey value) {
         return constNumber;
     }
     public<T1 extends Ins> T1 createAt(int index, Opcode<T1> opcode) {
+        if(index == getCount()) {
+            return createNext(opcode);
+        }
         T1 item = opcode.newInstance();
         add(index, item);
         return item;
@@ -645,12 +648,26 @@ public void merge(InstructionList instructionList){
         getInsArray().trimToSize();
         updateAddresses();
     }
-    public void fromSmali(SmaliMethod smaliMethod) throws IOException {
-        Iterator<SmaliInstruction> iterator = smaliMethod.getInstructions();
+    public void fromSmali(SmaliCodeSet smaliCodeSet) throws IOException {
+        int index = 0;
+        int offset = smaliCodeSet.getAddressOffset();
+        if(offset != 0) {
+            Ins ins = getAtAddress(smaliCodeSet.getAddressOffset());
+            if(ins != null) {
+                index = ins.getIndex();
+            }else if(offset >= getCodeUnits()) {
+                index = getCount();
+            }
+        }
+        fromSmali(index, smaliCodeSet);
+    }
+    public void fromSmali(int index, SmaliCodeSet smaliCodeSet) throws IOException {
+        Iterator<SmaliInstruction> iterator = smaliCodeSet.getInstructions();
         while (iterator.hasNext()){
             SmaliInstruction smaliInstruction = iterator.next();
-            Ins ins = createNext(smaliInstruction.getOpcode());
+            Ins ins = createAt(index, smaliInstruction.getOpcode());
             ins.fromSmali(smaliInstruction);
+            index ++;
         }
         getInsArray().trimToSize();
         updateAddresses();

diff --git a/src/main/java/com/reandroid/dex/debug/DebugSequence.java b/src/main/java/com/reandroid/dex/debug/DebugSequence.java
@@ -21,6 +21,7 @@
 import com.reandroid.dex.base.FixedDexContainer;
 import com.reandroid.dex.data.InstructionList;
 import com.reandroid.dex.id.IdItem;
+import com.reandroid.dex.smali.model.SmaliCodeSet;
 import com.reandroid.dex.smali.model.SmaliDebug;
 import com.reandroid.dex.smali.model.SmaliMethod;
 import com.reandroid.utils.collection.*;
@@ -293,8 +294,8 @@ public void merge(DebugSequence sequence){
         cacheValues();
         getElementList().trimToSize();
     }
-    public void fromSmali(SmaliMethod smaliMethod) throws IOException {
-        Iterator<SmaliDebug> iterator = smaliMethod.getDebugs();
+    public void fromSmali(SmaliCodeSet smaliCodeSet) throws IOException {
+        Iterator<SmaliDebug> iterator = smaliCodeSet.getDebugs();
         while (iterator.hasNext()){
             SmaliDebug smaliDebug = iterator.next();
             DebugElementType<?> type = smaliDebug.getDebugElementType();

diff --git a/src/main/java/com/reandroid/dex/model/DexClass.java b/src/main/java/com/reandroid/dex/model/DexClass.java
@@ -21,7 +21,10 @@
 import com.reandroid.dex.id.IdItem;
 import com.reandroid.dex.key.*;
 import com.reandroid.dex.reference.TypeListReference;
+import com.reandroid.dex.smali.SmaliReader;
 import com.reandroid.dex.smali.SmaliWriter;
+import com.reandroid.dex.smali.model.SmaliField;
+import com.reandroid.dex.smali.model.SmaliMethod;
 import com.reandroid.dex.value.DexValueBlock;
 import com.reandroid.dex.value.DexValueType;
 import com.reandroid.dex.value.NullValue;
@@ -302,6 +305,12 @@ public Iterator<DexField> getInstanceFields() {
         }
         return ComputeIterator.of(classData.getInstanceFields(), this::createField);
     }
+    public DexField getOrCreateInstanceField(FieldKey fieldKey){
+        return createField(getOrCreateInstance(fieldKey));
+    }
+    public FieldDef getOrCreateInstance(FieldKey fieldKey){
+        return getOrCreateClassData().getOrCreateInstance(fieldKey);
+    }
     public Iterator<DexMethod> getDeclaredMethods(Predicate<DexMethod> filter) {
         Iterator<DexMethod> iterator = getDeclaredMethods();
         if(filter == null){
@@ -620,6 +629,49 @@ public void append(SmaliWriter writer) throws IOException {
         getClassData();
         getId().append(writer);
     }
+
+    public DexMethod parseMethod(SmaliReader reader) throws IOException {
+        SmaliMethod smaliMethod = new SmaliMethod();
+        smaliMethod.setDefining(getKey());
+        smaliMethod.parse(reader);
+        MethodKey methodKey = smaliMethod.getKey();
+        DexMethod dexMethod;
+        if(smaliMethod.isDirect()) {
+            dexMethod = getOrCreateDirectMethod(methodKey);
+        } else {
+            dexMethod = getOrCreateVirtualMethod(methodKey);
+        }
+        dexMethod.getDefinition().fromSmali(smaliMethod);
+        return dexMethod;
+    }
+    public DexField parseField(SmaliReader reader) throws IOException {
+        SmaliField smaliField = new SmaliField();
+        smaliField.setDefining(getKey());
+        smaliField.parse(reader);
+        FieldKey fieldKey = smaliField.getKey();
+        DexField dexField;
+        if(smaliField.isInstance()) {
+            dexField = getOrCreateInstanceField(fieldKey);
+        } else {
+            dexField = getOrCreateStaticField(fieldKey);
+        }
+        dexField.getDefinition().fromSmali(smaliField);
+        return dexField;
+    }
+    public DexField parseInterfaces(SmaliReader reader) throws IOException {
+        SmaliField smaliField = new SmaliField();
+        smaliField.setDefining(getKey());
+        smaliField.parse(reader);
+        FieldKey fieldKey = smaliField.getKey();
+        DexField dexField;
+        if(smaliField.isInstance()) {
+            dexField = getOrCreateInstanceField(fieldKey);
+        } else {
+            dexField = getOrCreateStaticField(fieldKey);
+        }
+        dexField.getDefinition().fromSmali(smaliField);
+        return dexField;
+    }
     public void writeSmali(SmaliWriter writer, File dir) throws IOException {
         File file = toSmaliFile(dir);
         IOUtil.writeUtf8(toSmali(writer), file);

diff --git a/src/main/java/com/reandroid/dex/smali/model/SmaliAnnotationDirectory.java b/src/main/java/com/reandroid/dex/smali/model/SmaliAnnotationDirectory.java
diff --git a/src/main/java/com/reandroid/dex/smali/model/SmaliAnnotationItem.java b/src/main/java/com/reandroid/dex/smali/model/SmaliAnnotationItem.java
@@ -29,6 +29,7 @@ public class SmaliAnnotationItem extends SmaliSet<SmaliAnnotationElement> implem
 
     public SmaliAnnotationItem(){
         super();
+        this.smaliDirective = SmaliDirective.ANNOTATION;
     }
 
     public AnnotationVisibility getVisibility() {
@@ -50,8 +51,16 @@ public SmaliDirective getSmaliDirective() {
         return smaliDirective;
     }
 
-    public void setSmaliDirective(SmaliDirective smaliDirective) {
-        this.smaliDirective = smaliDirective;
+    public void setSmaliDirective(SmaliDirective annotationDirective) {
+        if(annotationDirective == null) {
+            throw new NullPointerException("Null annotation directive");
+        }
+        if(annotationDirective != SmaliDirective.ANNOTATION &&
+                annotationDirective != SmaliDirective.SUB_ANNOTATION) {
+            throw new IllegalArgumentException("Invalid annotation directive: "
+                    + annotationDirective);
+        }
+        this.smaliDirective = annotationDirective;
     }
 
     @Override
@@ -77,19 +86,33 @@ public void parse(SmaliReader reader) throws IOException{
         setSmaliDirective(directive);
         setVisibility(AnnotationVisibility.parse(reader));
         setType(TypeKey.read(reader));
-        while (parseElements(reader)){
+        while (parseNext(reader) != null){
             reader.skipWhitespacesOrComment();
         }
         SmaliParseException.expect(reader, getSmaliDirective(), true);
     }
-    private boolean parseElements(SmaliReader reader) throws IOException {
+    @Override
+    SmaliAnnotationElement createNext(SmaliReader reader) {
         reader.skipWhitespacesOrComment();
+        if(reader.finished()) {
+            return null;
+        }
         if(getSmaliDirective().isEnd(reader)){
-            return false;
+            return null;
+        }
+        return new SmaliAnnotationElement();
+    }
+
+    public static SmaliAnnotationItem read(SmaliReader reader) throws IOException {
+        reader.skipWhitespacesOrComment();
+        if(reader.finished()) {
+            return null;
+        }
+        SmaliAnnotationItem smali = new SmaliAnnotationItem();
+        smali.parse(reader);
+        if(!smali.isEmpty()) {
+            return smali;
         }
-        SmaliAnnotationElement element = new SmaliAnnotationElement();
-        add(element);
-        element.parse(reader);
-        return true;
+        return null;
     }
 }
diff --git a/src/main/java/com/reandroid/dex/smali/model/SmaliAnnotationSet.java b/src/main/java/com/reandroid/dex/smali/model/SmaliAnnotationSet.java
@@ -33,14 +33,7 @@ public void append(SmaliWriter writer) throws IOException {
     }
 
     @Override
-    public void parse(SmaliReader reader) throws IOException{
-        SmaliAnnotationItem item;
-        while ((item = createNext(reader)) != null){
-            add(item);
-            item.parse(reader);
-        }
-    }
-    private SmaliAnnotationItem createNext(SmaliReader reader){
+    SmaliAnnotationItem createNext(SmaliReader reader) {
         reader.skipWhitespacesOrComment();
         SmaliDirective directive = SmaliDirective.parse(reader, false);
         if(directive != SmaliDirective.ANNOTATION && directive != SmaliDirective.SUB_ANNOTATION){
@@ -52,4 +45,12 @@ private SmaliAnnotationItem createNext(SmaliReader reader){
         }
         return new SmaliAnnotationItem();
     }
+    public static SmaliAnnotationSet read(SmaliReader reader) throws IOException {
+        SmaliAnnotationSet smali = new SmaliAnnotationSet();
+        smali.parse(reader);
+        if(!smali.isEmpty()) {
+            return smali;
+        }
+        return null;
+    }
 }
diff --git a/src/main/java/com/reandroid/dex/smali/model/SmaliClass.java b/src/main/java/com/reandroid/dex/smali/model/SmaliClass.java
@@ -189,6 +189,13 @@ private void parseSource(SmaliReader reader) throws IOException{
         setSourceFile(StringKey.read(reader));
     }
 
+    public SmaliField parseField(SmaliReader reader) throws IOException {
+        return fields.parseNext(reader);
+    }
+    public SmaliMethod parseMethod(SmaliReader reader) throws IOException {
+        return methods.parseNext(reader);
+    }
+
     @Override
     public String toDebugString() {
         return "class = " + getKey();

diff --git a/src/main/java/com/reandroid/dex/smali/model/SmaliCodeSet.java b/src/main/java/com/reandroid/dex/smali/model/SmaliCodeSet.java
@@ -19,19 +19,27 @@
 import com.reandroid.dex.smali.SmaliDirective;
 import com.reandroid.dex.smali.SmaliReader;
 import com.reandroid.dex.smali.SmaliWriter;
-import com.reandroid.utils.collection.InstanceIterator;
 
 import java.io.IOException;
 import java.util.Iterator;
 
-public class SmaliCodeSet extends SmaliSet<SmaliCode>{
+public class SmaliCodeSet extends SmaliSet<SmaliCode> {
+
+    private int addressOffset;
 
     public SmaliCodeSet(){
         super();
     }
 
-    public void updateAddresses(){
-        int address = 0;
+    public int getAddressOffset() {
+        return addressOffset;
+    }
+    public void setAddressOffset(int addressOffset) {
+        this.addressOffset = addressOffset;
+    }
+
+    public void updateAddresses() {
+        int address = getAddressOffset();
         Iterator<SmaliInstruction> iterator = getInstructions();
         while (iterator.hasNext()) {
             SmaliInstruction ins = iterator.next();
@@ -40,10 +48,19 @@ public void updateAddresses(){
         }
     }
     public Iterator<SmaliInstruction> getInstructions() {
-        return InstanceIterator.of(iterator(), SmaliInstruction.class);
+        return iterator(SmaliInstruction.class);
+    }
+    public Iterator<SmaliCodeTryItem> getTryItems() {
+        return iterator(SmaliCodeTryItem.class);
     }
     public Iterator<SmaliDebug> getDebugs() {
-        return InstanceIterator.of(iterator(), SmaliDebug.class);
+        return iterator(SmaliDebug.class);
+    }
+    public void clearInstructions() {
+        removeInstances(SmaliDebug.class);
+    }
+    public void clearDebugs() {
+        removeInstances(SmaliDebug.class);
     }
 
     @Override
@@ -57,15 +74,11 @@ public void append(SmaliWriter writer) throws IOException {
 
     @Override
     public void parse(SmaliReader reader) throws IOException {
-        SmaliCode code;
-        while ((code = createNext(reader)) != null){
-            add(code);
-            code.parse(reader);
-            reader.skipWhitespacesOrComment();
-        }
+        super.parse(reader);
         updateAddresses();
     }
-    private SmaliCode createNext(SmaliReader reader){
+    @Override
+    SmaliCode createNext(SmaliReader reader) {
         SmaliDirective directive = SmaliDirective.parse(reader, false);
         if(directive != null){
             return createFor(directive);
@@ -80,7 +93,7 @@ private SmaliCode createNext(SmaliReader reader){
         }
         return null;
     }
-    public static SmaliCode createFor(SmaliDirective directive){
+    private static SmaliCode createFor(SmaliDirective directive){
         if(directive == SmaliDirective.LINE){
             return new SmaliLineNumber();
         }