Allow multidex of instrumentations

Allow installation of instrumentation secondary dex files. Both
the instrumentation and the instrumented application are installed (if
necessary) by MultiDex.installInstrumentation. Instrumentation
secondary dex files are extracted in the Application code cache folder
because it generally doesn't have access to its own folder.
Instrumentation preferences are saved in the Application preferences
with a prefix.

Bug: 31383194

Test: frameworks/base/core/tests/hosttests/test-apps/MultiDexLegacy*

Change-Id: I705ed87162326fd64128454aa144a359b09436cd
diff --git a/instrumentation/src/com/android/test/runner/MultiDexTestRunner.java b/instrumentation/src/com/android/test/runner/MultiDexTestRunner.java
index 3de3632..d6931d1 100644
--- a/instrumentation/src/com/android/test/runner/MultiDexTestRunner.java
+++ b/instrumentation/src/com/android/test/runner/MultiDexTestRunner.java
@@ -32,7 +32,7 @@
 
     @Override
     public void onCreate(Bundle arguments) {
-        MultiDex.install(getTargetContext());
+        MultiDex.installInstrumentation(getContext(), getTargetContext());
         super.onCreate(arguments);
     }
 
diff --git a/library/src/android/support/multidex/MultiDex.java b/library/src/android/support/multidex/MultiDex.java
index d35da96..ab7f668 100644
--- a/library/src/android/support/multidex/MultiDex.java
+++ b/library/src/android/support/multidex/MultiDex.java
@@ -17,6 +17,7 @@
 package android.support.multidex;
 
 import android.app.Application;
+import android.app.Instrumentation;
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.os.Build;
@@ -70,7 +71,9 @@
 
     private static final int VM_WITH_MULTIDEX_VERSION_MINOR = 1;
 
-    private static final Set<String> installedApk = new HashSet<String>();
+    private static final String NO_KEY_PREFIX = "";
+
+    private static final Set<File> installedApk = new HashSet<File>();
 
     private static final boolean IS_VM_MULTIDEX_CAPABLE =
             isVMMultidexCapable(System.getProperty("java.vm.version"));
@@ -88,83 +91,164 @@
      *         extension.
      */
     public static void install(Context context) {
-        Log.i(TAG, "install");
+        Log.i(TAG, "Installing application");
         if (IS_VM_MULTIDEX_CAPABLE) {
             Log.i(TAG, "VM has multidex support, MultiDex support library is disabled.");
             return;
         }
 
         if (Build.VERSION.SDK_INT < MIN_SDK_VERSION) {
-            throw new RuntimeException("Multi dex installation failed. SDK " + Build.VERSION.SDK_INT
+            throw new RuntimeException("MultiDex installation failed. SDK " + Build.VERSION.SDK_INT
                     + " is unsupported. Min SDK version is " + MIN_SDK_VERSION + ".");
         }
 
         try {
             ApplicationInfo applicationInfo = getApplicationInfo(context);
             if (applicationInfo == null) {
-                // Looks like running on a test Context, so just return without patching.
+              Log.i(TAG, "No ApplicationInfo available, i.e. running on a test Context:"
+                  + " MultiDex support library is disabled.");
+              return;
+            }
+
+            doInstallation(context,
+                    new File(applicationInfo.sourceDir),
+                    new File(applicationInfo.dataDir),
+                    CODE_CACHE_SECONDARY_FOLDER_NAME,
+                    NO_KEY_PREFIX);
+
+        } catch (Exception e) {
+            Log.e(TAG, "MultiDex installation failure", e);
+            throw new RuntimeException("MultiDex installation failed (" + e.getMessage() + ").");
+        }
+        Log.i(TAG, "install done");
+    }
+
+    /**
+     * Patches the instrumentation context class loader by appending extra dex files
+     * loaded from the instrumentation apk and the application apk. This method should be called in
+     * the onCreate of your {@link Instrumentation}, see
+     * {@link com.android.test.runner.MultiDexTestRunner} for an example.
+     *
+     * @param instrumentationContext instrumentation context.
+     * @param targetContext target application context.
+     * @throws RuntimeException if an error occurred preventing the classloader
+     *         extension.
+     */
+    public static void installInstrumentation(Context instrumentationContext,
+            Context targetContext) {
+        Log.i(TAG, "Installing instrumentation");
+
+        if (IS_VM_MULTIDEX_CAPABLE) {
+            Log.i(TAG, "VM has multidex support, MultiDex support library is disabled.");
+            return;
+        }
+
+        if (Build.VERSION.SDK_INT < MIN_SDK_VERSION) {
+            throw new RuntimeException("MultiDex installation failed. SDK " + Build.VERSION.SDK_INT
+                    + " is unsupported. Min SDK version is " + MIN_SDK_VERSION + ".");
+        }
+        try {
+
+            ApplicationInfo instrumentationInfo = getApplicationInfo(instrumentationContext);
+            if (instrumentationInfo == null) {
+                Log.i(TAG, "No ApplicationInfo available for instrumentation, i.e. running on a"
+                    + " test Context: MultiDex support library is disabled.");
                 return;
             }
 
-            synchronized (installedApk) {
-                String apkPath = applicationInfo.sourceDir;
-                if (installedApk.contains(apkPath)) {
-                    return;
-                }
-                installedApk.add(apkPath);
-
-                if (Build.VERSION.SDK_INT > MAX_SUPPORTED_SDK_VERSION) {
-                    Log.w(TAG, "MultiDex is not guaranteed to work in SDK version "
-                            + Build.VERSION.SDK_INT + ": SDK version higher than "
-                            + MAX_SUPPORTED_SDK_VERSION + " should be backed by "
-                            + "runtime with built-in multidex capabilty but it's not the "
-                            + "case here: java.vm.version=\""
-                            + System.getProperty("java.vm.version") + "\"");
-                }
-
-                /* The patched class loader is expected to be a descendant of
-                 * dalvik.system.BaseDexClassLoader. We modify its
-                 * dalvik.system.DexPathList pathList field to append additional DEX
-                 * file entries.
-                 */
-                ClassLoader loader;
-                try {
-                    loader = context.getClassLoader();
-                } catch (RuntimeException e) {
-                    /* Ignore those exceptions so that we don't break tests relying on Context like
-                     * a android.test.mock.MockContext or a android.content.ContextWrapper with a
-                     * null base Context.
-                     */
-                    Log.w(TAG, "Failure while trying to obtain Context class loader. " +
-                            "Must be running in test mode. Skip patching.", e);
-                    return;
-                }
-                if (loader == null) {
-                    // Note, the context class loader is null when running Robolectric tests.
-                    Log.e(TAG,
-                            "Context class loader is null. Must be running in test mode. "
-                            + "Skip patching.");
-                    return;
-                }
-
-                try {
-                  clearOldDexDir(context);
-                } catch (Throwable t) {
-                  Log.w(TAG, "Something went wrong when trying to clear old MultiDex extraction, "
-                      + "continuing without cleaning.", t);
-                }
-
-                File dexDir = getDexDir(context, applicationInfo);
-                List<? extends File> files =
-                    MultiDexExtractor.load(context, applicationInfo, dexDir, false);
-                installSecondaryDexes(loader, dexDir, files);
+            ApplicationInfo applicationInfo = getApplicationInfo(targetContext);
+            if (applicationInfo == null) {
+                Log.i(TAG, "No ApplicationInfo available, i.e. running on a test Context:"
+                    + " MultiDex support library is disabled.");
+                return;
             }
 
+            String instrumentationPrefix = instrumentationContext.getPackageName() + ".";
+
+            File dataDir = new File(applicationInfo.dataDir);
+
+            doInstallation(targetContext,
+                    new File(instrumentationInfo.sourceDir),
+                    dataDir,
+                    instrumentationPrefix + CODE_CACHE_SECONDARY_FOLDER_NAME,
+                    instrumentationPrefix);
+
+            doInstallation(targetContext,
+                    new File(applicationInfo.sourceDir),
+                    dataDir,
+                    CODE_CACHE_SECONDARY_FOLDER_NAME,
+                    NO_KEY_PREFIX);
         } catch (Exception e) {
-            Log.e(TAG, "Multidex installation failure", e);
-            throw new RuntimeException("Multi dex installation failed (" + e.getMessage() + ").");
+            Log.e(TAG, "MultiDex installation failure", e);
+            throw new RuntimeException("MultiDex installation failed (" + e.getMessage() + ").");
         }
-        Log.i(TAG, "install done");
+        Log.i(TAG, "Installation done");
+    }
+
+    /**
+     * @param mainContext context used to get filesDir, to save preference and to get the
+     * classloader to patch.
+     * @param sourceApk Apk file.
+     * @param dataDir data directory to use for code cache simulation.
+     * @param secondaryFolderName name of the folder for storing extractions.
+     * @param prefsKeyPrefix prefix of all stored preference keys.
+     */
+    private static void doInstallation(Context mainContext, File sourceApk, File dataDir,
+            String secondaryFolderName, String prefsKeyPrefix) throws IOException,
+                IllegalArgumentException, IllegalAccessException, NoSuchFieldException,
+                InvocationTargetException, NoSuchMethodException {
+        synchronized (installedApk) {
+            if (installedApk.contains(sourceApk)) {
+                return;
+            }
+            installedApk.add(sourceApk);
+
+            if (Build.VERSION.SDK_INT > MAX_SUPPORTED_SDK_VERSION) {
+                Log.w(TAG, "MultiDex is not guaranteed to work in SDK version "
+                        + Build.VERSION.SDK_INT + ": SDK version higher than "
+                        + MAX_SUPPORTED_SDK_VERSION + " should be backed by "
+                        + "runtime with built-in multidex capabilty but it's not the "
+                        + "case here: java.vm.version=\""
+                        + System.getProperty("java.vm.version") + "\"");
+            }
+
+            /* The patched class loader is expected to be a descendant of
+             * dalvik.system.BaseDexClassLoader. We modify its
+             * dalvik.system.DexPathList pathList field to append additional DEX
+             * file entries.
+             */
+            ClassLoader loader;
+            try {
+                loader = mainContext.getClassLoader();
+            } catch (RuntimeException e) {
+                /* Ignore those exceptions so that we don't break tests relying on Context like
+                 * a android.test.mock.MockContext or a android.content.ContextWrapper with a
+                 * null base Context.
+                 */
+                Log.w(TAG, "Failure while trying to obtain Context class loader. " +
+                        "Must be running in test mode. Skip patching.", e);
+                return;
+            }
+            if (loader == null) {
+                // Note, the context class loader is null when running Robolectric tests.
+                Log.e(TAG,
+                        "Context class loader is null. Must be running in test mode. "
+                        + "Skip patching.");
+                return;
+            }
+
+            try {
+              clearOldDexDir(mainContext);
+            } catch (Throwable t) {
+              Log.w(TAG, "Something went wrong when trying to clear old MultiDex extraction, "
+                  + "continuing without cleaning.", t);
+            }
+
+            File dexDir = getDexDir(mainContext, dataDir, secondaryFolderName);
+            List<? extends File> files =
+                    MultiDexExtractor.load(mainContext, sourceApk, dexDir, prefsKeyPrefix, false);
+            installSecondaryDexes(loader, dexDir, files);
+        }
     }
 
     private static ApplicationInfo getApplicationInfo(Context context) {
@@ -335,9 +419,9 @@
         }
     }
 
-    private static File getDexDir(Context context, ApplicationInfo applicationInfo)
+    private static File getDexDir(Context context, File dataDir, String secondaryFolderName)
             throws IOException {
-        File cache = new File(applicationInfo.dataDir, CODE_CACHE_NAME);
+        File cache = new File(dataDir, CODE_CACHE_NAME);
         try {
             mkdirChecked(cache);
         } catch (IOException e) {
@@ -348,7 +432,7 @@
             cache = new File(context.getFilesDir(), CODE_CACHE_NAME);
             mkdirChecked(cache);
         }
-        File dexDir = new File(cache, CODE_CACHE_SECONDARY_FOLDER_NAME);
+        File dexDir = new File(cache, secondaryFolderName);
         mkdirChecked(dexDir);
         return dexDir;
     }
diff --git a/library/src/android/support/multidex/MultiDexExtractor.java b/library/src/android/support/multidex/MultiDexExtractor.java
index 2d7402a..39b6bf7 100644
--- a/library/src/android/support/multidex/MultiDexExtractor.java
+++ b/library/src/android/support/multidex/MultiDexExtractor.java
@@ -18,7 +18,6 @@
 
 import android.content.Context;
 import android.content.SharedPreferences;
-import android.content.pm.ApplicationInfo;
 import android.os.Build;
 import android.util.Log;
 import java.io.BufferedOutputStream;
@@ -93,10 +92,11 @@
      * @throws IOException if encounters a problem while reading or writing
      *         secondary dex files
      */
-    static List<? extends File> load(Context context, ApplicationInfo applicationInfo, File dexDir,
+    static List<? extends File> load(Context context, File sourceApk, File dexDir,
+            String prefsKeyPrefix,
             boolean forceReload) throws IOException {
-        Log.i(TAG, "MultiDexExtractor.load(" + applicationInfo.sourceDir + ", " + forceReload + ")");
-        final File sourceApk = new File(applicationInfo.sourceDir);
+        Log.i(TAG, "MultiDexExtractor.load(" + sourceApk.getPath() + ", " + forceReload + ", " +
+                prefsKeyPrefix + ")");
 
         long currentCrc = getZipCrc(sourceApk);
 
@@ -113,19 +113,21 @@
             cacheLock = lockChannel.lock();
             Log.i(TAG, lockFile.getPath() + " locked");
 
-            if (!forceReload && !isModified(context, sourceApk, currentCrc)) {
+            if (!forceReload && !isModified(context, sourceApk, currentCrc, prefsKeyPrefix)) {
                 try {
-                    files = loadExistingExtractions(context, sourceApk, dexDir);
+                    files = loadExistingExtractions(context, sourceApk, dexDir, prefsKeyPrefix);
                 } catch (IOException ioe) {
                     Log.w(TAG, "Failed to reload existing extracted secondary dex files,"
                             + " falling back to fresh extraction", ioe);
                     files = performExtractions(sourceApk, dexDir);
-                    putStoredApkInfo(context, getTimeStamp(sourceApk), currentCrc, files);
+                    putStoredApkInfo(context, prefsKeyPrefix, getTimeStamp(sourceApk), currentCrc,
+                            files);
                 }
             } else {
                 Log.i(TAG, "Detected that extraction must be performed.");
                 files = performExtractions(sourceApk, dexDir);
-                putStoredApkInfo(context, getTimeStamp(sourceApk), currentCrc, files);
+                putStoredApkInfo(context, prefsKeyPrefix, getTimeStamp(sourceApk), currentCrc,
+                        files);
             }
         } finally {
             if (cacheLock != null) {
@@ -157,13 +159,14 @@
      * {@link #LOCK_FILENAME}.
      */
     private static List<ExtractedDex> loadExistingExtractions(
-            Context context, File sourceApk, File dexDir)
+            Context context, File sourceApk, File dexDir,
+            String prefsKeyPrefix)
             throws IOException {
         Log.i(TAG, "loading existing secondary dex files");
 
         final String extractedFilePrefix = sourceApk.getName() + EXTRACTED_NAME_EXT;
         SharedPreferences multiDexPreferences = getMultiDexPreferences(context);
-        int totalDexNumber = multiDexPreferences.getInt(KEY_DEX_NUMBER, 1);
+        int totalDexNumber = multiDexPreferences.getInt(prefsKeyPrefix + KEY_DEX_NUMBER, 1);
         final List<ExtractedDex> files = new ArrayList<ExtractedDex>(totalDexNumber - 1);
 
         for (int secondaryNumber = 2; secondaryNumber <= totalDexNumber; secondaryNumber++) {
@@ -171,15 +174,15 @@
             ExtractedDex extractedFile = new ExtractedDex(dexDir, fileName);
             if (extractedFile.isFile()) {
                 extractedFile.crc = getZipCrc(extractedFile);
-                long expectedCrc =
-                        multiDexPreferences.getLong(KEY_DEX_CRC + secondaryNumber, NO_VALUE);
-                long expectedModTime =
-                        multiDexPreferences.getLong(KEY_DEX_TIME + secondaryNumber, NO_VALUE);
+                long expectedCrc = multiDexPreferences.getLong(
+                        prefsKeyPrefix + KEY_DEX_CRC + secondaryNumber, NO_VALUE);
+                long expectedModTime = multiDexPreferences.getLong(
+                        prefsKeyPrefix + KEY_DEX_TIME + secondaryNumber, NO_VALUE);
                 long lastModified = extractedFile.lastModified();
                 if ((expectedModTime != lastModified)
                         || (expectedCrc != extractedFile.crc)) {
                     throw new IOException("Invalid extracted dex: " + extractedFile +
-                            ", expected modification time: "
+                            " (key \"" + prefsKeyPrefix + "\"), expected modification time: "
                             + expectedModTime + ", modification time: "
                             + lastModified + ", expected crc: "
                             + expectedCrc + ", file crc: " + extractedFile.crc);
@@ -199,10 +202,11 @@
      * Compare current archive and crc with values stored in {@link SharedPreferences}. Should be
      * called only while owning the lock on {@link #LOCK_FILENAME}.
      */
-    private static boolean isModified(Context context, File archive, long currentCrc) {
+    private static boolean isModified(Context context, File archive, long currentCrc,
+            String prefsKeyPrefix) {
         SharedPreferences prefs = getMultiDexPreferences(context);
-        return (prefs.getLong(KEY_TIME_STAMP, NO_VALUE) != getTimeStamp(archive))
-                || (prefs.getLong(KEY_CRC, NO_VALUE) != currentCrc);
+        return (prefs.getLong(prefsKeyPrefix + KEY_TIME_STAMP, NO_VALUE) != getTimeStamp(archive))
+                || (prefs.getLong(prefsKeyPrefix + KEY_CRC, NO_VALUE) != currentCrc);
     }
 
     private static long getTimeStamp(File archive) {
@@ -303,18 +307,18 @@
      * Save {@link SharedPreferences}. Should be called only while owning the lock on
      * {@link #LOCK_FILENAME}.
      */
-    private static void putStoredApkInfo(Context context, long timeStamp, long crc,
-            List<ExtractedDex> extractedDexes) {
+    private static void putStoredApkInfo(Context context, String keyPrefix, long timeStamp,
+            long crc, List<ExtractedDex> extractedDexes) {
         SharedPreferences prefs = getMultiDexPreferences(context);
         SharedPreferences.Editor edit = prefs.edit();
-        edit.putLong(KEY_TIME_STAMP, timeStamp);
-        edit.putLong(KEY_CRC, crc);
-        edit.putInt(KEY_DEX_NUMBER, extractedDexes.size() + 1);
+        edit.putLong(keyPrefix + KEY_TIME_STAMP, timeStamp);
+        edit.putLong(keyPrefix + KEY_CRC, crc);
+        edit.putInt(keyPrefix + KEY_DEX_NUMBER, extractedDexes.size() + 1);
 
         int extractedDexId = 2;
         for (ExtractedDex dex : extractedDexes) {
-            edit.putLong(KEY_DEX_CRC + extractedDexId, dex.crc);
-            edit.putLong(KEY_DEX_TIME + extractedDexId, dex.lastModified());
+            edit.putLong(keyPrefix + KEY_DEX_CRC + extractedDexId, dex.crc);
+            edit.putLong(keyPrefix + KEY_DEX_TIME + extractedDexId, dex.lastModified());
             extractedDexId++;
         }
         /* Use commit() and not apply() as advised by the doc because we need synchronous writing of