Merge "Disk for custom VM can be composed from partitions" into main
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);
                     }
                 }
             }