Add test api for setting vendor disk image

So far only a single test asserting that setting vendor disk image
requires caller to have USE_CUSTOM_VIRTUAL_MACHINE permission. More
tests to come in the follow up changes.

Bug: 285855433
Test: atest MicrodroidTestApp
Change-Id: Ic68fccbb8f916f30cab1d772f45714d1e75369a2
diff --git a/javalib/api/test-current.txt b/javalib/api/test-current.txt
index 1298000..cf95770 100644
--- a/javalib/api/test-current.txt
+++ b/javalib/api/test-current.txt
@@ -13,6 +13,7 @@
 
   public static final class VirtualMachineConfig.Builder {
     method @NonNull @RequiresPermission(android.system.virtualmachine.VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION) public android.system.virtualmachine.VirtualMachineConfig.Builder setPayloadConfigPath(@NonNull String);
+    method @NonNull @RequiresPermission(android.system.virtualmachine.VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION) public android.system.virtualmachine.VirtualMachineConfig.Builder setVendorDiskImage(@NonNull java.io.File);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setVmConsoleInputSupported(boolean);
   }
 
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index b400eeb..4cad2e3 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -76,6 +76,7 @@
     private static final String KEY_ENCRYPTED_STORAGE_BYTES = "encryptedStorageBytes";
     private static final String KEY_VM_OUTPUT_CAPTURED = "vmOutputCaptured";
     private static final String KEY_VM_CONSOLE_INPUT_SUPPORTED = "vmConsoleInputSupported";
+    private static final String KEY_VENDOR_DISK_IMAGE_PATH = "vendorDiskImagePath";
 
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
@@ -167,6 +168,8 @@
     /** Whether the app can write console input to the VM */
     private final boolean mVmConsoleInputSupported;
 
+    @Nullable private final File mVendorDiskImage;
+
     private VirtualMachineConfig(
             @Nullable String packageName,
             @Nullable String apkPath,
@@ -178,7 +181,8 @@
             @CpuTopology int cpuTopology,
             long encryptedStorageBytes,
             boolean vmOutputCaptured,
-            boolean vmConsoleInputSupported) {
+            boolean vmConsoleInputSupported,
+            @Nullable File vendorDiskImage) {
         // This is only called from Builder.build(); the builder handles parameter validation.
         mPackageName = packageName;
         mApkPath = apkPath;
@@ -191,6 +195,7 @@
         mEncryptedStorageBytes = encryptedStorageBytes;
         mVmOutputCaptured = vmOutputCaptured;
         mVmConsoleInputSupported = vmConsoleInputSupported;
+        mVendorDiskImage = vendorDiskImage;
     }
 
     /** Loads a config from a file. */
@@ -267,6 +272,11 @@
         builder.setVmOutputCaptured(b.getBoolean(KEY_VM_OUTPUT_CAPTURED));
         builder.setVmConsoleInputSupported(b.getBoolean(KEY_VM_CONSOLE_INPUT_SUPPORTED));
 
+        String vendorDiskImagePath = b.getString(KEY_VENDOR_DISK_IMAGE_PATH);
+        if (vendorDiskImagePath != null) {
+            builder.setVendorDiskImage(new File(vendorDiskImagePath));
+        }
+
         return builder.build();
     }
 
@@ -302,6 +312,9 @@
         }
         b.putBoolean(KEY_VM_OUTPUT_CAPTURED, mVmOutputCaptured);
         b.putBoolean(KEY_VM_CONSOLE_INPUT_SUPPORTED, mVmConsoleInputSupported);
+        if (mVendorDiskImage != null) {
+            b.putString(KEY_VENDOR_DISK_IMAGE_PATH, mVendorDiskImage.getAbsolutePath());
+        }
         b.writeToStream(output);
     }
 
@@ -501,6 +514,20 @@
                 vsConfig.cpuTopology = android.system.virtualizationservice.CpuTopology.ONE_CPU;
                 break;
         }
+        if (mVendorDiskImage != null) {
+            VirtualMachineAppConfig.CustomConfig customConfig =
+                    new VirtualMachineAppConfig.CustomConfig();
+            customConfig.taskProfiles = new String[0];
+            try {
+                customConfig.vendorImage =
+                        ParcelFileDescriptor.open(mVendorDiskImage, MODE_READ_ONLY);
+            } catch (FileNotFoundException e) {
+                throw new VirtualMachineException(
+                        "Failed to open vendor disk image " + mVendorDiskImage.getAbsolutePath(),
+                        e);
+            }
+            vsConfig.customConfig = customConfig;
+        }
         return vsConfig;
     }
 
@@ -572,6 +599,7 @@
         private long mEncryptedStorageBytes;
         private boolean mVmOutputCaptured = false;
         private boolean mVmConsoleInputSupported = false;
+        @Nullable private File mVendorDiskImage;
 
         /**
          * Creates a builder for the given context.
@@ -645,7 +673,8 @@
                     mCpuTopology,
                     mEncryptedStorageBytes,
                     mVmOutputCaptured,
-                    mVmConsoleInputSupported);
+                    mVmConsoleInputSupported,
+                    mVendorDiskImage);
         }
 
         /**
@@ -863,5 +892,18 @@
             mVmConsoleInputSupported = supported;
             return this;
         }
+
+        /**
+         * Sets the path to the disk image with vendor-specific modules.
+         *
+         * @hide
+         */
+        @TestApi
+        @RequiresPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION)
+        @NonNull
+        public Builder setVendorDiskImage(@NonNull File vendorDiskImage) {
+            mVendorDiskImage = vendorDiskImage;
+            return this;
+        }
     }
 }
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index fe8f5c9..8a31c21 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -43,7 +43,10 @@
     ],
     min_sdk_version: "33",
     // Defined in ../vmshareapp/Android.bp
-    data: [":MicrodroidVmShareApp"],
+    data: [
+        ":MicrodroidVmShareApp",
+        ":test_microdroid_vendor_image",
+    ],
 }
 
 // Defaults shared between MicrodroidTestNativeLib and MicrodroidPayloadInOtherAppNativeLib shared
diff --git a/tests/testapk/AndroidManifest.xml b/tests/testapk/AndroidManifest.xml
index 2ea3f6c..d6e6004 100644
--- a/tests/testapk/AndroidManifest.xml
+++ b/tests/testapk/AndroidManifest.xml
@@ -22,8 +22,7 @@
     <queries>
         <package android:name="com.android.microdroid.vmshare_app" />
     </queries>
-    <application>
-    </application>
+    <application />
     <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
         android:targetPackage="com.android.microdroid.test"
         android:label="Microdroid Test" />
diff --git a/tests/testapk/AndroidTest.xml b/tests/testapk/AndroidTest.xml
index 929dd31..e72a2e3 100644
--- a/tests/testapk/AndroidTest.xml
+++ b/tests/testapk/AndroidTest.xml
@@ -23,6 +23,14 @@
         <option name="test-file-name" value="MicrodroidTestApp.apk" />
         <option name="test-file-name" value="MicrodroidVmShareApp.apk" />
     </target_preparer>
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="run-command" value="mkdir -p /data/local/tmp/cts/microdroid" />
+        <option name="teardown-command" value="rm -rf /data/local/tmp/cts/microdroid" />
+    </target_preparer>
+    <target_preparer class="com.android.compatibility.common.tradefed.targetprep.FilePusher">
+        <option name="cleanup" value="true" />
+        <option name="push" value="test_microdroid_vendor_image.img->/data/local/tmp/cts/microdroid/test_microdroid_vendor_image.img" />
+    </target_preparer>
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="com.android.microdroid.test" />
         <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
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 dc045a1..4cebd4c 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -2016,6 +2016,30 @@
         }
     }
 
+    @Test
+    public void configuringVendorDiskImageRequiresCustomPermission() throws Exception {
+        assumeSupportedDevice();
+
+        File vendorDiskImage =
+                new File("/data/local/tmp/cts/microdroid/test_microdroid_vendor_image.img");
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryName("MicrodroidTestNativeLib.so")
+                        .setVendorDiskImage(vendorDiskImage)
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+
+        VirtualMachine vm =
+                forceCreateNewVirtualMachine("test_vendor_image_req_custom_permission", config);
+
+        SecurityException e =
+                assertThrows(
+                        SecurityException.class, () -> runVmTestService(TAG, vm, (ts, tr) -> {}));
+        assertThat(e)
+                .hasMessageThat()
+                .contains("android.permission.USE_CUSTOM_VIRTUAL_MACHINE permission");
+    }
+
     private static class VmShareServiceConnection implements ServiceConnection {
 
         private final CountDownLatch mLatch = new CountDownLatch(1);