Disk for custom VM can be composed from partitions

This is needed for Ferrochrome in order to have multiple copies of the
"state" partition per user while keeping to have one copy of the OS
itself (kernel, rootfs, etc.)

Bug: 351973725
Test: N/A Ferrochrome doesn't yet use this feature due to the lack of
support for specifying GUID for the partitions.

Change-Id: If6fde37f0dedd4af9a237cadfcb4ecf43b57ccd2
diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java b/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
index b3c0746..0efaa18 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -687,12 +687,27 @@
         for (int i = 0; i < config.disks.length; i++) {
             config.disks[i] = new DiskImage();
             config.disks[i].writable = customImageConfig.getDisks()[i].isWritable();
+            String diskImagePath = customImageConfig.getDisks()[i].getImagePath();
+            if (diskImagePath != null) {
+                config.disks[i].image =
+                        ParcelFileDescriptor.open(
+                                new File(diskImagePath),
+                                config.disks[i].writable ? MODE_READ_WRITE : MODE_READ_ONLY);
+            }
 
-            config.disks[i].image =
-                    ParcelFileDescriptor.open(
-                            new File(customImageConfig.getDisks()[i].getImagePath()),
-                            config.disks[i].writable ? MODE_READ_WRITE : MODE_READ_ONLY);
-            config.disks[i].partitions = new Partition[0];
+            List<Partition> partitions = new ArrayList<>();
+            for (VirtualMachineCustomImageConfig.Partition p :
+                    customImageConfig.getDisks()[i].getPartitions()) {
+                Partition part = new Partition();
+                part.label = p.name;
+                part.image =
+                        ParcelFileDescriptor.open(
+                                new File(p.imagePath),
+                                p.writable ? MODE_READ_WRITE : MODE_READ_ONLY);
+                part.writable = p.writable;
+                partitions.add(part);
+            }
+            config.disks[i].partitions = partitions.toArray(new Partition[0]);
         }
 
         config.displayConfig =
diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
index 125e01c..aa22fc2 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
@@ -32,6 +32,9 @@
     private static final String KEY_PARAMS = "params";
     private static final String KEY_DISK_WRITABLES = "disk_writables";
     private static final String KEY_DISK_IMAGES = "disk_images";
+    private static final String KEY_PARTITION_LABELS = "partition_labels_";
+    private static final String KEY_PARTITION_IMAGES = "partition_images_";
+    private static final String KEY_PARTITION_WRITABLES = "partition_writables_";
     private static final String KEY_DISPLAY_CONFIG = "display_config";
     private static final String KEY_TOUCH = "touch";
     private static final String KEY_KEYBOARD = "keyboard";
@@ -155,8 +158,20 @@
         if (writables != null && diskImages != null) {
             if (writables.length == diskImages.length) {
                 for (int i = 0; i < writables.length; i++) {
-                    builder.addDisk(
-                            writables[i] ? Disk.RWDisk(diskImages[i]) : Disk.RODisk(diskImages[i]));
+                    String diskImage = diskImages[i];
+                    diskImage = diskImage.equals("") ? null : diskImage;
+                    Disk disk = writables[i] ? Disk.RWDisk(diskImage) : Disk.RODisk(diskImage);
+                    String[] labels =
+                            customImageConfigBundle.getStringArray(KEY_PARTITION_LABELS + i);
+                    String[] images =
+                            customImageConfigBundle.getStringArray(KEY_PARTITION_IMAGES + i);
+                    boolean[] partitionWritables =
+                            customImageConfigBundle.getBooleanArray(KEY_PARTITION_WRITABLES + i);
+                    for (int j = 0; j < labels.length; j++) {
+                        disk.addPartition(
+                                new Partition(labels[j], images[j], partitionWritables[j]));
+                    }
+                    builder.addDisk(disk);
                 }
             }
         }
@@ -189,7 +204,25 @@
             String[] images = new String[disks.length];
             for (int i = 0; i < disks.length; i++) {
                 writables[i] = disks[i].writable;
-                images[i] = disks[i].imagePath;
+                String imagePath = disks[i].imagePath;
+                images[i] = imagePath == null ? "" : imagePath;
+
+                List<String> partitionLabels = new ArrayList<>();
+                List<String> partitionImages = new ArrayList<>();
+                List<Boolean> partitionWritables = new ArrayList<>();
+                for (Partition p : disks[i].getPartitions()) {
+                    partitionLabels.add(p.name);
+                    partitionImages.add(p.imagePath);
+                    partitionWritables.add(p.writable);
+                }
+                pb.putStringArray(KEY_PARTITION_LABELS + i, partitionLabels.toArray(new String[0]));
+                pb.putStringArray(KEY_PARTITION_IMAGES + i, partitionImages.toArray(new String[0]));
+                boolean[] arr = new boolean[partitionWritables.size()];
+                int index = 0;
+                for (Boolean b : partitionWritables) {
+                    arr[index++] = b;
+                }
+                pb.putBooleanArray(KEY_PARTITION_WRITABLES + i, arr);
             }
             pb.putBooleanArray(KEY_DISK_WRITABLES, writables);
             pb.putStringArray(KEY_DISK_IMAGES, images);
@@ -232,10 +265,12 @@
     public static final class Disk {
         private final boolean writable;
         private final String imagePath;
+        private final List<Partition> partitions;
 
         private Disk(boolean writable, String imagePath) {
             this.writable = writable;
             this.imagePath = imagePath;
+            this.partitions = new ArrayList<>();
         }
 
         /** @hide */
@@ -257,6 +292,30 @@
         public String getImagePath() {
             return imagePath;
         }
+
+        /** @hide */
+        public Disk addPartition(Partition p) {
+            this.partitions.add(p);
+            return this;
+        }
+
+        /** @hide */
+        public List<Partition> getPartitions() {
+            return partitions;
+        }
+    }
+
+    /** @hide */
+    public static final class Partition {
+        public final String name;
+        public final String imagePath;
+        public final boolean writable;
+
+        public Partition(String name, String imagePath, boolean writable) {
+            this.name = name;
+            this.imagePath = imagePath;
+            this.writable = writable;
+        }
     }
 
     /** @hide */
diff --git a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
index 258aeea..74c7f0b 100644
--- a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
+++ b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
@@ -131,6 +131,24 @@
                                     VirtualMachineCustomImageConfig.Disk.RODisk(
                                             item.getString("image")));
                         }
+                    } else if (item.has("partitions")) {
+                        boolean writable = item.optBoolean("writable", false);
+                        VirtualMachineCustomImageConfig.Disk disk =
+                                writable
+                                        ? VirtualMachineCustomImageConfig.Disk.RWDisk(null)
+                                        : VirtualMachineCustomImageConfig.Disk.RODisk(null);
+                        JSONArray partitions = item.getJSONArray("partitions");
+                        for (int j = 0; j < partitions.length(); j++) {
+                            JSONObject partition = partitions.getJSONObject(j);
+                            String label = partition.getString("label");
+                            String path = partition.getString("path");
+                            boolean partitionWritable = partition.optBoolean("writable", false);
+                            VirtualMachineCustomImageConfig.Partition p =
+                                    new VirtualMachineCustomImageConfig.Partition(
+                                            label, path, partitionWritable);
+                            disk.addPartition(p);
+                        }
+                        customImageConfigBuilder.addDisk(disk);
                     }
                 }
             }