Determinism and hardening improvements

dexlib2/src/main/java/com/android/tools/smali/dexlib2/base/reference/BaseCallSiteReference.java

@@ -48,7 +48,8 @@ public int hashCode() {
     public boolean equals(Object o) {
         if (o != null && o instanceof CallSiteReference) {
             CallSiteReference other = (CallSiteReference) o;
-            return getMethodHandle().equals(other.getMethodHandle()) &&
+            return getName().equals(other.getName()) &&
+                    getMethodHandle().equals(other.getMethodHandle()) &&
                     getMethodName().equals(other.getMethodName()) &&
                     getMethodProto().equals(other.getMethodProto()) &&
                     getExtraArguments().equals(other.getExtraArguments());

dexlib2/src/main/java/com/android/tools/smali/dexlib2/dexbacked/DexBackedDexFile.java

@@ -89,9 +89,12 @@ public class DexBackedDexFile implements DexFile {
     private final int classCount;
     private final int classStartOffset;
     private final int mapOffset;
-    private final int hiddenApiRestrictionsOffset;
+    private int hiddenApiRestrictionsOffset = -1;
+    private final int headerOffset;
+    private final int containerSize;
 
-    protected DexBackedDexFile(@Nullable Opcodes opcodes, @Nonnull byte[] buf, int offset, boolean verifyMagic) {
+    protected DexBackedDexFile(@Nullable Opcodes opcodes, @Nonnull byte[] buf, int offset,
+                               boolean verifyMagic) {
         this(opcodes, buf, offset, verifyMagic, 0);
     }
 
@@ -100,18 +103,45 @@ protected DexBackedDexFile(@Nullable Opcodes opcodes,
                                int offset,
                                boolean verifyMagic,
                                int header_offset) {
+        if (offset < 0 || offset > buf.length) {
+            throw new IndexOutOfBoundsException(
+                    String.format("Invalid offset %d for buffer of length %d", offset, buf.length));
+        }
+        if (header_offset < 0 || buf.length - offset < header_offset) {
+            throw new IndexOutOfBoundsException(String.format(
+                    "Invalid header_offset %d for buffer of length %d at offset %d",
+                    header_offset, buf.length, offset));
+        }
+
         dexBuffer = new DexBuffer(buf, offset);
         dataBuffer = new DexBuffer(buf, offset + getBaseDataOffset());
 
-        int dexVersion = getVersion(buf, offset, verifyMagic);
+        int dexVersion = getVersion(buf, offset + header_offset, verifyMagic);
 
         if (opcodes == null) {
             this.opcodes = getDefaultOpcodes(dexVersion);
         } else {
             this.opcodes = opcodes;
         }
 
+        int requiredHeaderSize = (dexVersion >= 41) ? 120 : 112;
+        if (buf.length - offset - header_offset < requiredHeaderSize) {
+            throw new DexUtil.InvalidFile(String.format(
+                    "File is too short to contain a valid dex header (size: %d, required: %d)",
+                    buf.length - offset - header_offset, requiredHeaderSize));
+        }
+
         fileSize = dexBuffer.readSmallUint(header_offset + HeaderItem.FILE_SIZE_OFFSET);
+        if (fileSize < requiredHeaderSize) {
+            throw new DexUtil.InvalidFile(String.format(
+                    "Invalid file_size %d (smaller than required header size %d)",
+                    fileSize, requiredHeaderSize));
+        }
+        if (fileSize > buf.length - offset - header_offset) {
+            throw new DexUtil.InvalidFile(String.format(
+                    "Invalid file_size %d (exceeds remaining buffer length %d)",
+                    fileSize, buf.length - offset - header_offset));
+        }
         stringCount = dexBuffer.readSmallUint(header_offset + HeaderItem.STRING_COUNT_OFFSET);
         stringStartOffset = dexBuffer.readSmallUint(header_offset + HeaderItem.STRING_START_OFFSET);
         typeCount = dexBuffer.readSmallUint(header_offset + HeaderItem.TYPE_COUNT_OFFSET);
@@ -126,19 +156,28 @@ protected DexBackedDexFile(@Nullable Opcodes opcodes,
         classStartOffset = dexBuffer.readSmallUint(header_offset + HeaderItem.CLASS_START_OFFSET);
         mapOffset = dexBuffer.readSmallUint(header_offset + HeaderItem.MAP_OFFSET);
 
-        MapItem mapItem = getMapItemForSection(ItemType.HIDDENAPI_CLASS_DATA_ITEM);
-        if (mapItem != null) {
-            hiddenApiRestrictionsOffset = mapItem.getOffset();
-        } else {
-            hiddenApiRestrictionsOffset = DexWriter.NO_OFFSET;
-        }
-
+        this.headerOffset = header_offset;
         int container_off = 0;
         if (dexVersion >= 41) {
-          container_off = dexBuffer.readSmallUint(header_offset + HeaderItem.CONTAINER_OFF_OFFSET);
+            container_off = dexBuffer.readSmallUint(
+                    header_offset + HeaderItem.HEADER_OFFSET_OFFSET);
+            this.containerSize = dexBuffer.readSmallUint(
+                    header_offset + HeaderItem.CONTAINER_SIZE_OFFSET);
+            if (this.containerSize < header_offset + fileSize) {
+                throw new DexUtil.InvalidFile(String.format(
+                        "DEX entry (header_offset: %d, file_size: %d) exceeds container_size %d",
+                        header_offset, fileSize, containerSize));
+            }
+            if (this.containerSize > buf.length - offset) {
+                throw new DexUtil.InvalidFile(String.format(
+                        "Invalid container_size %d (exceeds buffer length %d)",
+                        containerSize, buf.length - offset));
+            }
+        } else {
+            this.containerSize = this.fileSize;
         }
         if (container_off != header_offset) {
-          throw new DexUtil.InvalidFile(String.format("Unexpected container offset in header"));
+            throw new DexUtil.InvalidFile(String.format("Unexpected container offset in header"));
         }
     }
 
@@ -157,6 +196,20 @@ public int getFileSize() {
         return fileSize;
     }
 
+    /**
+     * @return True if this is the first entry in a DEX container (or classic DEX).
+     */
+    public boolean isDexContainerFirstEntry() {
+        return headerOffset == 0;
+    }
+
+    /**
+     * @return True if this is the last entry in a DEX container, ignoring trailing garbage.
+     */
+    public boolean isDexContainerLastEntry() {
+        return headerOffset + fileSize >= containerSize;
+    }
+
     protected int getVersion(byte[] buf, int offset, boolean verifyMagic) {
         if (verifyMagic) {
             return DexUtil.verifyDexHeader(buf, offset);
@@ -563,20 +616,33 @@ protected DexBackedMethodImplementation createMethodImplementation(
         return new DexBackedMethodImplementation(dexFile, method, codeOffset);
     }
 
+    private int getHiddenApiRestrictionsOffset() {
+        if (hiddenApiRestrictionsOffset == -1) {
+            MapItem mapItem = getMapItemForSection(ItemType.HIDDENAPI_CLASS_DATA_ITEM);
+            if (mapItem != null) {
+                hiddenApiRestrictionsOffset = mapItem.getOffset();
+            } else {
+                hiddenApiRestrictionsOffset = DexWriter.NO_OFFSET;
+            }
+        }
+        return hiddenApiRestrictionsOffset;
+    }
+
     private int readHiddenApiRestrictionsOffset(int classIndex) {
-        if (hiddenApiRestrictionsOffset == DexWriter.NO_OFFSET) {
+        int restrictionsOffset = getHiddenApiRestrictionsOffset();
+        if (restrictionsOffset == DexWriter.NO_OFFSET) {
             return DexWriter.NO_OFFSET;
         }
 
         int offset = dexBuffer.readInt(
-                hiddenApiRestrictionsOffset +
+                restrictionsOffset +
                         HiddenApiClassDataItem.OFFSETS_LIST_OFFSET +
                         classIndex * HiddenApiClassDataItem.OFFSET_ITEM_SIZE);
         if (offset == DexWriter.NO_OFFSET) {
             return DexWriter.NO_OFFSET;
         }
 
-        return hiddenApiRestrictionsOffset + offset;
+        return restrictionsOffset + offset;
     }
 
     public static abstract class OptionalIndexedSection<T> extends IndexedSection<T> {

dexlib2/src/main/java/com/android/tools/smali/dexlib2/dexbacked/ZipDexContainer.java

@@ -45,10 +45,12 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Enumeration;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.TreeMap;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
 import javax.annotation.Nonnull;
@@ -59,9 +61,22 @@
  */
 public class ZipDexContainer implements MultiDexContainer<DexBackedDexFile> {
 
+    /**
+     * Represents a dex entry in a zip container, including the uncompressed CRC-32 checksum
+     * of the zip entry.
+     */
+    public interface ZipDexEntry extends DexEntry<DexBackedDexFile> {
+        /**
+         * Gets the uncompressed CRC-32 checksum of the zip entry containing this dex file.
+         *
+         * @return The uncompressed CRC-32 checksum of the zip entry
+         */
+        long getCrc();
+    }
+
     private final File zipFilePath;
     @Nullable private final Opcodes opcodes;
-    private TreeMap<String, DexBackedDexFile> entries;
+    private Map<String, ZipDexEntry> entries;
 
     /**
      * Constructs a new ZipDexContainer for the given zip file
@@ -75,76 +90,106 @@ public ZipDexContainer(@Nonnull File zipFilePath, @Nullable Opcodes opcodes) {
 
     /**
      * Gets a list of the names of dex files in this zip file.
+     * <p>
+     * The returned list follows a strict deterministic order: it starts with all dex files in
+     * classes.dex, then classes2.dex, etc. sequentially until classesN.dex is missing, followed by
+     * any remaining dex entries in lexicographically sorted order.
      *
      * @return A list of the names of dex files in this zip file
      */
     @Nonnull @Override public List<String> getDexEntryNames() throws IOException {
         return new ArrayList<>(getEntries().keySet());
     }
 
-    private Map<String, DexBackedDexFile> getEntries() throws IOException {
+    private Map<String, ZipDexEntry> getEntries() throws IOException {
         if (entries != null) {
           return entries;
         }
-        entries = new TreeMap<String, DexBackedDexFile>();
+        // LinkedHashMap is used to preserve the strict deterministic insertion ordering of entries.
+        entries = new LinkedHashMap<String, ZipDexEntry>();
         try (ZipFile zipFile = getZipFile()) {
+            Map<String, ZipEntry> validZipEntries = new HashMap<>();
             Enumeration<? extends ZipEntry> entriesEnumeration = zipFile.entries();
 
             while (entriesEnumeration.hasMoreElements()) {
                 ZipEntry entry = entriesEnumeration.nextElement();
-
-                if (!isDex(zipFile, entry)) {
-                    continue;
+                if (isDex(zipFile, entry)) {
+                    validZipEntries.put(entry.getName(), entry);
                 }
+            }
 
-                // There might be several dex files in zip entry since DEX v41.
-                try (InputStream inputStream = zipFile.getInputStream(entry)) {
-                    byte[] buf = InputStreamUtil.toByteArray(inputStream);
-                    for (int offset = 0, i = 1; offset < buf.length; i++) {
-                      DexBackedDexFile dex = new DexBackedDexFile(opcodes, buf, 0, true, offset);
-                      entries.put(entry.getName() + (i > 1 ? ("/" + i) : ""), dex);
-                      offset += dex.getFileSize();
-                    };
-                } catch (ArrayIndexOutOfBoundsException ex) {
-                    // This is expected if the dex file is invalid.
+            for (int n = 1; ; n++) {
+                String targetName = (n == 1) ? "classes.dex" : ("classes" + n + ".dex");
+                ZipEntry entry = validZipEntries.remove(targetName);
+                if (entry == null) {
+                    break;
                 }
+                loadDexFilesFromZipEntry(zipFile, entry);
+            }
+
+            List<String> remainingNames = new ArrayList<>(validZipEntries.keySet());
+            Collections.sort(remainingNames);
+            for (String remainingName : remainingNames) {
+                ZipEntry entry = validZipEntries.get(remainingName);
+                loadDexFilesFromZipEntry(zipFile, entry);
             }
 
             return entries;
         }
     }
 
+    private void loadDexFilesFromZipEntry(@Nonnull ZipFile zipFile, @Nonnull ZipEntry entry)
+            throws IOException {
+        long entryCrc = entry.getCrc();
+        try (InputStream inputStream = zipFile.getInputStream(entry)) {
+            byte[] buf = InputStreamUtil.toByteArray(inputStream);
+            for (int offset = 0, i = 0; offset < buf.length; i++) {
+              DexBackedDexFile dex = new DexBackedDexFile(opcodes, buf, 0, true, offset);
+              boolean hasSingleEntry = dex.isDexContainerFirstEntry() &&
+                      dex.isDexContainerLastEntry();
+              String entryName = entry.getName() + (hasSingleEntry ? "" : ("/" + i));
+              ZipDexEntry dexEntry = new ZipDexEntry() {
+                  @Nonnull
+                  @Override
+                  public String getEntryName() {
+                      return entryName;
+                  }
+
+                  @Nonnull
+                  @Override
+                  public DexBackedDexFile getDexFile() {
+                      return dex;
+                  }
+
+                  @Nonnull
+                  @Override
+                  public MultiDexContainer<DexBackedDexFile> getContainer() {
+                      return ZipDexContainer.this;
+                  }
+
+                  @Override
+                  public long getCrc() {
+                      return entryCrc;
+                  }
+              };
+              entries.put(entryName, dexEntry);
+              offset += dex.getFileSize();
+              if (dex.isDexContainerLastEntry()) {
+                  break;
+              }
+            };
+        }
+    }
+
     /**
      * Loads a dex file from a specific named entry.
      *
      * @param entryName The name of the entry
      * @return A DexEntry, or null if there is no entry with the given name
      * @throws NotADexFile If the entry isn't a dex file
      */
-    @Nullable @Override public DexEntry<DexBackedDexFile> getEntry(@Nonnull String entryName) throws IOException {
-         DexFile dexFile = getEntries().get(entryName);
-         if (dexFile == null) {
-             return null;
-         }
-         return new DexEntry() {
-             @Nonnull
-             @Override
-             public String getEntryName() {
-                 return entryName;
-             }
-
-             @Nonnull
-             @Override
-             public DexFile getDexFile() {
-                 return dexFile;
-             }
-
-             @Nonnull
-             @Override
-             public MultiDexContainer getContainer() {
-                 return ZipDexContainer.this;
-             }
-         };
+    @Nullable @Override public ZipDexEntry getEntry(@Nonnull String entryName) throws IOException {
+         return getEntries().get(entryName);
     }
 
     public boolean isZipFile() {
@@ -158,7 +203,8 @@ public boolean isZipFile() {
         // just eat it
     }
 
-    protected boolean isDex(@Nonnull ZipFile zipFile, @Nonnull ZipEntry zipEntry) throws IOException {
+    protected boolean isDex(@Nonnull ZipFile zipFile, @Nonnull ZipEntry zipEntry)
+            throws IOException {
         try (InputStream inputStream = new BufferedInputStream(zipFile.getInputStream(zipEntry))) {
             DexUtil.verifyDexHeader(inputStream);
         } catch (NotADexFile ex) {

dexlib2/src/main/java/com/android/tools/smali/dexlib2/dexbacked/raw/HeaderItem.java

@@ -43,6 +43,7 @@
 
 public class HeaderItem {
     public static final int ITEM_SIZE = 0x70;
+    public static final int MAGIC_SIZE = 8;
 
     private static final byte[] MAGIC_VALUE = new byte[] { 0x64, 0x65, 0x78, 0x0a, 0x00, 0x00, 0x00, 0x00 };
 
@@ -88,7 +89,7 @@ public class HeaderItem {
     public static final int DATA_START_OFFSET = 108;
 
     public static final int CONTAINER_SIZE_OFFSET = 112;
-    public static final int CONTAINER_OFF_OFFSET = 116;
+    public static final int HEADER_OFFSET_OFFSET = 116;
 
     @Nonnull private DexBackedDexFile dexFile;
 
@@ -267,7 +268,7 @@ public static byte[] getMagicForDexVersion(int dexVersion) {
      * @return True if the magic value is valid
      */
     public static boolean verifyMagic(byte[] buf, int offset) {
-        if (buf.length - offset < 8) {
+        if (offset < 0 || offset > buf.length || buf.length - offset < MAGIC_SIZE) {
             return false;
         }
 
@@ -297,6 +298,9 @@ public static boolean verifyMagic(byte[] buf, int offset) {
      * @return The dex version if the header is valid or -1 if the header is invalid
      */
     public static int getVersion(byte[] buf, int offset) {
+        if (offset < 0 || offset > buf.length || buf.length - offset < MAGIC_SIZE) {
+            throw new DexBackedDexFile.NotADexFile("File is too short");
+        }
         if (!verifyMagic(buf, offset)) {
             return -1;
         }