Merge "Cache package name not APK path"
diff --git a/javalib/api/system-current.txt b/javalib/api/system-current.txt
index 1977321..fe9943d 100644
--- a/javalib/api/system-current.txt
+++ b/javalib/api/system-current.txt
@@ -56,7 +56,7 @@
   }
 
   public final class VirtualMachineConfig {
-    method @NonNull public String getApkPath();
+    method @Nullable public String getApkPath();
     method @NonNull public int getDebugLevel();
     method @IntRange(from=0) public long getEncryptedStorageKib();
     method @IntRange(from=0) public int getMemoryMib();
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index e66cf29..7c7f4b5 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -779,7 +779,8 @@
                     createVmPipes();
                 }
 
-                VirtualMachineAppConfig appConfig = getConfig().toVsConfig();
+                VirtualMachineAppConfig appConfig =
+                        getConfig().toVsConfig(mContext.getPackageManager());
                 appConfig.name = mName;
 
                 try {
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index 1fd49c8..b358f9e 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -29,6 +29,8 @@
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
 import android.content.Context;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
 import android.os.ParcelFileDescriptor;
 import android.os.PersistableBundle;
 import android.sysprop.HypervisorProperties;
@@ -58,8 +60,9 @@
     private static final String[] EMPTY_STRING_ARRAY = {};
 
     // These define the schema of the config file persisted on disk.
-    private static final int VERSION = 3;
+    private static final int VERSION = 4;
     private static final String KEY_VERSION = "version";
+    private static final String KEY_PACKAGENAME = "packageName";
     private static final String KEY_APKPATH = "apkPath";
     private static final String KEY_PAYLOADCONFIGPATH = "payloadConfigPath";
     private static final String KEY_PAYLOADBINARYNAME = "payloadBinaryPath";
@@ -94,8 +97,11 @@
      */
     @SystemApi public static final int DEBUG_LEVEL_FULL = 1;
 
+    /** Name of a package whose primary APK contains the VM payload. */
+    @Nullable private final String mPackageName;
+
     /** Absolute path to the APK file containing the VM payload. */
-    @NonNull private final String mApkPath;
+    @Nullable private final String mApkPath;
 
     @DebugLevel private final int mDebugLevel;
 
@@ -129,7 +135,8 @@
     private final boolean mVmOutputCaptured;
 
     private VirtualMachineConfig(
-            @NonNull String apkPath,
+            @Nullable String packageName,
+            @Nullable String apkPath,
             @Nullable String payloadConfigPath,
             @Nullable String payloadBinaryName,
             @DebugLevel int debugLevel,
@@ -139,6 +146,7 @@
             long encryptedStorageKib,
             boolean vmOutputCaptured) {
         // This is only called from Builder.build(); the builder handles parameter validation.
+        mPackageName = packageName;
         mApkPath = apkPath;
         mPayloadConfigPath = payloadConfigPath;
         mPayloadBinaryName = payloadBinaryName;
@@ -191,8 +199,13 @@
                     "Version " + version + " too high; current is " + VERSION);
         }
 
-        Builder builder = new Builder();
-        builder.setApkPath(b.getString(KEY_APKPATH));
+        String packageName = b.getString(KEY_PACKAGENAME);
+        Builder builder = new Builder(packageName);
+
+        String apkPath = b.getString(KEY_APKPATH);
+        if (apkPath != null) {
+            builder.setApkPath(apkPath);
+        }
 
         String payloadConfigPath = b.getString(KEY_PAYLOADCONFIGPATH);
         if (payloadConfigPath == null) {
@@ -234,7 +247,12 @@
     private void serializeOutputStream(@NonNull OutputStream output) throws IOException {
         PersistableBundle b = new PersistableBundle();
         b.putInt(KEY_VERSION, VERSION);
-        b.putString(KEY_APKPATH, mApkPath);
+        if (mPackageName != null) {
+            b.putString(KEY_PACKAGENAME, mPackageName);
+        }
+        if (mApkPath != null) {
+            b.putString(KEY_APKPATH, mApkPath);
+        }
         b.putString(KEY_PAYLOADCONFIGPATH, mPayloadConfigPath);
         b.putString(KEY_PAYLOADBINARYNAME, mPayloadBinaryName);
         b.putInt(KEY_DEBUGLEVEL, mDebugLevel);
@@ -252,12 +270,13 @@
 
     /**
      * Returns the absolute path of the APK which should contain the binary payload that will
-     * execute within the VM.
+     * execute within the VM. Returns null if no specific path has been set, so the primary APK will
+     * be used.
      *
      * @hide
      */
     @SystemApi
-    @NonNull
+    @Nullable
     public String getApkPath() {
         return mApkPath;
     }
@@ -383,7 +402,8 @@
                 && this.mVmOutputCaptured == other.mVmOutputCaptured
                 && Objects.equals(this.mPayloadConfigPath, other.mPayloadConfigPath)
                 && Objects.equals(this.mPayloadBinaryName, other.mPayloadBinaryName)
-                && this.mApkPath.equals(other.mApkPath);
+                && Objects.equals(this.mPackageName, other.mPackageName)
+                && Objects.equals(this.mApkPath, other.mApkPath);
     }
 
     /**
@@ -393,11 +413,25 @@
      * app-owned files and that could be abused to run a VM with software that the calling
      * application doesn't own.
      */
-    VirtualMachineAppConfig toVsConfig() throws VirtualMachineException {
+    VirtualMachineAppConfig toVsConfig(@NonNull PackageManager packageManager)
+            throws VirtualMachineException {
         VirtualMachineAppConfig vsConfig = new VirtualMachineAppConfig();
 
+        String apkPath = mApkPath;
+        if (apkPath == null) {
+            try {
+                ApplicationInfo appInfo =
+                        packageManager.getApplicationInfo(
+                                mPackageName, PackageManager.ApplicationInfoFlags.of(0));
+                // This really is the path to the APK, not a directory.
+                apkPath = appInfo.sourceDir;
+            } catch (PackageManager.NameNotFoundException e) {
+                throw new VirtualMachineException("Package not found", e);
+            }
+        }
+
         try {
-            vsConfig.apk = ParcelFileDescriptor.open(new File(mApkPath), MODE_READ_ONLY);
+            vsConfig.apk = ParcelFileDescriptor.open(new File(apkPath), MODE_READ_ONLY);
         } catch (FileNotFoundException e) {
             throw new VirtualMachineException("Failed to open APK", e);
         }
@@ -433,7 +467,7 @@
      */
     @SystemApi
     public static final class Builder {
-        @Nullable private final Context mContext;
+        @Nullable private final String mPackageName;
         @Nullable private String mApkPath;
         @Nullable private String mPayloadConfigPath;
         @Nullable private String mPayloadBinaryName;
@@ -452,15 +486,15 @@
          */
         @SystemApi
         public Builder(@NonNull Context context) {
-            mContext = requireNonNull(context, "context must not be null");
+            mPackageName = requireNonNull(context, "context must not be null").getPackageName();
         }
 
         /**
-         * Creates a builder with no associated context; {@link #setApkPath} must be called to
-         * specify which APK contains the payload.
+         * Creates a builder for a specific package. If packageName is null, {@link #setApkPath}
+         * must be called to specify the APK containing the payload.
          */
-        private Builder() {
-            mContext = null;
+        private Builder(@Nullable String packageName) {
+            mPackageName = packageName;
         }
 
         /**
@@ -471,14 +505,16 @@
         @SystemApi
         @NonNull
         public VirtualMachineConfig build() {
-            String apkPath;
-            if (mApkPath == null) {
-                if (mContext == null) {
-                    throw new IllegalStateException("apkPath must be specified");
-                }
-                apkPath = mContext.getPackageCodePath();
-            } else {
+            String apkPath = null;
+            String packageName = null;
+
+            if (mApkPath != null) {
                 apkPath = mApkPath;
+            } else if (mPackageName != null) {
+                packageName = mPackageName;
+            } else {
+                // This should never happen, unless we're deserializing a bad config
+                throw new IllegalStateException("apkPath or packageName must be specified");
             }
 
             if (mPayloadBinaryName == null) {
@@ -501,6 +537,7 @@
             }
 
             return new VirtualMachineConfig(
+                    packageName,
                     apkPath,
                     mPayloadConfigPath,
                     mPayloadBinaryName,
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index e1a2e40..7bd5f08 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -31,6 +31,7 @@
 import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
 
 import android.content.Context;
+import android.content.ContextWrapper;
 import android.os.Build;
 import android.os.ParcelFileDescriptor;
 import android.os.ParcelFileDescriptor.AutoCloseInputStream;
@@ -342,7 +343,7 @@
         VirtualMachineConfig.Builder minimalBuilder = newVmConfigBuilder();
         VirtualMachineConfig minimal = minimalBuilder.setPayloadBinaryName("binary.so").build();
 
-        assertThat(minimal.getApkPath()).isEqualTo(getContext().getPackageCodePath());
+        assertThat(minimal.getApkPath()).isNull();
         assertThat(minimal.getDebugLevel()).isEqualTo(DEBUG_LEVEL_NONE);
         assertThat(minimal.getMemoryMib()).isEqualTo(0);
         assertThat(minimal.getNumCpus()).isEqualTo(1);
@@ -425,13 +426,9 @@
         assertThat(e).hasMessageThat().contains("debug level must be FULL to capture output");
     }
 
-    private VirtualMachineConfig.Builder newBaselineBuilder() {
-        return newVmConfigBuilder().setPayloadBinaryName("binary.so").setApkPath("/apk/path");
-    }
-
     @Test
     @CddTest(requirements = {"9.17/C-1-1"})
-    public void compatibleConfigTests() throws Exception {
+    public void compatibleConfigTests() {
         int maxCpus = Runtime.getRuntime().availableProcessors();
 
         VirtualMachineConfig baseline = newBaselineBuilder().build();
@@ -467,6 +464,31 @@
                 newBaselineBuilder().setDebugLevel(DEBUG_LEVEL_FULL);
         VirtualMachineConfig debuggable = debuggableBuilder.build();
         assertConfigCompatible(debuggable, debuggableBuilder.setVmOutputCaptured(true)).isFalse();
+
+        VirtualMachineConfig currentContextConfig =
+                new VirtualMachineConfig.Builder(getContext())
+                        .setProtectedVm(isProtectedVm())
+                        .setPayloadBinaryName("binary.so")
+                        .build();
+
+        // packageName is not directly exposed by the config, so we have to be a bit creative
+        // to modify it.
+        Context otherContext =
+                new ContextWrapper(getContext()) {
+                    @Override
+                    public String getPackageName() {
+                        return "other.package.name";
+                    }
+                };
+        VirtualMachineConfig.Builder otherContextBuilder =
+                new VirtualMachineConfig.Builder(otherContext)
+                        .setProtectedVm(isProtectedVm())
+                        .setPayloadBinaryName("binary.so");
+        assertConfigCompatible(currentContextConfig, otherContextBuilder).isFalse();
+    }
+
+    private VirtualMachineConfig.Builder newBaselineBuilder() {
+        return newVmConfigBuilder().setPayloadBinaryName("binary.so").setApkPath("/apk/path");
     }
 
     private BooleanSubject assertConfigCompatible(