Merge "Migrate ComposKeyTestCase to unit tests"
diff --git a/.prebuilt_info/prebuilt_info_pvmfw_pvmfw_img.asciipb b/.prebuilt_info/prebuilt_info_pvmfw_pvmfw_img.asciipb
new file mode 100644
index 0000000..4199b2b
--- /dev/null
+++ b/.prebuilt_info/prebuilt_info_pvmfw_pvmfw_img.asciipb
@@ -0,0 +1,12 @@
+drops {
+  android_build_drop {
+    build_id: "8134799"
+    target: "u-boot_pvmfw"
+    source_file: "pvmfw.img"
+  }
+  dest_file: "pvmfw/pvmfw.img"
+  version: ""
+  version_group: ""
+  git_project: "platform/packages/modules/Virtualization"
+  git_branch: "master"
+}
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index 9a0fe1a..8fe3403 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -294,6 +294,8 @@
         input_dir, 'etc', 'microdroid_bootconfig.app_debuggable')
     bootconfig_full_debuggable = os.path.join(
         input_dir, 'etc', 'microdroid_bootconfig.full_debuggable')
+    uboot_env_img = os.path.join(
+        input_dir, 'etc', 'uboot_env.img')
 
     # Key(pubkey) for bootloader should match with the one used to make VBmeta below
     # while it's okay to use different keys for other image files.
@@ -330,17 +332,21 @@
         MakeVbmetaImage(args, key, vbmeta_img, images=[
                         boot_img, vendor_boot_img, init_boot_img, system_a_img, vendor_a_img])
 
-    # Re-sign bootconfigs with the same key
+    # Re-sign bootconfigs and the uboot_env with the same key
     bootconfig_sign_key = key
     AddHashFooter(args, bootconfig_sign_key, bootconfig_normal)
     AddHashFooter(args, bootconfig_sign_key, bootconfig_app_debuggable)
     AddHashFooter(args, bootconfig_sign_key, bootconfig_full_debuggable)
+    AddHashFooter(args, bootconfig_sign_key, uboot_env_img)
 
-    # Re-sign vbmeta_bootconfig with a chained_partition to "bootconfig"
-    # Note that, for now, `key` and `bootconfig_sign_key` are the same, but technically they
-    # can be different. Vbmeta records pubkeys which signed chained partitions.
+    # Re-sign vbmeta_bootconfig with chained_partitions to "bootconfig" and
+    # "uboot_env". Note that, for now, `key` and `bootconfig_sign_key` are the
+    # same, but technically they can be different. Vbmeta records pubkeys which
+    # signed chained partitions.
     MakeVbmetaImage(args, key, vbmeta_bootconfig_img, chained_partitions={
-                    'bootconfig': bootconfig_sign_key})
+                    'bootconfig': bootconfig_sign_key,
+                    'uboot_env': bootconfig_sign_key,
+    })
 
 
 def VerifyVirtApex(args):
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index f4b3440..6a35fb0 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -122,6 +122,7 @@
             configPath: config_path.to_owned(),
             debugLevel: debug_level,
             extraIdsigs: vec![idsig_manifest_apk_fd],
+            protectedVm: false,
             memoryMib: VM_MEMORY_MIB,
             numCpus: parameters.cpus.map_or(1, NonZeroU32::get) as i32,
             cpuAffinity: parameters.cpu_set.clone(),
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index 7d1f9b0..04a90e0 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -51,6 +51,7 @@
     private static final String KEY_APKPATH = "apkPath";
     private static final String KEY_PAYLOADCONFIGPATH = "payloadConfigPath";
     private static final String KEY_DEBUGLEVEL = "debugLevel";
+    private static final String KEY_PROTECTED_VM = "protectedVm";
     private static final String KEY_MEMORY_MIB = "memoryMib";
     private static final String KEY_NUM_CPUS = "numCpus";
     private static final String KEY_CPU_AFFINITY = "cpuAffinity";
@@ -83,6 +84,11 @@
     private final DebugLevel mDebugLevel;
 
     /**
+     * Whether to run the VM in protected mode, so the host can't access its memory.
+     */
+    private final boolean mProtectedVm;
+
+    /**
      * The amount of RAM to give the VM, in MiB. If this is 0 or negative the default will be used.
      */
     private final int mMemoryMib;
@@ -111,6 +117,7 @@
             @NonNull Signature[] certs,
             @NonNull String payloadConfigPath,
             DebugLevel debugLevel,
+            boolean protectedVm,
             int memoryMib,
             int numCpus,
             String cpuAffinity) {
@@ -118,6 +125,7 @@
         mCerts = certs;
         mPayloadConfigPath = payloadConfigPath;
         mDebugLevel = debugLevel;
+        mProtectedVm = protectedVm;
         mMemoryMib = memoryMib;
         mNumCpus = numCpus;
         mCpuAffinity = cpuAffinity;
@@ -149,11 +157,12 @@
             throw new VirtualMachineException("No payloadConfigPath");
         }
         final DebugLevel debugLevel = DebugLevel.values()[b.getInt(KEY_DEBUGLEVEL)];
+        final boolean protectedVm = b.getBoolean(KEY_PROTECTED_VM);
         final int memoryMib = b.getInt(KEY_MEMORY_MIB);
         final int numCpus = b.getInt(KEY_NUM_CPUS);
         final String cpuAffinity = b.getString(KEY_CPU_AFFINITY);
-        return new VirtualMachineConfig(apkPath, certs, payloadConfigPath, debugLevel, memoryMib,
-                numCpus, cpuAffinity);
+        return new VirtualMachineConfig(apkPath, certs, payloadConfigPath, debugLevel, protectedVm,
+                memoryMib, numCpus, cpuAffinity);
     }
 
     /** Persists this config to a stream, for example a file. */
@@ -169,6 +178,8 @@
         b.putStringArray(KEY_CERTS, certs);
         b.putString(KEY_PAYLOADCONFIGPATH, mPayloadConfigPath);
         b.putInt(KEY_DEBUGLEVEL, mDebugLevel.ordinal());
+        b.putBoolean(KEY_PROTECTED_VM, mProtectedVm);
+        b.putInt(KEY_NUM_CPUS, mNumCpus);
         if (mMemoryMib > 0) {
             b.putInt(KEY_MEMORY_MIB, mMemoryMib);
         }
@@ -219,6 +230,7 @@
                 parcel.debugLevel = VirtualMachineAppConfig.DebugLevel.FULL;
                 break;
         }
+        parcel.protectedVm = mProtectedVm;
         parcel.memoryMib = mMemoryMib;
         parcel.numCpus = mNumCpus;
         parcel.cpuAffinity = mCpuAffinity;
@@ -230,16 +242,17 @@
         private Context mContext;
         private String mPayloadConfigPath;
         private DebugLevel mDebugLevel;
+        private boolean mProtectedVm;
         private int mMemoryMib;
         private int mNumCpus;
         private String mCpuAffinity;
-        // TODO(jiyong): add more items like # of cpu, size of ram, debuggability, etc.
 
         /** Creates a builder for the given context (APK), and the payload config file in APK. */
         public Builder(@NonNull Context context, @NonNull String payloadConfigPath) {
             mContext = context;
             mPayloadConfigPath = payloadConfigPath;
             mDebugLevel = DebugLevel.NONE;
+            mProtectedVm = false;
             mNumCpus = 1;
             mCpuAffinity = null;
         }
@@ -250,6 +263,12 @@
             return this;
         }
 
+        /** Sets whether to protect the VM memory from the host. Defaults to false. */
+        public Builder protectedVm(boolean protectedVm) {
+            mProtectedVm = protectedVm;
+            return this;
+        }
+
         /**
          * Sets the amount of RAM to give the VM. If this is zero or negative then the default will
          * be used.
@@ -309,8 +328,8 @@
             }
 
             return new VirtualMachineConfig(
-                    apkPath, certs, mPayloadConfigPath, mDebugLevel, mMemoryMib, mNumCpus,
-                    mCpuAffinity);
+                    apkPath, certs, mPayloadConfigPath, mDebugLevel, mProtectedVm, mMemoryMib,
+                    mNumCpus, mCpuAffinity);
         }
     }
 }
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 6f27ce1..3566bd2 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -348,6 +348,10 @@
             name: "bootconfig",
             private_key: ":microdroid_sign_key",
         },
+        {
+            name: "uboot_env",
+            private_key: ":microdroid_sign_key",
+        },
     ],
 }
 
@@ -524,18 +528,42 @@
 
 genrule {
     name: "microdroid_uboot_env_gen",
-    tools: ["mkenvimage_host"],
-    srcs: ["uboot-env.txt"],
+    tools: [
+        "mkenvimage_host",
+        "avbtool",
+    ],
+    srcs: [
+        "uboot-env.txt",
+        ":microdroid_sign_key",
+    ],
     out: ["output.img"],
-    cmd: "$(location mkenvimage_host) -s 4096 -o $(out) $(in)",
+    cmd: "$(location mkenvimage_host) -s 4096 -o $(out) $(location uboot-env.txt) && " +
+        "$(location avbtool) add_hash_footer " +
+        "--algorithm SHA256_RSA4096 " +
+        "--partition_name uboot_env " +
+        "--key $(location :microdroid_sign_key) " +
+        "--partition_size $$(( " + avb_hash_footer_kb + " * 1024 + ( $$(stat --format=%s $(out)) + 4096 - 1 ) / 4096 * 4096 )) " +
+        "--image $(out)",
 }
 
 genrule {
     name: "microdroid_uboot_env_gen_x86_64",
-    tools: ["mkenvimage_host"],
-    srcs: ["uboot-env-x86_64.txt"],
+    tools: [
+        "mkenvimage_host",
+        "avbtool",
+    ],
+    srcs: [
+        "uboot-env-x86_64.txt",
+        ":microdroid_sign_key",
+    ],
     out: ["output.img"],
-    cmd: "$(location mkenvimage_host) -s 4096 -o $(out) $(in)",
+    cmd: "$(location mkenvimage_host) -s 4096 -o $(out) $(location uboot-env-x86_64.txt) && " +
+        "$(location avbtool) add_hash_footer " +
+        "--algorithm SHA256_RSA4096 " +
+        "--partition_name uboot_env " +
+        "--key $(location :microdroid_sign_key) " +
+        "--partition_size $$(( " + avb_hash_footer_kb + " * 1024 + ( $$(stat --format=%s $(out)) + 4096 - 1 ) / 4096 * 4096 )) " +
+        "--image $(out)",
 }
 
 // Note that keys can be different for filesystem images even though we're using the same key
diff --git a/microdroid/microdroid.json b/microdroid/microdroid.json
index 0c294e9..cb27a24 100644
--- a/microdroid/microdroid.json
+++ b/microdroid/microdroid.json
@@ -37,6 +37,5 @@
       "writable": true
     }
   ],
-  "memory_mib": 256,
-  "protected": false
+  "memory_mib": 256
 }
diff --git a/microdroid_manager/src/instance.rs b/microdroid_manager/src/instance.rs
index 5a77198..f3bbf16 100644
--- a/microdroid_manager/src/instance.rs
+++ b/microdroid_manager/src/instance.rs
@@ -316,6 +316,7 @@
 
 #[derive(Debug, Serialize, Deserialize, PartialEq)]
 pub struct MicrodroidData {
+    pub salt: Vec<u8>, // Should be [u8; 64] but that isn't serializable.
     pub apk_data: ApkData,
     pub extra_apks_data: Vec<ApkData>,
     pub apex_data: Vec<ApexData>,
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 064933d..827f9ff 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -36,6 +36,7 @@
 use microdroid_payload_config::{Task, TaskType, VmPayloadConfig};
 use once_cell::sync::OnceCell;
 use payload::{get_apex_data_from_payload, load_metadata, to_metadata};
+use rand::Fill;
 use ring::digest;
 use rustutils::system_properties;
 use rustutils::system_properties::PropertyWatcher;
@@ -195,7 +196,7 @@
             authorityHash: authority_hash,
             authorityDescriptor: None,
             mode: if is_debuggable()? { Mode::DEBUG } else { Mode::NORMAL },
-            hidden: [0; 64],
+            hidden: verified_data.salt.try_into().unwrap(),
         }])
         .context("IDiceMaintenance::demoteSelf failed")?;
     Ok(())
@@ -438,9 +439,19 @@
 
     info!("payload verification successful. took {:#?}", start_time.elapsed().unwrap());
 
+    // Use the salt from a verified instance, or generate a salt for a new instance.
+    let salt = if let Some(saved_data) = saved_data {
+        saved_data.salt.clone()
+    } else {
+        let mut salt = vec![0u8; 64];
+        salt.as_mut_slice().try_fill(&mut rand::thread_rng())?;
+        salt
+    };
+
     // At this point, we can ensure that the root_hash from the idsig file is trusted, either by
     // fully verifying the APK or by comparing it with the saved root_hash.
     Ok(MicrodroidData {
+        salt,
         apk_data: ApkData { root_hash: root_hash_from_idsig, pubkey: main_apk_pubkey },
         extra_apks_data,
         apex_data: apex_data_from_payload,
diff --git a/pvmfw/pvmfw.img b/pvmfw/pvmfw.img
index 510b2c4..a2541a3 100644
--- a/pvmfw/pvmfw.img
+++ b/pvmfw/pvmfw.img
Binary files differ
diff --git a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
index 208d61f..f15036c 100644
--- a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
@@ -24,4 +24,7 @@
 
     /* read a system property. */
     String readProperty(String prop);
+
+    /* get the VM's stable secret. */
+    byte[] insecurelyExposeSecret();
 }
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index 4cca538..40d72fe 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -10,6 +10,7 @@
         "androidx.test.runner",
         "androidx.test.ext.junit",
         "com.android.microdroid.testservice-java",
+        "truth-prebuilt",
     ],
     libs: ["android.system.virtualmachine"],
     jni_libs: ["MicrodroidTestNativeLib"],
@@ -22,6 +23,7 @@
     name: "MicrodroidTestNativeLib",
     srcs: ["src/native/testbinary.cpp"],
     shared_libs: [
+        "android.security.dice-ndk",
         "android.system.virtualmachineservice-ndk",
         "com.android.microdroid.testservice-ndk",
         "libbase",
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 bd44a3c..803bdc6 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -15,14 +15,14 @@
  */
 package com.android.microdroid.test;
 
-import static org.hamcrest.core.Is.is;
-import static org.hamcrest.core.IsNot.not;
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.TruthJUnit.assume;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeNoException;
-import static org.junit.Assume.assumeThat;
 
 import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
 
@@ -52,6 +52,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.nio.file.Files;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -212,8 +213,10 @@
     @Test
     public void changingDebugLevelInvalidatesVmIdentity()
             throws VirtualMachineException, InterruptedException, IOException {
-        assumeThat("Skip on Cuttlefish. b/195765441",
-                android.os.Build.DEVICE, is(not("vsoc_x86_64")));
+        assume()
+            .withMessage("Skip on Cuttlefish. b/195765441")
+            .that(android.os.Build.DEVICE)
+            .isNotEqualTo("vsoc_x86_64");
 
         VirtualMachineConfig.Builder builder =
                 new VirtualMachineConfig.Builder(mInner.mContext, "assets/vm_config.json");
@@ -269,4 +272,64 @@
                 };
         listener.runToFinish(mInner.mVm);
     }
+
+    private byte[] launchVmAndGetSecret(String instanceName)
+            throws VirtualMachineException, InterruptedException {
+        VirtualMachineConfig.Builder builder =
+                new VirtualMachineConfig.Builder(mInner.mContext, "assets/vm_config.json");
+        VirtualMachineConfig normalConfig = builder.debugLevel(DebugLevel.NONE).build();
+        mInner.mVm = mInner.mVmm.getOrCreate(instanceName, normalConfig);
+        final CompletableFuture<byte[]> secret = new CompletableFuture<>();
+        VmEventListener listener =
+                new VmEventListener() {
+                    @Override
+                    public void onPayloadReady(VirtualMachine vm) {
+                        try {
+                            ITestService testService = ITestService.Stub.asInterface(
+                                    vm.connectToVsockServer(ITestService.SERVICE_PORT).get());
+                            secret.complete(testService.insecurelyExposeSecret());
+                        } catch (Exception e) {
+                            fail("Exception while connecting to service: " + e.toString());
+                        }
+                        // TODO(b/208639280): remove this sleep. For now, we need to wait for a few
+                        // seconds so that crosvm can actually persist instance.img.
+                        try {
+                            Thread.sleep(30 * 1000);
+                        } catch (InterruptedException e) { }
+                        forceStop(vm);
+                    }
+                };
+        listener.runToFinish(mInner.mVm);
+        return secret.getNow(null);
+    }
+
+    @Test
+    public void instancesOfSameVmHaveDifferentSecrets()
+            throws VirtualMachineException, InterruptedException {
+        assume()
+            .withMessage("Skip on Cuttlefish. b/195765441")
+            .that(android.os.Build.DEVICE)
+            .isNotEqualTo("vsoc_x86_64");
+
+        byte[] vm_a_secret = launchVmAndGetSecret("test_vm_a");
+        byte[] vm_b_secret = launchVmAndGetSecret("test_vm_b");
+        assertThat(vm_a_secret).isNotNull();
+        assertThat(vm_b_secret).isNotNull();
+        assertThat(vm_a_secret).isNotEqualTo(vm_b_secret);
+    }
+
+    @Test
+    public void sameInstanceKeepsSameSecrets()
+            throws VirtualMachineException, InterruptedException {
+        assume()
+            .withMessage("Skip on Cuttlefish. b/195765441")
+            .that(android.os.Build.DEVICE)
+            .isNotEqualTo("vsoc_x86_64");
+
+        byte[] vm_secret_first_boot = launchVmAndGetSecret("test_vm");
+        byte[] vm_secret_second_boot = launchVmAndGetSecret("test_vm");
+        assertThat(vm_secret_first_boot).isNotNull();
+        assertThat(vm_secret_second_boot).isNotNull();
+        assertThat(vm_secret_first_boot).isEqualTo(vm_secret_second_boot);
+    }
 }
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index 301328a..417ff4a 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include <aidl/android/security/dice/IDiceNode.h>
 #include <aidl/android/system/virtualmachineservice/IVirtualMachineService.h>
 #include <aidl/com/android/microdroid/testservice/BnTestService.h>
 #include <android-base/file.h>
@@ -32,6 +33,9 @@
 #include <binder_rpc_unstable.hpp>
 #include <string>
 
+using aidl::android::hardware::security::dice::BccHandover;
+using aidl::android::security::dice::IDiceNode;
+
 using aidl::android::system::virtualmachineservice::IVirtualMachineService;
 
 using android::base::ErrnoError;
@@ -74,6 +78,23 @@
 
             return ndk::ScopedAStatus::ok();
         }
+
+        ndk::ScopedAStatus insecurelyExposeSecret(std::vector<uint8_t>* out) override {
+            ndk::SpAIBinder binder(AServiceManager_getService("android.security.dice.IDiceNode"));
+            auto service = IDiceNode::fromBinder(binder);
+            if (service == nullptr) {
+                return ndk::ScopedAStatus::
+                        fromServiceSpecificErrorWithMessage(0, "Failed to find diced");
+            }
+            BccHandover handover;
+            auto deriveStatus = service->derive({}, &handover);
+            if (!deriveStatus.isOk()) {
+                return ndk::ScopedAStatus::fromServiceSpecificErrorWithMessage(0,
+                                                                               "Failed call diced");
+            }
+            *out = {handover.cdiSeal.begin(), handover.cdiSeal.end()};
+            return ndk::ScopedAStatus::ok();
+        }
     };
     auto testService = ndk::SharedRefBase::make<TestService>();
 
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 653524e..b82064b 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -31,7 +31,6 @@
         "libcommand_fds",
         "libdisk",
         "libidsig",
-        "libkvm",
         "liblog_rust",
         "libmicrodroid_metadata",
         "libmicrodroid_payload_config",
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
index 8265f96..c36e561 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
@@ -47,6 +47,9 @@
     /** Debug level of the VM */
     DebugLevel debugLevel;
 
+    /** Whether the VM should be a protected VM. */
+    boolean protectedVm;
+
     /**
      * The amount of RAM to give the VM, in MiB. If this is 0 or negative then it will default to
      * the value in microdroid.json, if any, or the crosvm default.
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 42eb1e6..5b0c9b7 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -50,7 +50,6 @@
 use binder_common::{lazy_service::LazyServiceGuard, new_binder_exception};
 use disk::QcowFile;
 use idsig::{HashAlgorithm, V4Signature};
-use kvm::{Kvm, Cap};
 use log::{debug, error, info, warn};
 use microdroid_payload_config::VmPayloadConfig;
 use rustutils::system_properties;
@@ -190,9 +189,7 @@
             VirtualMachineConfig::AppConfig(config) => BorrowedOrOwned::Owned(
                 load_app_config(config, &temporary_directory).map_err(|e| {
                     error!("Failed to load app config from {}: {}", &config.configPath, e);
-                    // At this point, we do not know the protected status of Vm
-                    // setting it to false, though this may not be correct.
-                    write_vm_creation_stats(false, false);
+                    write_vm_creation_stats(config.protectedVm, false);
                     new_binder_exception(
                         ExceptionCode::SERVICE_SPECIFIC,
                         format!("Failed to load app config from {}: {}", &config.configPath, e),
@@ -202,7 +199,15 @@
             VirtualMachineConfig::RawConfig(config) => BorrowedOrOwned::Borrowed(config),
         };
         let config = config.as_ref();
-        let protected_vm = config.protectedVm;
+        let protected = config.protectedVm;
+
+        // Debug level FULL is only supported for non-protected VMs.
+        if is_debug_level_full && protected {
+            return Err(new_binder_exception(
+                ExceptionCode::SERVICE_SPECIFIC,
+                "FULL debug level not supported for protected VMs.",
+            ));
+        };
 
         // Check if partition images are labeled incorrectly. This is to prevent random images
         // which are not protected by the Android Verified Boot (e.g. bits downloaded by apps) from
@@ -226,7 +231,7 @@
         let zero_filler_path = temporary_directory.join("zero.img");
         write_zero_filler(&zero_filler_path).map_err(|e| {
             error!("Failed to make composite image: {}", e);
-            write_vm_creation_stats(protected_vm, false);
+            write_vm_creation_stats(protected, false);
             new_binder_exception(
                 ExceptionCode::SERVICE_SPECIFIC,
                 format!("Failed to make composite image: {}", e),
@@ -248,24 +253,6 @@
             })
             .collect::<Result<Vec<DiskFile>, _>>()?;
 
-        let protected_vm_supported = Kvm::new()
-            .map_err(|e| new_binder_exception(ExceptionCode::SERVICE_SPECIFIC, e.to_string()))?
-            .check_extension(Cap::ArmProtectedVm);
-        let protected = config.protectedVm && protected_vm_supported;
-        if config.protectedVm && !protected_vm_supported {
-            warn!("Protected VM was requested, but it isn't supported on this machine. Ignored.");
-        }
-
-        // And force run in non-protected mode when debug level is FULL
-        let protected = if is_debug_level_full {
-            if protected {
-                warn!("VM will run in FULL debug level. Running in non-protected mode");
-            }
-            false
-        } else {
-            protected
-        };
-
         // Actually start the VM.
         let crosvm_config = CrosvmConfig {
             cid,
@@ -292,7 +279,7 @@
             )
             .map_err(|e| {
                 error!("Failed to create VM with config {:?}: {}", config, e);
-                write_vm_creation_stats(protected_vm, false);
+                write_vm_creation_stats(protected, false);
                 new_binder_exception(
                     ExceptionCode::SERVICE_SPECIFIC,
                     format!("Failed to create VM: {}", e),
@@ -300,7 +287,7 @@
             })?,
         );
         state.add_vm(Arc::downgrade(&instance));
-        write_vm_creation_stats(protected_vm, true);
+        write_vm_creation_stats(protected, true);
         Ok(VirtualMachine::create(instance))
     }
 
@@ -587,6 +574,7 @@
         vm_config.memoryMib = config.memoryMib;
     }
 
+    vm_config.protectedVm = config.protectedVm;
     vm_config.numCpus = config.numCpus;
     vm_config.cpuAffinity = config.cpuAffinity.clone();
 
diff --git a/virtualizationservice/src/payload.rs b/virtualizationservice/src/payload.rs
index c661d44..7b8cb7f 100644
--- a/virtualizationservice/src/payload.rs
+++ b/virtualizationservice/src/payload.rs
@@ -53,7 +53,7 @@
     list: Vec<ApexInfo>,
 }
 
-#[derive(Clone, Debug, Deserialize)]
+#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq)]
 struct ApexInfo {
     #[serde(rename = "moduleName")]
     name: String,
@@ -69,6 +69,12 @@
 
     #[serde(rename = "isFactory")]
     is_factory: bool,
+
+    #[serde(rename = "isActive")]
+    is_active: bool,
+
+    #[serde(rename = "provideSharedApexLibs")]
+    provide_shared_apex_libs: bool,
 }
 
 impl ApexInfoList {
@@ -94,21 +100,19 @@
             Ok(apex_info_list)
         })
     }
+}
 
-    /// Returns the list of apex names matching with the predicate
-    fn get_matching(&self, predicate: fn(&ApexInfo) -> bool) -> Vec<String> {
-        self.list.iter().filter(|info| predicate(info)).map(|info| info.name.clone()).collect()
-    }
-
-    fn get(&self, apex_name: &str) -> Result<&ApexInfo> {
-        self.list
-            .iter()
-            .find(|apex| apex.name == apex_name)
-            .ok_or_else(|| anyhow!("{} not found.", apex_name))
-    }
-
-    fn get_path_for(&self, apex_name: &str) -> Result<PathBuf> {
-        Ok(self.get(apex_name)?.path.clone())
+impl ApexInfo {
+    fn matches(&self, apex_config: &ApexConfig) -> bool {
+        // Match with pseudo name "{CLASSPATH}" which represents APEXes contributing
+        // to any derive_classpath environment variable
+        if apex_config.name == "{CLASSPATH}" && self.has_classpath_jar {
+            return true;
+        }
+        if apex_config.name == self.name {
+            return true;
+        }
+        false
     }
 }
 
@@ -151,20 +155,18 @@
 
 fn make_metadata_file(
     config_path: &str,
-    apex_names: &[String],
+    apex_infos: &[&ApexInfo],
     temporary_directory: &Path,
-    apex_list: &ApexInfoList,
 ) -> Result<ParcelFileDescriptor> {
     let metadata_path = temporary_directory.join("metadata");
     let metadata = Metadata {
         version: 1,
-        apexes: apex_names
+        apexes: apex_infos
             .iter()
             .enumerate()
-            .map(|(i, apex_name)| {
-                let apex_info = apex_list.get(apex_name)?;
+            .map(|(i, apex_info)| {
                 Ok(ApexPayload {
-                    name: apex_name.clone(),
+                    name: apex_info.name.clone(),
                     partition_name: format!("microdroid-apex-{}", i),
                     last_update_seconds: apex_info.last_update_seconds,
                     is_factory: apex_info.is_factory,
@@ -226,12 +228,13 @@
     let pm = PackageManager::new()?;
     let apex_list = pm.get_apex_list(vm_payload_config.prefer_staged)?;
 
-    // collect APEX names from config
-    let apexes = collect_apex_names(&apex_list, &vm_payload_config.apexes, app_config.debugLevel);
-    info!("Microdroid payload APEXes: {:?}", apexes);
+    // collect APEXes from config
+    let apex_infos =
+        collect_apex_infos(&apex_list, &vm_payload_config.apexes, app_config.debugLevel);
+    info!("Microdroid payload APEXes: {:?}", apex_infos.iter().map(|ai| &ai.name));
 
     let metadata_file =
-        make_metadata_file(&app_config.configPath, &apexes, temporary_directory, &apex_list)?;
+        make_metadata_file(&app_config.configPath, &apex_infos, temporary_directory)?;
     // put metadata at the first partition
     let mut partitions = vec![Partition {
         label: "payload-metadata".to_owned(),
@@ -239,9 +242,8 @@
         writable: false,
     }];
 
-    for (i, apex) in apexes.iter().enumerate() {
-        let apex_path = apex_list.get_path_for(apex)?;
-        let apex_file = open_parcel_file(&apex_path, false)?;
+    for (i, apex_info) in apex_infos.iter().enumerate() {
+        let apex_file = open_parcel_file(&apex_info.path, false)?;
         partitions.push(Partition {
             label: format!("microdroid-apex-{}", i),
             image: Some(apex_file),
@@ -316,30 +318,26 @@
     Ok(apexes)
 }
 
-// Collect APEX names from config
-fn collect_apex_names(
-    apex_list: &ApexInfoList,
-    apexes: &[ApexConfig],
+// Collect ApexInfos from VM config
+fn collect_apex_infos<'a>(
+    apex_list: &'a ApexInfoList,
+    apex_configs: &[ApexConfig],
     debug_level: DebugLevel,
-) -> Vec<String> {
-    // Process pseudo names like "{CLASSPATH}".
-    // For now we have following pseudo APEX names:
-    // - {CLASSPATH}: represents APEXes contributing to any derive_classpath environment variable
-    let mut apex_names: Vec<String> = apexes
-        .iter()
-        .flat_map(|apex| match apex.name.as_str() {
-            "{CLASSPATH}" => apex_list.get_matching(|apex| apex.has_classpath_jar),
-            _ => vec![apex.name.clone()],
-        })
-        .collect();
-    // Add required APEXes
-    apex_names.extend(MICRODROID_REQUIRED_APEXES.iter().map(|name| name.to_string()));
+) -> Vec<&'a ApexInfo> {
+    let mut additional_apexes: Vec<&str> = MICRODROID_REQUIRED_APEXES.to_vec();
     if debug_level != DebugLevel::NONE {
-        apex_names.extend(MICRODROID_REQUIRED_APEXES_DEBUG.iter().map(|name| name.to_string()));
+        additional_apexes.extend(MICRODROID_REQUIRED_APEXES_DEBUG.to_vec());
     }
-    apex_names.sort();
-    apex_names.dedup();
-    apex_names
+
+    apex_list
+        .list
+        .iter()
+        .filter(|ai| {
+            apex_configs.iter().any(|cfg| ai.matches(cfg) && ai.is_active)
+                || additional_apexes.iter().any(|name| name == &ai.name && ai.is_active)
+                || ai.provide_shared_apex_libs
+        })
+        .collect()
 }
 
 pub fn add_microdroid_images(
@@ -409,36 +407,125 @@
     }
 
     #[test]
-    fn test_collect_apex_names() {
-        let apex_list = ApexInfoList {
+    fn test_collect_apexes() {
+        let apex_info_list = ApexInfoList {
             list: vec![
                 ApexInfo {
-                    name: "hasnt_classpath".to_string(),
-                    path: PathBuf::from("path0"),
+                    // 0
+                    name: "com.android.adbd".to_string(),
+                    path: PathBuf::from("adbd"),
                     has_classpath_jar: false,
                     last_update_seconds: 12345678,
                     is_factory: true,
+                    is_active: true,
+                    ..Default::default()
                 },
                 ApexInfo {
+                    // 1
+                    name: "com.android.os.statsd".to_string(),
+                    path: PathBuf::from("statsd"),
+                    has_classpath_jar: false,
+                    last_update_seconds: 12345678,
+                    is_factory: true,
+                    is_active: false,
+                    ..Default::default()
+                },
+                ApexInfo {
+                    // 2
+                    name: "com.android.os.statsd".to_string(),
+                    path: PathBuf::from("statsd/updated"),
+                    has_classpath_jar: false,
+                    last_update_seconds: 12345678 + 1,
+                    is_factory: false,
+                    is_active: true,
+                    ..Default::default()
+                },
+                ApexInfo {
+                    // 3
+                    name: "no_classpath".to_string(),
+                    path: PathBuf::from("no_classpath"),
+                    has_classpath_jar: false,
+                    last_update_seconds: 12345678,
+                    is_factory: true,
+                    is_active: true,
+                    ..Default::default()
+                },
+                ApexInfo {
+                    // 4
                     name: "has_classpath".to_string(),
-                    path: PathBuf::from("path1"),
+                    path: PathBuf::from("has_classpath"),
                     has_classpath_jar: true,
                     last_update_seconds: 87654321,
+                    is_factory: true,
+                    is_active: false,
+                    ..Default::default()
+                },
+                ApexInfo {
+                    // 5
+                    name: "has_classpath".to_string(),
+                    path: PathBuf::from("has_classpath/updated"),
+                    has_classpath_jar: true,
+                    last_update_seconds: 87654321 + 1,
                     is_factory: false,
+                    is_active: true,
+                    ..Default::default()
+                },
+                ApexInfo {
+                    // 6
+                    name: "apex-foo".to_string(),
+                    path: PathBuf::from("apex-foo"),
+                    has_classpath_jar: false,
+                    last_update_seconds: 87654321,
+                    is_factory: true,
+                    is_active: false,
+                    ..Default::default()
+                },
+                ApexInfo {
+                    // 7
+                    name: "apex-foo".to_string(),
+                    path: PathBuf::from("apex-foo/updated"),
+                    has_classpath_jar: false,
+                    last_update_seconds: 87654321 + 1,
+                    is_factory: false,
+                    is_active: true,
+                    ..Default::default()
+                },
+                ApexInfo {
+                    // 8
+                    name: "sharedlibs".to_string(),
+                    path: PathBuf::from("apex-foo"),
+                    last_update_seconds: 87654321,
+                    is_factory: true,
+                    provide_shared_apex_libs: true,
+                    ..Default::default()
+                },
+                ApexInfo {
+                    // 9
+                    name: "sharedlibs".to_string(),
+                    path: PathBuf::from("apex-foo/updated"),
+                    last_update_seconds: 87654321 + 1,
+                    is_active: true,
+                    provide_shared_apex_libs: true,
+                    ..Default::default()
                 },
             ],
         };
-        let apexes = vec![
-            ApexConfig { name: "config_name".to_string() },
+        let apex_configs = vec![
+            ApexConfig { name: "apex-foo".to_string() },
             ApexConfig { name: "{CLASSPATH}".to_string() },
         ];
         assert_eq!(
-            collect_apex_names(&apex_list, &apexes, DebugLevel::FULL),
+            collect_apex_infos(&apex_info_list, &apex_configs, DebugLevel::FULL),
             vec![
-                "com.android.adbd".to_string(),
-                "com.android.os.statsd".to_string(),
-                "config_name".to_string(),
-                "has_classpath".to_string(),
+                // Pass active/required APEXes
+                &apex_info_list.list[0],
+                &apex_info_list.list[2],
+                // Pass active APEXes specified in the config
+                &apex_info_list.list[5],
+                &apex_info_list.list[7],
+                // Pass both preinstalled(inactive) and updated(active) for "sharedlibs" APEXes
+                &apex_info_list.list[8],
+                &apex_info_list.list[9],
             ]
         );
     }
diff --git a/vm/src/main.rs b/vm/src/main.rs
index ad8c201..25f9bfb 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -72,6 +72,10 @@
         #[structopt(long, default_value = "none", parse(try_from_str=parse_debug_level))]
         debug: DebugLevel,
 
+        /// Run VM in protected mode.
+        #[structopt(short, long)]
+        protected: bool,
+
         /// Memory size (in MiB) of the VM. If unspecified, defaults to the value of `memory_mib`
         /// in the VM config file.
         #[structopt(short, long)]
@@ -174,6 +178,7 @@
             console,
             log,
             debug,
+            protected,
             mem,
             cpus,
             cpu_affinity,
@@ -188,6 +193,7 @@
             console.as_deref(),
             log.as_deref(),
             debug,
+            protected,
             mem,
             cpus,
             cpu_affinity,
diff --git a/vm/src/run.rs b/vm/src/run.rs
index 8583fe2..d558add 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -50,6 +50,7 @@
     console_path: Option<&Path>,
     log_path: Option<&Path>,
     debug_level: DebugLevel,
+    protected: bool,
     mem: Option<u32>,
     cpus: Option<u32>,
     cpu_affinity: Option<String>,
@@ -100,6 +101,7 @@
         instanceImage: open_parcel_file(instance, true /* writable */)?.into(),
         configPath: config_path.to_owned(),
         debugLevel: debug_level,
+        protectedVm: protected,
         memoryMib: mem.unwrap_or(0) as i32, // 0 means use the VM default
         numCpus: cpus.unwrap_or(1) as i32,
         cpuAffinity: cpu_affinity,