Add a dex file rewriter to the rewriters module

Previously, if you needed to provide a custom dex file rewriter, you
would have to override the DexRewriter class. This provides a more
consistent interface, allowing the dex file rewriter to be specified in
the module.

4 files changed, 75 insertions, 23 deletions

diff --git a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/DexFileRewriter.java b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/DexFileRewriter.java
new file mode 100644
index 00000000..faa991c7
--- /dev/null
+++ b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/DexFileRewriter.java
@@ -0,0 +1,67 @@
+/*
+ * Copyright 2020, 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.rewriter;
+
+import org.jf.dexlib2.Opcodes;
+import org.jf.dexlib2.iface.ClassDef;
+import org.jf.dexlib2.iface.DexFile;
+
+import javax.annotation.Nonnull;
+import java.util.Set;
+
+public class DexFileRewriter implements Rewriter<DexFile> {
+    @Nonnull protected final Rewriters rewriters;
+
+    public DexFileRewriter(@Nonnull Rewriters rewriters) {
+        this.rewriters = rewriters;
+    }
+
+    @Nonnull @Override public DexFile rewrite(@Nonnull DexFile value) {
+        return new RewrittenDexFile(value);
+    }
+
+    protected class RewrittenDexFile implements DexFile {
+        @Nonnull protected final DexFile dexFile;
+
+        public RewrittenDexFile(@Nonnull DexFile dexFile) {
+            this.dexFile = dexFile;
+        }
+
+        @Override @Nonnull public Set<? extends ClassDef> getClasses() {
+            return RewriterUtils.rewriteSet(rewriters.getClassDefRewriter(), dexFile.getClasses());
+        }
+
+        @Nonnull @Override public Opcodes getOpcodes() {
+            return dexFile.getOpcodes();
+        }
+    }
+}
diff --git a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/DexRewriter.java b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/DexRewriter.java
index 839ac096..84963aa0 100644
--- a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/DexRewriter.java
+++ b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/DexRewriter.java
@@ -31,7 +31,6 @@
 
 package org.jf.dexlib2.rewriter;
 
-import org.jf.dexlib2.Opcodes;
 import org.jf.dexlib2.iface.*;
 import org.jf.dexlib2.iface.debug.DebugItem;
 import org.jf.dexlib2.iface.instruction.Instruction;
@@ -40,7 +39,6 @@ import org.jf.dexlib2.iface.reference.MethodReference;
 import org.jf.dexlib2.iface.value.EncodedValue;
 
 import javax.annotation.Nonnull;
-import java.util.Set;
 
 /**
  * Out-of-the box, this class does nothing except make a picture-perfect copy of a dex file.
@@ -68,6 +66,7 @@ import java.util.Set;
  * </pre>
  */
 public class DexRewriter implements Rewriters {
+    private final Rewriter<DexFile> dexFileRewriter;
     private final Rewriter<ClassDef> classDefRewriter;
     private final Rewriter<Field> fieldRewriter;
     private final Rewriter<Method> methodRewriter;
@@ -85,6 +84,7 @@ public class DexRewriter implements Rewriters {
     private final Rewriter<EncodedValue> encodedValueRewriter;
 
     public DexRewriter(RewriterModule module) {
+        this.dexFileRewriter = module.getDexFileRewriter(this);
         this.classDefRewriter = module.getClassDefRewriter(this);
         this.fieldRewriter = module.getFieldRewriter(this);
         this.methodRewriter = module.getMethodRewriter(this);
@@ -102,27 +102,7 @@ public class DexRewriter implements Rewriters {
         this.encodedValueRewriter = module.getEncodedValueRewriter(this);
     }
 
-    @Nonnull
-    public DexFile rewriteDexFile(@Nonnull DexFile dexFile) {
-        return new RewrittenDexFile(dexFile);
-    }
-
-    protected class RewrittenDexFile implements DexFile {
-        @Nonnull protected final DexFile dexFile;
-
-        public RewrittenDexFile(@Nonnull DexFile dexFile) {
-            this.dexFile = dexFile;
-        }
-
-        @Override @Nonnull public Set<? extends ClassDef> getClasses() {
-            return RewriterUtils.rewriteSet(getClassDefRewriter(), dexFile.getClasses());
-        }
-
-        @Nonnull @Override public Opcodes getOpcodes() {
-            return dexFile.getOpcodes();
-        }
-    }
-
+    @Nonnull @Override public Rewriter<DexFile> getDexFileRewriter() { return dexFileRewriter; }
     @Nonnull @Override public Rewriter<ClassDef> getClassDefRewriter() { return classDefRewriter; }
     @Nonnull @Override public Rewriter<Field> getFieldRewriter() { return fieldRewriter; }
     @Nonnull @Override public Rewriter<Method> getMethodRewriter() { return methodRewriter; }
diff --git a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/RewriterModule.java b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/RewriterModule.java
index 2913a676..5723b6df 100644
--- a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/RewriterModule.java
+++ b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/RewriterModule.java
@@ -41,6 +41,10 @@ import org.jf.dexlib2.iface.value.EncodedValue;
 import javax.annotation.Nonnull;
 
 public class RewriterModule {
+    @Nonnull public Rewriter<DexFile> getDexFileRewriter(@Nonnull Rewriters rewriters) {
+        return new DexFileRewriter(rewriters);
+    }
+
     @Nonnull public Rewriter<ClassDef> getClassDefRewriter(@Nonnull Rewriters rewriters) {
         return new ClassDefRewriter(rewriters);
     }
diff --git a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/Rewriters.java b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/Rewriters.java
index 97fd8792..e64c48f3 100644
--- a/dexlib2/src/main/java/org/jf/dexlib2/rewriter/Rewriters.java
+++ b/dexlib2/src/main/java/org/jf/dexlib2/rewriter/Rewriters.java
@@ -41,6 +41,7 @@ import org.jf.dexlib2.iface.value.EncodedValue;
 import javax.annotation.Nonnull;
 
 public interface Rewriters {
+    @Nonnull Rewriter<DexFile> getDexFileRewriter();
     @Nonnull Rewriter<ClassDef> getClassDefRewriter();
     @Nonnull Rewriter<Field> getFieldRewriter();