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 {
diff --git a/build/microdroid/Android.bp b/build/microdroid/Android.bp
index 059077a..10b492b 100644
--- a/build/microdroid/Android.bp
+++ b/build/microdroid/Android.bp
@@ -82,7 +82,9 @@
         "microdroid_file_contexts",
         "microdroid_manifest",
         "microdroid_property_contexts",
+        "e2fsck.microdroid",
         "mke2fs.microdroid",
+        "resize2fs.microdroid",
         "microdroid_fstab",
 
         "libvm_payload", // used by payload to interact with microdroid manager
diff --git a/guest/encryptedstore/src/main.rs b/guest/encryptedstore/src/main.rs
index dd4ee3b..8647003 100644
--- a/guest/encryptedstore/src/main.rs
+++ b/guest/encryptedstore/src/main.rs
@@ -30,7 +30,9 @@
 use std::path::{Path, PathBuf};
 use std::process::Command;
 
+const E2FSCK_BIN: &str = "/system/bin/e2fsck";
 const MK2FS_BIN: &str = "/system/bin/mke2fs";
+const RESIZE2FS_BIN: &str = "/system/bin/resize2fs";
 const UNFORMATTED_STORAGE_MAGIC: &str = "UNFORMATTED-STORAGE";
 
 fn main() {
@@ -91,6 +93,8 @@
     if needs_formatting {
         info!("Freshly formatting the crypt device");
         format_ext4(&crypt_device)?;
+    } else {
+        resize_fs(&crypt_device)?;
     }
     mount(&crypt_device, mountpoint)
         .with_context(|| format!("Unable to mount {:?}", crypt_device))?;
@@ -174,6 +178,27 @@
     Ok(())
 }
 
+fn resize_fs(device: &Path) -> Result<()> {
+    // Check the partition
+    Command::new(E2FSCK_BIN)
+        .arg("-fvy")
+        .arg(device)
+        .status()
+        .context("failed to execute e2fsck")?;
+
+    // Resize the filesystem to the size of the device.
+    Command::new(RESIZE2FS_BIN).arg(device).status().context("failed to execute resize2fs")?;
+
+    // Finally check again if we were successful.
+    Command::new(E2FSCK_BIN)
+        .arg("-fvy")
+        .arg(device)
+        .status()
+        .context("failed to execute e2fsck")?;
+
+    Ok(())
+}
+
 fn mount(source: &Path, mountpoint: &Path) -> Result<()> {
     create_dir_all(mountpoint).with_context(|| format!("Failed to create {:?}", &mountpoint))?;
     let mount_options = CString::new(
diff --git a/guest/microdroid_manager/src/main.rs b/guest/microdroid_manager/src/main.rs
index 57ad35d..d665c87 100644
--- a/guest/microdroid_manager/src/main.rs
+++ b/guest/microdroid_manager/src/main.rs
@@ -368,14 +368,6 @@
         umount2("/microdroid_resources", MntFlags::MNT_DETACH)?;
     }
 
-    // Run encryptedstore binary to prepare the storage
-    let encryptedstore_child = if Path::new(ENCRYPTEDSTORE_BACKING_DEVICE).exists() {
-        info!("Preparing encryptedstore ...");
-        Some(prepare_encryptedstore(&vm_secret).context("encryptedstore run")?)
-    } else {
-        None
-    };
-
     let mut zipfuse = Zipfuse::default();
 
     // Before reading a file from the APK, start zipfuse
@@ -410,6 +402,19 @@
     );
     mount_extra_apks(&config, &mut zipfuse)?;
 
+    // Wait until apex config is done. (e.g. linker configuration for apexes)
+    wait_for_property_true(APEX_CONFIG_DONE_PROP).context("Failed waiting for apex config done")?;
+
+    // Run encryptedstore binary to prepare the storage
+    // Postpone initialization until apex mount completes to ensure e2fsck and resize2fs binaries
+    // are accessible.
+    let encryptedstore_child = if Path::new(ENCRYPTEDSTORE_BACKING_DEVICE).exists() {
+        info!("Preparing encryptedstore ...");
+        Some(prepare_encryptedstore(&vm_secret).context("encryptedstore run")?)
+    } else {
+        None
+    };
+
     register_vm_payload_service(
         allow_restricted_apis,
         service.clone(),
@@ -425,9 +430,6 @@
             .context("set microdroid_manager.export_tombstones.enabled")?;
     }
 
-    // Wait until apex config is done. (e.g. linker configuration for apexes)
-    wait_for_property_true(APEX_CONFIG_DONE_PROP).context("Failed waiting for apex config done")?;
-
     // Trigger init post-fs-data. This will start authfs if we wask it to.
     if config.enable_authfs {
         system_properties::write("microdroid_manager.authfs.enabled", "1")
diff --git a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java
index 5f634ef..b1b7a90 100644
--- a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java
+++ b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java
@@ -1575,6 +1575,12 @@
                                 ? createVirtualMachineConfigForRawFrom(vmConfig)
                                 : createVirtualMachineConfigForAppFrom(vmConfig, service);
 
+                if (vmConfig.isEncryptedStorageEnabled()) {
+                    service.setEncryptedStorageSize(
+                        ParcelFileDescriptor.open(mEncryptedStoreFilePath, MODE_READ_WRITE),
+                        vmConfig.getEncryptedStorageBytes());
+                }
+
                 mVirtualMachine =
                         service.createVm(
                                 vmConfigParcel, consoleOutFd, consoleInFd, mLogWriter, null);
@@ -1930,6 +1936,8 @@
      * <p>The new config must be {@linkplain VirtualMachineConfig#isCompatibleWith compatible with}
      * the existing config.
      *
+     * <p>NOTE: Modification of the encrypted storage size is restricted to expansion only and is an
+     * irreversible operation.
      * <p>NOTE: This method may block and should not be called on the main thread.
      *
      * @return the old config
@@ -1944,7 +1952,9 @@
             throws VirtualMachineException {
         synchronized (mLock) {
             VirtualMachineConfig oldConfig = mConfig;
-            if (!oldConfig.isCompatibleWith(newConfig)) {
+            if (!oldConfig.isCompatibleWith(newConfig)
+                    || oldConfig.getEncryptedStorageBytes()
+                            > newConfig.getEncryptedStorageBytes()) {
                 throw new VirtualMachineException("incompatible config");
             }
             checkStopped();
diff --git a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachineConfig.java b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachineConfig.java
index 83b234d..a543c19 100644
--- a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -634,7 +634,6 @@
         }
         return this.mDebugLevel == other.mDebugLevel
                 && this.mProtectedVm == other.mProtectedVm
-                && this.mEncryptedStorageBytes == other.mEncryptedStorageBytes
                 && this.mVmOutputCaptured == other.mVmOutputCaptured
                 && this.mVmConsoleInputSupported == other.mVmConsoleInputSupported
                 && this.mConnectVmConsole == other.mConnectVmConsole
diff --git a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
index 4c824f0..5d92976 100644
--- a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
@@ -47,6 +47,9 @@
     /* get the encrypted storage path. */
     String getEncryptedStoragePath();
 
+    /* get the size of the encrypted storage in bytes. */
+    long getEncryptedStorageSize();
+
     /* start a simple vsock server on ECHO_REVERSE_PORT that reads a line at a time and echoes
      * each line reverse.
      */
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index 67249b4..6256cfb 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -619,6 +619,7 @@
         public String mExtraApkTestProp;
         public String mApkContentsPath;
         public String mEncryptedStoragePath;
+        public long mEncryptedStorageSize;
         public String[] mEffectiveCapabilities;
         public int mUid;
         public String mFileContent;
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 e4a3ff6..6fd0885 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -129,6 +129,7 @@
     private static final String TEST_APP_PACKAGE_NAME = "com.android.microdroid.test";
     private static final String VM_ATTESTATION_PAYLOAD_PATH = "libvm_attestation_test_payload.so";
     private static final String VM_ATTESTATION_MESSAGE = "Hello RKP from AVF!";
+    private static final long TOLERANCE_BYTES = 400_000;
     private static final int ENCRYPTED_STORAGE_BYTES = 4_000_000;
 
     private static final String RELAXED_ROLLBACK_PROTECTION_SCHEME_TEST_PACKAGE_NAME =
@@ -741,11 +742,6 @@
         // so in the API spec.
         assertConfigCompatible(baseline, newBaselineBuilder().setApkPath("/different")).isTrue();
 
-        // Changes that are currently incompatible for ease of implementation, but this might change
-        // in the future.
-        assertConfigCompatible(baseline, newBaselineBuilder().setEncryptedStorageBytes(100_000))
-                .isFalse();
-
         VirtualMachineConfig.Builder debuggableBuilder =
                 newBaselineBuilder().setDebugLevel(DEBUG_LEVEL_FULL);
         VirtualMachineConfig debuggable = debuggableBuilder.build();
@@ -1866,6 +1862,169 @@
         assertThat(testResults.mFileContent).isEqualTo(EXAMPLE_STRING);
     }
 
+    @Test
+    @CddTest
+    public void encryptedStorageSupportsExpansion() throws Exception {
+        assumeSupportedDevice();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                        .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES)
+                        .build();
+
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
+        TestResults testResults =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (ts, tr) -> {
+                            tr.mEncryptedStorageSize = ts.getEncryptedStorageSize();
+                        });
+        testResults.assertNoException();
+        assertThat(testResults.mEncryptedStorageSize)
+            .isWithin(TOLERANCE_BYTES)
+            .of(ENCRYPTED_STORAGE_BYTES);
+
+        // Re-run the VM with more storage size & verify the file persisted.
+        // Note, the previous `runVmTestService` stopped the VM
+        config = newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                    .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES * 2)
+                    .build();
+        vm.setConfig(config);
+        assertThat(vm.getConfig().getEncryptedStorageBytes())
+            .isEqualTo(ENCRYPTED_STORAGE_BYTES * 2);
+
+        testResults =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (ts, tr) -> {
+                            tr.mEncryptedStorageSize = ts.getEncryptedStorageSize();
+                        });
+        testResults.assertNoException();
+        assertThat(testResults.mEncryptedStorageSize)
+            .isWithin(TOLERANCE_BYTES)
+            .of(ENCRYPTED_STORAGE_BYTES * 2);
+    }
+
+    @Test
+    @CddTest
+    public void encryptedStorageExpansionIsPersistent() throws Exception {
+        assumeSupportedDevice();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                        .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES)
+                        .build();
+
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
+        TestResults testResults =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (ts, tr) -> {
+                            ts.writeToFile(
+                                    /* content= */ EXAMPLE_STRING,
+                                    /* path= */ "/mnt/encryptedstore/test_file");
+                        });
+        testResults.assertNoException();
+
+        // Re-run the VM with more storage size & verify the file persisted.
+        // Note, the previous `runVmTestService` stopped the VM
+        config = newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                    .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES * 2)
+                    .build();
+        vm.setConfig(config);
+
+        testResults =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (ts, tr) -> {
+                            tr.mFileContent = ts.readFromFile("/mnt/encryptedstore/test_file");
+                        });
+        testResults.assertNoException();
+        assertThat(testResults.mFileContent).isEqualTo(EXAMPLE_STRING);
+    }
+
+    @Test
+    @CddTest
+    public void encryptedStorageSizeUnchanged() throws Exception {
+        assumeSupportedDevice();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                        .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES)
+                        .build();
+
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
+        TestResults testResults =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (ts, tr) -> {
+                            tr.mEncryptedStorageSize = ts.getEncryptedStorageSize();
+                        });
+        testResults.assertNoException();
+        assertThat(testResults.mEncryptedStorageSize)
+            .isWithin(TOLERANCE_BYTES)
+            .of(ENCRYPTED_STORAGE_BYTES);
+
+        // Re-run the VM with more storage size & verify the file persisted.
+        // Note, the previous `runVmTestService` stopped the VM
+        config = newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                    .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES)
+                    .build();
+        vm.setConfig(config);
+        assertThat(vm.getConfig().getEncryptedStorageBytes())
+            .isEqualTo(ENCRYPTED_STORAGE_BYTES);
+
+        testResults =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (ts, tr) -> {
+                            tr.mEncryptedStorageSize = ts.getEncryptedStorageSize();
+                        });
+        testResults.assertNoException();
+        assertThat(testResults.mEncryptedStorageSize)
+            .isWithin(TOLERANCE_BYTES)
+            .of(ENCRYPTED_STORAGE_BYTES);
+    }
+
+    @Test
+    @CddTest
+    public void encryptedStorageShrinkFails() throws Exception {
+        assumeSupportedDevice();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                        .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES)
+                        .build();
+
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
+        TestResults testResults =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (ts, tr) -> {
+                            tr.mEncryptedStorageSize = ts.getEncryptedStorageSize();
+                        });
+        testResults.assertNoException();
+        assertThat(testResults.mEncryptedStorageSize)
+            .isWithin(TOLERANCE_BYTES)
+            .of(ENCRYPTED_STORAGE_BYTES);
+
+        // Re-run the VM with more storage size & verify the file persisted.
+        // Note, the previous `runVmTestService` stopped the VM
+        VirtualMachineConfig newConfig =
+            newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                    .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES / 2)
+                    .build();
+        assertThrowsVmExceptionContaining(
+            () -> vm.setConfig(newConfig), "incompatible config");
+    }
+
     private boolean deviceCapableOfProtectedVm() {
         int capabilities = getVirtualMachineManager().getCapabilities();
         if ((capabilities & CAPABILITY_PROTECTED_VM) != 0) {
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index 355cfb1..13eafce 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -29,6 +29,7 @@
 #include <stdint.h>
 #include <stdio.h>
 #include <sys/capability.h>
+#include <sys/statvfs.h>
 #include <sys/system_properties.h>
 #ifdef __MICRODROID_TEST_PAYLOAD_USES_LIBICU__
 #include <unicode/uchar.h>
@@ -232,6 +233,23 @@
             return ScopedAStatus::ok();
         }
 
+        ScopedAStatus getEncryptedStorageSize(int64_t *out) override {
+            const char* path_c = AVmPayload_getEncryptedStoragePath();
+            if (path_c == nullptr) {
+                *out = 0;
+                return ScopedAStatus::ok();
+            }
+            struct statvfs buffer;
+            if (statvfs(path_c, &buffer) != 0) {
+                std::string msg = "statvfs " + std::string(path_c) + " failed :  " +
+                    std::strerror(errno);
+                return ScopedAStatus::fromExceptionCodeWithMessage(EX_SERVICE_SPECIFIC,
+                                                                   msg.c_str());
+            }
+            *out= buffer.f_blocks * buffer.f_frsize;
+            return ScopedAStatus::ok();
+        }
+
         ScopedAStatus getEffectiveCapabilities(std::vector<std::string>* out) override {
             if (out == nullptr) {
                 return ScopedAStatus::ok();
diff --git a/tests/testapk/src/native/testbinary.rs b/tests/testapk/src/native/testbinary.rs
index 3900cad..6a7d96e 100644
--- a/tests/testapk/src/native/testbinary.rs
+++ b/tests/testapk/src/native/testbinary.rs
@@ -87,6 +87,10 @@
 
     // Everything below here is unimplemented. Implementations may be added as needed.
 
+    fn getEncryptedStorageSize(&self) -> BinderResult<i64> {
+        unimplemented()
+    }
+
     fn readProperty(&self, _: &str) -> BinderResult<String> {
         unimplemented()
     }
diff --git a/tests/vmshareapp/src/java/com/android/microdroid/test/sharevm/VmShareServiceImpl.java b/tests/vmshareapp/src/java/com/android/microdroid/test/sharevm/VmShareServiceImpl.java
index 7f44fa5..9489aed 100644
--- a/tests/vmshareapp/src/java/com/android/microdroid/test/sharevm/VmShareServiceImpl.java
+++ b/tests/vmshareapp/src/java/com/android/microdroid/test/sharevm/VmShareServiceImpl.java
@@ -223,6 +223,11 @@
         }
 
         @Override
+        public long getEncryptedStorageSize() throws RemoteException {
+            throw new UnsupportedOperationException("Not supported");
+        }
+
+        @Override
         public void runEchoReverseServer() throws RemoteException {
             throw new UnsupportedOperationException("Not supported");
         }