Add support for encrypted storage expansion

Capability to configure the encrypted storage size
Partition resizing to the required size upon boot
New unit tests to validate this functionality

Bug: 381067202
Test: atest MicrodroidTests
Change-Id: I6f5737ee601e7c511bdd316b180bf50e3d102ab1
diff --git a/android/virtmgr/src/aidl.rs b/android/virtmgr/src/aidl.rs
index 6aecc75..e954d21 100644
--- a/android/virtmgr/src/aidl.rs
+++ b/android/virtmgr/src/aidl.rs
@@ -303,6 +303,34 @@
         Ok(())
     }
 
+    fn setEncryptedStorageSize(
+        &self,
+        image_fd: &ParcelFileDescriptor,
+        size: i64,
+    ) -> binder::Result<()> {
+        check_manage_access()?;
+
+        let size = size
+            .try_into()
+            .with_context(|| format!("Invalid size: {}", size))
+            .or_binder_exception(ExceptionCode::ILLEGAL_ARGUMENT)?;
+        let size = round_up(size, PARTITION_GRANULARITY_BYTES);
+
+        let image = clone_file(image_fd)?;
+        let image_size = image.metadata().unwrap().len();
+
+        if image_size > size {
+            return Err(Status::new_exception_str(
+                ExceptionCode::ILLEGAL_ARGUMENT,
+                Some("Can't shrink encrypted storage"),
+            ));
+        }
+        // Reset the file length. In most filesystems, this will not allocate any physical disk
+        // space, it will only change the logical size.
+        image.set_len(size).context("Failed to extend file").or_service_specific_exception(-1)?;
+        Ok(())
+    }
+
     /// Creates or update the idsig file by digesting the input APK file.
     fn createOrUpdateIdsigFile(
         &self,
diff --git a/android/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl b/android/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
index 169c3dc..37222ba 100644
--- a/android/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
+++ b/android/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
@@ -58,6 +58,13 @@
             in ParcelFileDescriptor imageFd, long sizeBytes, PartitionType type);
 
     /**
+     * Set the encrypted storage size.
+     *
+     * The file must be open with both read and write permissions.
+     */
+    void setEncryptedStorageSize(in ParcelFileDescriptor imageFd, long size);
+
+    /**
      * Create or update an idsig file that digests the given APK file. The idsig file follows the
      * idsig format that is defined by the APK Signature Scheme V4. The idsig file is not updated
      * when it is up to date with the input file, which is checked by comparing the
diff --git a/android/vm/src/run.rs b/android/vm/src/run.rs
index a362b8e..8385fb4 100644
--- a/android/vm/src/run.rs
+++ b/android/vm/src/run.rs
@@ -35,6 +35,7 @@
 use rand::{distributions::Alphanumeric, Rng};
 use std::fs;
 use std::fs::File;
+use std::fs::OpenOptions;
 use std::io;
 use std::io::{Read, Write};
 use std::os::fd::AsFd;
@@ -112,6 +113,8 @@
                 config.microdroid.storage_size.unwrap_or(10 * 1024 * 1024),
                 PartitionType::ENCRYPTEDSTORE,
             )?;
+        } else if let Some(storage_size) = config.microdroid.storage_size {
+            set_encrypted_storage(service.as_ref(), path, storage_size)?;
         }
         Some(open_parcel_file(path, true)?)
     } else {
@@ -370,6 +373,22 @@
     Ok(config.extra_apks.into_iter().map(|x| x.path.into()).collect())
 }
 
+fn set_encrypted_storage(
+    service: &dyn IVirtualizationService,
+    image_path: &Path,
+    size: u64,
+) -> Result<(), Error> {
+    let image = OpenOptions::new()
+        .create_new(false)
+        .read(true)
+        .write(true)
+        .open(image_path)
+        .with_context(|| format!("Failed to open {:?}", image_path))?;
+
+    service.setEncryptedStorageSize(&ParcelFileDescriptor::new(image), size.try_into()?)?;
+    Ok(())
+}
+
 struct Callback {}
 
 impl vmclient::VmCallback for Callback {