JesusFreke · andvgal · Jul 27, 2022
diff --git a/dexlib2/src/main/java/org/jf/dexlib2/dexbacked/DexBackedDexFile.java b/dexlib2/src/main/java/org/jf/dexlib2/dexbacked/DexBackedDexFile.java
@@ -257,6 +257,12 @@ public List<? extends Reference> getReferences(int referenceType) {
                 return getMethodSection();
             case ReferenceType.FIELD:
                 return getFieldSection();
+            case ReferenceType.METHOD_PROTO:
+                return getMethodSection();
+            case ReferenceType.METHOD_HANDLE:
+                return getMethodHandleSection();
+            case ReferenceType.CALL_SITE:
+                return getCallSiteSection();
             default:
                 throw new IllegalArgumentException(String.format("Invalid reference type: %d", referenceType));
         }

diff --git a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/CallSiteReferenceRewriter.java b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/CallSiteReferenceRewriter.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2014, Google Inc.
+ * Copyright 2022, Zimperium 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.rewriter;
+
+import com.google.common.base.Function;
+import com.google.common.collect.Lists;
+import org.jf.dexlib2.base.reference.BaseCallSiteReference;
+import org.jf.dexlib2.iface.reference.CallSiteReference;
+import org.jf.dexlib2.iface.reference.MethodHandleReference;
+import org.jf.dexlib2.iface.reference.MethodProtoReference;
+import org.jf.dexlib2.iface.value.EncodedValue;
+
+import javax.annotation.Nonnull;
+import java.util.List;
+
+public class CallSiteReferenceRewriter implements Rewriter<CallSiteReference> {
+    @Nonnull protected final Rewriters rewriters;
+
+    public CallSiteReferenceRewriter(@Nonnull Rewriters rewriters) {
+        this.rewriters = rewriters;
+    }
+
+    @Nonnull @Override public CallSiteReference rewrite(@Nonnull CallSiteReference callSiteReference) {
+        return new RewrittenCallSiteReference(callSiteReference);
+    }
+
+    protected class RewrittenCallSiteReference extends BaseCallSiteReference {
+        @Nonnull protected CallSiteReference callSiteReference;
+
+        public RewrittenCallSiteReference(@Nonnull CallSiteReference callSiteReference) {
+            this.callSiteReference = callSiteReference;
+        }
+
+        @Override @Nonnull public String getName() {
+            return callSiteReference.getName();
+        }
+
+        @Override @Nonnull public MethodHandleReference getMethodHandle() {
+                return RewriterUtils.rewriteMethodHandleReference(
+                        rewriters, callSiteReference.getMethodHandle());
+        }
+
+        @Override @Nonnull public String getMethodName() {
+            return callSiteReference.getMethodName();
+        }
+
+        @Override @Nonnull public MethodProtoReference getMethodProto() {
+            return RewriterUtils.rewriteMethodProtoReference(
+                        rewriters.getTypeRewriter(),
+                        callSiteReference.getMethodProto());
+        }
+
+        @Override @Nonnull public List<? extends EncodedValue> getExtraArguments() {
+            return Lists.transform(callSiteReference.getExtraArguments(),
+                    new Function<EncodedValue, EncodedValue>() {
+                        @Nonnull @Override public EncodedValue apply(EncodedValue encodedValue) {
+                            return RewriterUtils.rewriteValue(rewriters, encodedValue);
+                        }
+                    });
+        }
+    }
+}
diff --git a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/DexRewriter.java b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/DexRewriter.java
@@ -36,6 +36,7 @@
 import org.jf.dexlib2.iface.instruction.Instruction;
 import org.jf.dexlib2.iface.reference.FieldReference;
 import org.jf.dexlib2.iface.reference.MethodReference;
+import org.jf.dexlib2.iface.reference.CallSiteReference;
 import org.jf.dexlib2.iface.value.EncodedValue;
 
 import javax.annotation.Nonnull;
@@ -79,6 +80,7 @@ public class DexRewriter implements Rewriters {
     private final Rewriter<String> typeRewriter;
     private final Rewriter<FieldReference> fieldReferenceRewriter;
     private final Rewriter<MethodReference> methodReferenceRewriter;
+    private final Rewriter<CallSiteReference> callSiteReferenceRewriter;
     private final Rewriter<Annotation> annotationRewriter;
     private final Rewriter<AnnotationElement> annotationElementRewriter;
     private final Rewriter<EncodedValue> encodedValueRewriter;
@@ -97,6 +99,7 @@ public DexRewriter(RewriterModule module) {
         this.typeRewriter = module.getTypeRewriter(this);
         this.fieldReferenceRewriter = module.getFieldReferenceRewriter(this);
         this.methodReferenceRewriter = module.getMethodReferenceRewriter(this);
+        this.callSiteReferenceRewriter = module.getCallSiteReferenceRewriter(this);
         this.annotationRewriter = module.getAnnotationRewriter(this);
         this.annotationElementRewriter = module.getAnnotationElementRewriter(this);
         this.encodedValueRewriter = module.getEncodedValueRewriter(this);
@@ -115,6 +118,7 @@ public DexRewriter(RewriterModule module) {
     @Nonnull @Override public Rewriter<String> getTypeRewriter() { return typeRewriter; }
     @Nonnull @Override public Rewriter<FieldReference> getFieldReferenceRewriter() { return fieldReferenceRewriter; }
     @Nonnull @Override public Rewriter<MethodReference> getMethodReferenceRewriter() { return methodReferenceRewriter; }
+    @Nonnull @Override public Rewriter<CallSiteReference> getCallSiteReferenceRewriter() { return callSiteReferenceRewriter; }
     @Nonnull @Override public Rewriter<Annotation> getAnnotationRewriter() { return annotationRewriter; }
     @Nonnull @Override public Rewriter<AnnotationElement> getAnnotationElementRewriter() { return annotationElementRewriter; }
     @Nonnull @Override public Rewriter<EncodedValue> getEncodedValueRewriter() { return encodedValueRewriter; }

diff --git a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/InstructionRewriter.java b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/InstructionRewriter.java
@@ -33,6 +33,7 @@
 
 import org.jf.dexlib2.Opcode;
 import org.jf.dexlib2.ReferenceType;
+import org.jf.dexlib2.iface.instruction.DualReferenceInstruction;
 import org.jf.dexlib2.iface.instruction.Instruction;
 import org.jf.dexlib2.iface.instruction.ReferenceInstruction;
 import org.jf.dexlib2.iface.instruction.formats.*;
@@ -67,13 +68,42 @@ public InstructionRewriter(@Nonnull Rewriters rewriters) {
                     return new RewrittenInstruction3rc((Instruction3rc)instruction);
                 case Format45cc:
                     return new RewrittenInstruction45cc((Instruction45cc) instruction);
+                case Format4rcc:
+                    return new RewrittenInstruction4rcc((Instruction4rcc) instruction);
                 default:
                     throw new IllegalArgumentException();
             }
         }
         return instruction;
     }
 
+    @Nonnull private Reference rewriteReference(int type,
+                                                @Nonnull Reference reference) {
+        switch (type) {
+            case ReferenceType.TYPE:
+                return RewriterUtils.rewriteTypeReference(rewriters.getTypeRewriter(),
+                        (TypeReference)reference);
+            case ReferenceType.FIELD:
+                return rewriters.getFieldReferenceRewriter().rewrite((FieldReference)reference);
+            case ReferenceType.METHOD:
+                return rewriters.getMethodReferenceRewriter().rewrite((MethodReference)reference);
+            case ReferenceType.STRING:
+                return reference;
+            case ReferenceType.METHOD_PROTO:
+                return RewriterUtils.rewriteMethodProtoReference(
+                        rewriters.getTypeRewriter(),
+                        (MethodProtoReference)reference);
+            case ReferenceType.METHOD_HANDLE:
+                return RewriterUtils.rewriteMethodHandleReference(
+                        rewriters, (MethodHandleReference)reference);
+            case ReferenceType.CALL_SITE:
+                return rewriters.getCallSiteReferenceRewriter().rewrite((CallSiteReference)reference);
+            default:
+                throw new ExceptionWithContext("Invalid reference type: %d",
+                                type);
+        }
+    }
+
     protected class BaseRewrittenReferenceInstruction<T extends ReferenceInstruction>
             implements ReferenceInstruction {
         @Nonnull protected T instruction;
@@ -92,6 +122,9 @@ protected BaseRewrittenReferenceInstruction(@Nonnull T instruction) {
                 case ReferenceType.METHOD:
                     return rewriters.getMethodReferenceRewriter().rewrite((MethodReference)instruction.getReference());
                 case ReferenceType.STRING:
+                case ReferenceType.METHOD_PROTO:
+                case ReferenceType.METHOD_HANDLE:
+                case ReferenceType.CALL_SITE:
                     return instruction.getReference();
                 default:
                     throw new IllegalArgumentException();
@@ -205,9 +238,10 @@ public int getRegisterCount() {
         }
     }
 
-    protected class RewrittenInstruction45cc extends BaseRewrittenReferenceInstruction<Instruction45cc>
-            implements Instruction45cc {
-        public RewrittenInstruction45cc(@Nonnull Instruction45cc instruction) {
+    protected class BaseRewrittenDualReferenceInstruction<T extends DualReferenceInstruction>
+            extends BaseRewrittenReferenceInstruction<T>
+            implements DualReferenceInstruction {
+        public BaseRewrittenDualReferenceInstruction(@Nonnull T instruction) {
             super(instruction);
         }
 
@@ -221,8 +255,9 @@ public RewrittenInstruction45cc(@Nonnull Instruction45cc instruction) {
                 case ReferenceType.METHOD:
                     return rewriters.getMethodReferenceRewriter().rewrite((MethodReference)instruction.getReference2());
                 case ReferenceType.STRING:
-                    return instruction.getReference2();
                 case ReferenceType.METHOD_PROTO:
+                case ReferenceType.METHOD_HANDLE:
+                case ReferenceType.CALL_SITE:
                     return instruction.getReference2();
                 default:
                     throw new IllegalArgumentException();
@@ -232,6 +267,13 @@ public RewrittenInstruction45cc(@Nonnull Instruction45cc instruction) {
         public int getReferenceType2() {
             return instruction.getReferenceType2();
         }
+    }
+
+    protected class RewrittenInstruction45cc extends BaseRewrittenDualReferenceInstruction<Instruction45cc>
+            implements Instruction45cc {
+        public RewrittenInstruction45cc(@Nonnull Instruction45cc instruction) {
+            super(instruction);
+        }
 
         public int getRegisterC() {
             return instruction.getRegisterC();
@@ -257,4 +299,34 @@ public int getRegisterF() {
             return instruction.getRegisterF();
         }
     }
+
+    protected class RewrittenInstruction4rcc extends BaseRewrittenDualReferenceInstruction<Instruction4rcc>
+            implements Instruction4rcc {
+        public RewrittenInstruction4rcc(@Nonnull Instruction4rcc instruction) {
+            super(instruction);
+        }
+
+        public int getStartRegister() {
+            return instruction.getStartRegister();
+        }
+
+        public int getRegisterCount() {
+            return instruction.getRegisterCount();
+        }
+    }
+
+    protected class RewrittenInstruction4rcc extends BaseRewrittenDualReferenceInstruction<Instruction4rcc>
+            implements Instruction4rcc {
+        public RewrittenInstruction4rcc(@Nonnull Instruction4rcc instruction) {
+            super(instruction);
+        }
+
+        public int getStartRegister() {
+            return instruction.getStartRegister();
+        }
+
+        public int getRegisterCount() {
+            return instruction.getRegisterCount();
+        }
+    }
 }
diff --git a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/RewriterModule.java b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/RewriterModule.java
@@ -36,6 +36,7 @@
 import org.jf.dexlib2.iface.instruction.Instruction;
 import org.jf.dexlib2.iface.reference.FieldReference;
 import org.jf.dexlib2.iface.reference.MethodReference;
+import org.jf.dexlib2.iface.reference.CallSiteReference;
 import org.jf.dexlib2.iface.value.EncodedValue;
 
 import javax.annotation.Nonnull;
@@ -93,6 +94,10 @@ public class RewriterModule {
         return new MethodReferenceRewriter(rewriters);
     }
 
+    @Nonnull public Rewriter<CallSiteReference> getCallSiteReferenceRewriter(@Nonnull Rewriters rewriters) {
+        return new CallSiteReferenceRewriter(rewriters);
+    }
+
     @Nonnull public Rewriter<Annotation> getAnnotationRewriter(@Nonnull Rewriters rewriters) {
         return new AnnotationRewriter(rewriters);
     }