Enable transferring VM with encryptedstore enabled

VirtualMachineDescriptor should track encryptedstoreFd as well (this
will be null if encryptedstore in not enabled in config). Transfer of
this descriptor should transfer the storage image as well.

This patch also adds tests #importedVmIsEqualToTheOriginalVm_WithStorage
to test the functionality

Bug: 259383481
Test: atest #importedVmIsEqualToTheOriginalVm_WithStorage
Test: atest #importedVmIsEqualToTheOriginalVm_WithoutStorage
Change-Id: I7e06616910dea1f9b4c783bea5ce2755a9b79ada
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 1330afb..0893374 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -363,6 +363,16 @@
                 throw new VirtualMachineException("failed to create instance image", e);
             }
             vm.importInstanceFrom(vmDescriptor.getInstanceImgFd());
+
+            if (vmDescriptor.getEncryptedStoreFd() != null) {
+                try {
+                    vm.mEncryptedStoreFilePath.createNewFile();
+                } catch (IOException e) {
+                    throw new VirtualMachineException(
+                            "failed to create encrypted storage image", e);
+                }
+                vm.importEncryptedStoreFrom(vmDescriptor.getEncryptedStoreFd());
+            }
             return vm;
         } catch (VirtualMachineException | RuntimeException e) {
             // If anything goes wrong, delete any files created so far and the VM's directory
@@ -1044,7 +1054,10 @@
             try {
                 return new VirtualMachineDescriptor(
                         ParcelFileDescriptor.open(mConfigFilePath, MODE_READ_ONLY),
-                        ParcelFileDescriptor.open(mInstanceFilePath, MODE_READ_ONLY));
+                        ParcelFileDescriptor.open(mInstanceFilePath, MODE_READ_ONLY),
+                        mEncryptedStoreFilePath != null
+                                ? ParcelFileDescriptor.open(mEncryptedStoreFilePath, MODE_READ_ONLY)
+                                : null);
             } catch (IOException e) {
                 throw new VirtualMachineException(e);
             }
@@ -1210,4 +1223,14 @@
             throw new VirtualMachineException("failed to transfer instance image", e);
         }
     }
+
+    private void importEncryptedStoreFrom(@NonNull ParcelFileDescriptor encryptedStoreFd)
+            throws VirtualMachineException {
+        try (FileChannel storeOutput = new FileOutputStream(mEncryptedStoreFilePath).getChannel();
+                FileChannel storeInput = new AutoCloseInputStream(encryptedStoreFd).getChannel()) {
+            storeOutput.transferFrom(storeInput, /*position=*/ 0, storeInput.size());
+        } catch (IOException e) {
+            throw new VirtualMachineException("failed to transfer encryptedstore image", e);
+        }
+    }
 }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineDescriptor.java b/javalib/src/android/system/virtualmachine/VirtualMachineDescriptor.java
index edaf5b4..c9718aa 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineDescriptor.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineDescriptor.java
@@ -19,6 +19,7 @@
 import static java.util.Objects.requireNonNull;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
@@ -37,7 +38,9 @@
 public final class VirtualMachineDescriptor implements Parcelable {
     @NonNull private final ParcelFileDescriptor mConfigFd;
     @NonNull private final ParcelFileDescriptor mInstanceImgFd;
-    // TODO(b/243129654): Add trusted storage fd once it is available.
+    // File descriptor of the image backing the encrypted storage - Will be null if encrypted
+    // storage is not enabled. */
+    @Nullable private final ParcelFileDescriptor mEncryptedStoreFd;
 
     @Override
     public int describeContents() {
@@ -48,6 +51,7 @@
     public void writeToParcel(@NonNull Parcel out, int flags) {
         mConfigFd.writeToParcel(out, flags);
         mInstanceImgFd.writeToParcel(out, flags);
+        if (mEncryptedStoreFd != null) mEncryptedStoreFd.writeToParcel(out, flags);
     }
 
     @NonNull
@@ -78,14 +82,27 @@
         return mInstanceImgFd;
     }
 
+    /**
+     * @return File descriptor of image backing the encrypted storage.
+     *     <p>This method will return null if encrypted storage is not enabled.
+     */
+    @Nullable
+    ParcelFileDescriptor getEncryptedStoreFd() {
+        return mEncryptedStoreFd;
+    }
+
     VirtualMachineDescriptor(
-            @NonNull ParcelFileDescriptor configFd, @NonNull ParcelFileDescriptor instanceImgFd) {
+            @NonNull ParcelFileDescriptor configFd,
+            @NonNull ParcelFileDescriptor instanceImgFd,
+            @Nullable ParcelFileDescriptor encryptedStoreFd) {
         mConfigFd = configFd;
         mInstanceImgFd = instanceImgFd;
+        mEncryptedStoreFd = encryptedStoreFd;
     }
 
     private VirtualMachineDescriptor(Parcel in) {
         mConfigFd = requireNonNull(in.readFileDescriptor());
         mInstanceImgFd = requireNonNull(in.readFileDescriptor());
+        mEncryptedStoreFd = in.readFileDescriptor();
     }
 }
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 0bf65b3..8fe8fa2 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -841,13 +841,28 @@
     }
 
     @Test
-    public void importedVmIsEqualToTheOriginalVm() throws Exception {
+    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-1"})
+    public void importedVmIsEqualToTheOriginalVm_WithoutStorage() throws Exception {
+        TestResults testResults = importedVmIsEqualToTheOriginalVm(false);
+        assertThat(testResults.mEncryptedStoragePath).isEqualTo("");
+    }
+
+    @Test
+    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-1"})
+    public void importedVmIsEqualToTheOriginalVm_WithStorage() throws Exception {
+        TestResults testResults = importedVmIsEqualToTheOriginalVm(true);
+        assertThat(testResults.mEncryptedStoragePath).isEqualTo("/mnt/encryptedstore");
+    }
+
+    private TestResults importedVmIsEqualToTheOriginalVm(boolean encryptedStoreEnabled)
+            throws Exception {
         // Arrange
-        VirtualMachineConfig config =
+        VirtualMachineConfig.Builder builder =
                 newVmConfigBuilder()
                         .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                        .setDebugLevel(DEBUG_LEVEL_FULL)
-                        .build();
+                        .setDebugLevel(DEBUG_LEVEL_FULL);
+        if (encryptedStoreEnabled) builder = builder.setEncryptedStorageKib(4096);
+        VirtualMachineConfig config = builder.build();
         String vmNameOrig = "test_vm_orig";
         String vmNameImport = "test_vm_import";
         VirtualMachine vmOrig = forceCreateNewVirtualMachine(vmNameOrig, config);
@@ -867,12 +882,16 @@
         // Asserts
         assertFileContentsAreEqualInTwoVms("config.xml", vmNameOrig, vmNameImport);
         assertFileContentsAreEqualInTwoVms("instance.img", vmNameOrig, vmNameImport);
+        if (encryptedStoreEnabled) {
+            assertFileContentsAreEqualInTwoVms("storage.img", vmNameOrig, vmNameImport);
+        }
         assertThat(vmImport).isNotEqualTo(vmOrig);
         vmm.delete(vmNameOrig);
         assertThat(vmImport).isEqualTo(vmm.get(vmNameImport));
         TestResults testResults = runVmTestService(vmImport);
         assertThat(testResults.mException).isNull();
         assertThat(testResults.mAddInteger).isEqualTo(123 + 456);
+        return testResults;
     }
 
     @Test