Merge "[bssl] Replace ok_or with ok_or_else in bssl for lazy evaluation" into main
diff --git a/README.md b/README.md
index 827e55c..4905b56 100644
--- a/README.md
+++ b/README.md
@@ -20,6 +20,7 @@
 * [Microdroid kernel](microdroid/kernel/README.md)
 * [Microdroid payload](microdroid/payload/README.md)
 * [vmbase](vmbase/README.md)
+* [Encrypted Storage](encryptedstore/README.md)
 
 AVF APIs:
 * [Java API](java/framework/README.md)
diff --git a/encryptedstore/README.md b/encryptedstore/README.md
new file mode 100644
index 0000000..544d6eb
--- /dev/null
+++ b/encryptedstore/README.md
@@ -0,0 +1,31 @@
+# Encrypted Storage
+
+Since Android U, AVF (with Microdroid) supports Encrypted Storage which is the storage solution
+in a VM. Within a VM, this is mounted at a path that can be retrieved via the [`AVmPayload_getEncryptedStoragePath()`][vm_payload_api].
+Any data written in encrypted storage is persisted and is available next time the VM is run.
+
+Encrypted Storage is backed by a para-virtualized block device on the guest which is further
+backed by a qcow2 disk image in the host. The block device is formatted with an ext4 filesystem.
+
+## Security
+
+Encrypted Storage uses block level encryption layer (Device-Mapper's "crypt" target) using a key
+derived from the VM secret and AES256 cipher with HCTR2 mode. The Block level encryption ensures
+the filesystem is also encrypted.
+
+### Integrity
+Encrypted Storage does not offer the level of integrity offered by primitives such as
+authenticated encryption/dm-integrity/RPMB. That level of integrity comes with substantial
+disk/performance overhead. Instead, it uses HCTR2 which is a super-pseudorandom
+permutation encryption mode, this offers better resilience against malleability attacks (than other
+modes such as XTS).
+
+## Encrypted Storage and Updatable VMs
+
+With [Updatable VM feature][updatable_vm] shipping in Android V, Encrypted Storage can be accessed
+even after OTA/updates of boot images and apks. This requires chipsets to support [Secretkeeper HAL][sk_hal].
+
+
+[vm_payload_api]: https://cs.android.com/android/platform/superproject/main/+/main:packages/modules/Virtualization/vm_payload/include/vm_payload.h;l=2?q=vm_payload%2Finclude%2Fvm_payload.h&ss=android%2Fplatform%2Fsuperproject%2Fmain
+[updatable_vm]: https://cs.android.com/android/platform/superproject/main/+/main:packages/modules/Virtualization/docs/updatable_vm.md
+[sk_hal]: https://cs.android.com/android/platform/superproject/main/+/main:system/secretkeeper/README.md
diff --git a/java/framework/README.md b/java/framework/README.md
index cf7a6cb..bbcd0ef 100644
--- a/java/framework/README.md
+++ b/java/framework/README.md
@@ -339,6 +339,8 @@
 powerful attacker could delete it, wholly or partially roll it back to an
 earlier version, or modify it, corrupting the data.
 
+For more info see [README](https://cs.android.com/android/platform/superproject/main/+/main:packages/modules/Virtualization/java/framework/README.md)
+
 # Transferring a VM
 
 It is possible to make a copy of a VM instance. This can be used to transfer a
diff --git a/libs/devicemapper/src/crypt.rs b/libs/devicemapper/src/crypt.rs
index 36c45c7..3afd374 100644
--- a/libs/devicemapper/src/crypt.rs
+++ b/libs/devicemapper/src/crypt.rs
@@ -15,8 +15,8 @@
  */
 
 /// `crypt` module implements the "crypt" target in the device mapper framework. Specifically,
-/// it provides `DmCryptTargetBuilder` struct which is used to construct a `DmCryptTarget` struct
-/// which is then given to `DeviceMapper` to create a mapper device.
+/// it provides `DmCryptTargetBuilder` struct which is used to construct a `DmCryptTarget`
+/// struct which is then given to `DeviceMapper` to create a mapper device.
 use crate::DmTargetSpec;
 
 use anyhow::{ensure, Context, Result};
@@ -33,9 +33,14 @@
 /// Supported ciphers
 #[derive(Clone, Copy, Debug)]
 pub enum CipherType {
-    // AES-256-HCTR2 takes a 32-byte key
+    /// AES256 with HCTR2 mode. HCTR2 is a tweakable super-pseudorandom permutation
+    /// length-preserving encryption mode. It is the preferred mode in absence of other
+    /// dedicated integrity primitives (such as for encryptedstore in pVM) since it is less
+    /// malleable than other modes.
     AES256HCTR2,
-    // XTS requires key of twice the length of the underlying block cipher i.e., 64B for AES256
+    /// AES with XTS mode. This has slight performance benefits over HCTR2. In particular, XTS is
+    /// supported by inline encryption hardware. Note that (status quo) `encryptedstore` in VMs
+    /// is the only user of this module & inline encryption is not supported by guest kernel.
     AES256XTS,
 }
 impl CipherType {
@@ -50,7 +55,10 @@
 
     fn get_required_key_size(&self) -> usize {
         match *self {
+            // AES-256-HCTR2 takes a 32-byte key
             CipherType::AES256HCTR2 => 32,
+            // XTS requires key of twice the length of the underlying block cipher
+            // i.e., 64B for AES256
             CipherType::AES256XTS => 64,
         }
     }
diff --git a/microdroid_manager/src/dice.rs b/microdroid_manager/src/dice.rs
index 7f65159..cecf413 100644
--- a/microdroid_manager/src/dice.rs
+++ b/microdroid_manager/src/dice.rs
@@ -13,12 +13,12 @@
 // limitations under the License.
 
 use crate::instance::{ApexData, ApkData};
-use crate::{is_debuggable, MicrodroidData};
+use crate::{is_debuggable, is_strict_boot, MicrodroidData};
 use anyhow::{bail, Context, Result};
 use ciborium::{cbor, Value};
 use coset::CborSerializable;
 use dice_driver::DiceDriver;
-use diced_open_dice::OwnedDiceArtifacts;
+use diced_open_dice::{Hidden, OwnedDiceArtifacts, HIDDEN_SIZE};
 use microdroid_metadata::PayloadMetadata;
 use openssl::sha::{sha512, Sha512};
 use std::iter::once;
@@ -53,10 +53,37 @@
     let debuggable = is_debuggable()?;
 
     // Send the details to diced
-    let hidden = instance_data.salt.clone().try_into().unwrap();
+    let hidden = if cfg!(llpvm_changes) {
+        hidden_input_from_instance_id()?
+    } else {
+        instance_data.salt.clone().try_into().unwrap()
+    };
     dice.derive(code_hash, &config_descriptor, authority_hash, debuggable, hidden)
 }
 
+// Get the "Hidden input" for DICE derivation.
+// This provides differentiation of secrets for different VM instances with same payload.
+fn hidden_input_from_instance_id() -> Result<Hidden> {
+    // For protected VM: this is all 0s, pvmfw ensures differentiation is added early in secrets.
+    // For non-protected VM: this is derived from instance_id of the VM instance.
+    let hidden_input = if !is_strict_boot() {
+        if let Some(id) = super::get_instance_id()? {
+            sha512(&id)
+        } else {
+            // TODO(b/325094712): Absence of instance_id occurs due to missing DT in some
+            // x86_64 test devices (such as Cuttlefish). From security perspective, this is
+            // acceptable for non-protected VM.
+            log::warn!(
+                "Instance Id missing, this may lead to 2 non protected VMs having same secrets"
+            );
+            [0u8; HIDDEN_SIZE]
+        }
+    } else {
+        [0u8; HIDDEN_SIZE]
+    };
+    Ok(hidden_input)
+}
+
 struct Subcomponent {
     name: String,
     version: u64,
diff --git a/microdroid_manager/src/instance.rs b/microdroid_manager/src/instance.rs
index 888c451..2d39cd8 100644
--- a/microdroid_manager/src/instance.rs
+++ b/microdroid_manager/src/instance.rs
@@ -273,6 +273,8 @@
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct MicrodroidData {
+    // `salt` is obsolete, it was used as a differentiator for non-protected VM instances running
+    // same payload. Instance-id (present in DT) is used for that now.
     pub salt: Vec<u8>, // Should be [u8; 64] but that isn't serializable.
     pub apk_data: ApkData,
     pub extra_apks_data: Vec<ApkData>,
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 8d2c629..2386bd4 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -41,7 +41,7 @@
 use keystore2_crypto::ZVec;
 use libc::VMADDR_CID_HOST;
 use log::{error, info};
-use microdroid_metadata::PayloadMetadata;
+use microdroid_metadata::{Metadata, PayloadMetadata};
 use microdroid_payload_config::{ApkConfig, OsConfig, Task, TaskType, VmPayloadConfig};
 use nix::sys::signal::Signal;
 use payload::load_metadata;
@@ -72,6 +72,7 @@
     "/proc/device-tree/virtualization/guest/debug-microdroid,no-verified-boot";
 const SECRETKEEPER_KEY: &str = "/proc/device-tree/avf/secretkeeper_public_key";
 const INSTANCE_ID_PATH: &str = "/proc/device-tree/avf/untrusted/instance-id";
+const DEFER_ROLLBACK_PROTECTION: &str = "/proc/device-tree/avf/untrusted/defer-rollback-protection";
 
 const ENCRYPTEDSTORE_BIN: &str = "/system/bin/encryptedstore";
 const ZIPFUSE_BIN: &str = "/system/bin/zipfuse";
@@ -161,6 +162,10 @@
     Ok(instance_id)
 }
 
+fn should_defer_rollback_protection() -> bool {
+    Path::new(DEFER_ROLLBACK_PROTECTION).exists()
+}
+
 fn main() -> Result<()> {
     // If debuggable, print full backtrace to console log with stdio_to_kmsg
     if is_debuggable()? {
@@ -235,17 +240,12 @@
     }
 }
 
-fn try_run_payload(
-    service: &Strong<dyn IVirtualMachineService>,
-    vm_payload_service_fd: OwnedFd,
-) -> Result<i32> {
-    let metadata = load_metadata().context("Failed to load payload metadata")?;
-    let dice = DiceDriver::new(Path::new("/dev/open-dice0"), is_strict_boot())
-        .context("Failed to load DICE")?;
-
+fn verify_payload_with_instance_img(
+    metadata: &Metadata,
+    dice: &DiceDriver,
+) -> Result<MicrodroidData> {
     let mut instance = InstanceDisk::new().context("Failed to load instance.img")?;
-    let saved_data =
-        instance.read_microdroid_data(&dice).context("Failed to read identity data")?;
+    let saved_data = instance.read_microdroid_data(dice).context("Failed to read identity data")?;
 
     if is_strict_boot() {
         // Provisioning must happen on the first boot and never again.
@@ -265,7 +265,7 @@
     }
 
     // Verify the payload before using it.
-    let extracted_data = verify_payload(&metadata, saved_data.as_ref())
+    let extracted_data = verify_payload(metadata, saved_data.as_ref())
         .context("Payload verification failed")
         .map_err(|e| MicrodroidError::PayloadVerificationFailed(e.to_string()))?;
 
@@ -289,10 +289,29 @@
     } else {
         info!("Saving verified data.");
         instance
-            .write_microdroid_data(&extracted_data, &dice)
+            .write_microdroid_data(&extracted_data, dice)
             .context("Failed to write identity data")?;
         extracted_data
     };
+    Ok(instance_data)
+}
+
+fn try_run_payload(
+    service: &Strong<dyn IVirtualMachineService>,
+    vm_payload_service_fd: OwnedFd,
+) -> Result<i32> {
+    let metadata = load_metadata().context("Failed to load payload metadata")?;
+    let dice = DiceDriver::new(Path::new("/dev/open-dice0"), is_strict_boot())
+        .context("Failed to load DICE")?;
+
+    // Microdroid skips checking payload against instance image iff the device supports
+    // secretkeeper. In that case Microdroid use VmSecret::V2, which provide protection against
+    // rollback of boot images and packages.
+    let instance_data = if should_defer_rollback_protection() {
+        verify_payload(&metadata, None)?
+    } else {
+        verify_payload_with_instance_img(&metadata, &dice)?
+    };
 
     let payload_metadata = metadata.payload.ok_or_else(|| {
         MicrodroidError::PayloadInvalidConfig("No payload config in metadata".to_string())
diff --git a/microdroid_manager/src/verify.rs b/microdroid_manager/src/verify.rs
index 445c1ae..65c32b0 100644
--- a/microdroid_manager/src/verify.rs
+++ b/microdroid_manager/src/verify.rs
@@ -169,13 +169,14 @@
     // verified is consistent with the root hash) or because we have the saved APK data which will
     // be checked as identical to the data we have verified.
 
-    // 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 if is_strict_boot() {
-        // No need to add more entropy as a previous stage must have used a new, random salt.
+    let salt = if cfg!(llpvm_changes) || is_strict_boot() {
+        // Salt is obsolete with llpvm_changes.
         vec![0u8; 64]
+    } else if let Some(saved_data) = saved_data {
+        // Use the salt from a verified instance.
+        saved_data.salt.clone()
     } else {
+        // Generate a salt for a new instance.
         let mut salt = vec![0u8; 64];
         salt.as_mut_slice().try_fill(&mut rand::thread_rng())?;
         salt
diff --git a/microdroid_manager/src/vm_secret.rs b/microdroid_manager/src/vm_secret.rs
index 5ceedea..ed8ab1d 100644
--- a/microdroid_manager/src/vm_secret.rs
+++ b/microdroid_manager/src/vm_secret.rs
@@ -91,28 +91,19 @@
         vm_service: &Strong<dyn IVirtualMachineService>,
     ) -> Result<Self> {
         ensure!(dice_artifacts.bcc().is_some(), "Dice chain missing");
-
-        let Some(sk_service) =
-            is_sk_supported(vm_service).context("Failed to check if Secretkeeper is supported")?
-        else {
-            // Use V1 secrets if Secretkeeper is not supported.
+        if !crate::should_defer_rollback_protection() {
             return Ok(Self::V1 { dice_artifacts });
-        };
+        }
 
         let explicit_dice = OwnedDiceArtifactsWithExplicitKey::from_owned_artifacts(dice_artifacts)
             .context("Failed to get Dice artifacts in explicit key format")?;
         // For pVM, skp_secret are stored in Secretkeeper. For non-protected it is all 0s.
         let mut skp_secret = Zeroizing::new([0u8; SECRET_SIZE]);
         if super::is_strict_boot() {
-            let mut session = SkSession::new(
-                sk_service,
-                &explicit_dice,
-                Some(get_secretkeeper_identity().context("Failed to get secretkeeper identity")?),
-            )
-            .context("Failed to setup a Secretkeeper session")?;
-            let id = super::get_instance_id()
-                .context("Failed to get instance-id")?
-                .ok_or(anyhow!("Missing instance-id"))?;
+            let sk_service = get_secretkeeper_service(vm_service)?;
+            let mut session =
+                SkSession::new(sk_service, &explicit_dice, Some(get_secretkeeper_identity()?))?;
+            let id = super::get_instance_id()?.ok_or(anyhow!("Missing instance_id"))?;
             let explicit_dice_chain = explicit_dice
                 .explicit_key_dice_chain()
                 .ok_or(anyhow!("Missing explicit dice chain, this is unusual"))?;
@@ -279,22 +270,15 @@
     anyhow!("{:?}", err)
 }
 
-// Get the secretkeeper connection if supported. Host can be consulted whether the device supports
-// secretkeeper but that should be used with caution for protected VM.
-fn is_sk_supported(
+fn get_secretkeeper_service(
     host: &Strong<dyn IVirtualMachineService>,
-) -> Result<Option<Strong<dyn ISecretkeeper>>> {
-    let sk = if cfg!(llpvm_changes) {
-        host.getSecretkeeper()
-            // TODO rename this error!
-            .map_err(|e| {
-                super::MicrodroidError::FailedToConnectToVirtualizationService(format!(
-                    "Failed to get Secretkeeper: {e:?}"
-                ))
-            })?
-    } else {
-        // LLPVM flag is disabled
-        None
-    };
-    Ok(sk)
+) -> Result<Strong<dyn ISecretkeeper>> {
+    Ok(host
+        .getSecretkeeper()
+        // TODO rename this error!
+        .map_err(|e| {
+            super::MicrodroidError::FailedToConnectToVirtualizationService(format!(
+                "Failed to get Secretkeeper: {e:?}"
+            ))
+        })?)
 }
diff --git a/pvmfw/src/instance.rs b/pvmfw/src/instance.rs
index 6daadd9..43c7442 100644
--- a/pvmfw/src/instance.rs
+++ b/pvmfw/src/instance.rs
@@ -27,7 +27,6 @@
 use log::trace;
 use uuid::Uuid;
 use virtio_drivers::transport::{pci::bus::PciRoot, DeviceType, Transport};
-use vmbase::rand;
 use vmbase::util::ceiling_div;
 use vmbase::virtio::pci::{PciTransportIterator, VirtIOBlk};
 use vmbase::virtio::HalImpl;
@@ -38,8 +37,6 @@
 pub enum Error {
     /// Unexpected I/O error while accessing the underlying disk.
     FailedIo(gpt::Error),
-    /// Failed to generate a random salt to be stored.
-    FailedSaltGeneration(rand::Error),
     /// Impossible to create a new instance.img entry.
     InstanceImageFull,
     /// Badly formatted instance.img header block.
@@ -66,7 +63,6 @@
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Self::FailedIo(e) => write!(f, "Failed I/O to disk: {e}"),
-            Self::FailedSaltGeneration(e) => write!(f, "Failed to generate salt: {e}"),
             Self::InstanceImageFull => write!(f, "Failed to obtain a free instance.img partition"),
             Self::InvalidInstanceImageHeader => write!(f, "instance.img header is invalid"),
             Self::MissingInstanceImage => write!(f, "Failed to find the instance.img partition"),
@@ -93,27 +89,27 @@
 
 pub type Result<T> = core::result::Result<T, Error>;
 
-pub fn get_or_generate_instance_salt(
+fn aead_ctx_from_secret(secret: &[u8]) -> Result<AeadContext> {
+    let key = hkdf::<32>(secret, /* salt= */ &[], b"vm-instance", Digester::sha512())?;
+    Ok(AeadContext::new(Aead::aes_256_gcm_randnonce(), key.as_slice(), /* tag_len */ None)?)
+}
+
+/// Get the entry from instance.img. This method additionally returns Partition corresponding to
+/// pvmfw in the instance.img as well as index corresponding to empty header which can be used to
+/// record instance data with `record_instance_entry`.
+pub(crate) fn get_recorded_entry(
     pci_root: &mut PciRoot,
-    dice_inputs: &PartialInputs,
     secret: &[u8],
-) -> Result<(bool, Hidden)> {
+) -> Result<(Option<EntryBody>, Partition, usize)> {
     let mut instance_img = find_instance_img(pci_root)?;
 
     let entry = locate_entry(&mut instance_img)?;
     trace!("Found pvmfw instance.img entry: {entry:?}");
 
-    let key = hkdf::<32>(secret, /* salt= */ &[], b"vm-instance", Digester::sha512())?;
-    let tag_len = None;
-    let aead_ctx = AeadContext::new(Aead::aes_256_gcm_randnonce(), key.as_slice(), tag_len)?;
-    let ad = &[];
-    // The nonce is generated internally for `aes_256_gcm_randnonce`, so no additional
-    // nonce is required.
-    let nonce = &[];
-
-    let mut blk = [0; BLK_SIZE];
     match entry {
         PvmfwEntry::Existing { header_index, payload_size } => {
+            let aead_ctx = aead_ctx_from_secret(secret)?;
+            let mut blk = [0; BLK_SIZE];
             if payload_size > blk.len() {
                 // We currently only support single-blk entries.
                 return Err(Error::UnsupportedEntrySize(payload_size));
@@ -123,52 +119,41 @@
 
             let payload = &blk[..payload_size];
             let mut entry = [0; size_of::<EntryBody>()];
-            let decrypted = aead_ctx.open(payload, nonce, ad, &mut entry)?;
-
+            // The nonce is generated internally for `aes_256_gcm_randnonce`, so no additional
+            // nonce is required.
+            let decrypted =
+                aead_ctx.open(payload, /* nonce */ &[], /* ad */ &[], &mut entry)?;
             let body = EntryBody::read_from(decrypted).unwrap();
-            if dice_inputs.rkp_vm_marker {
-                // The RKP VM is allowed to run if it has passed the verified boot check and
-                // contains the expected version in its AVB footer.
-                // The comparison below with the previous boot information is skipped to enable the
-                // simultaneous update of the pvmfw and RKP VM.
-                // For instance, when both the pvmfw and RKP VM are updated, the code hash of the
-                // RKP VM will differ from the one stored in the instance image. In this case, the
-                // RKP VM is still allowed to run.
-                // This ensures that the updated RKP VM will retain the same CDIs in the next stage.
-                return Ok((false, body.salt));
-            }
-            if body.code_hash != dice_inputs.code_hash {
-                Err(Error::RecordedCodeHashMismatch)
-            } else if body.auth_hash != dice_inputs.auth_hash {
-                Err(Error::RecordedAuthHashMismatch)
-            } else if body.mode() != dice_inputs.mode {
-                Err(Error::RecordedDiceModeMismatch)
-            } else {
-                Ok((false, body.salt))
-            }
+            Ok((Some(body), instance_img, header_index))
         }
-        PvmfwEntry::New { header_index } => {
-            let salt = rand::random_array().map_err(Error::FailedSaltGeneration)?;
-            let body = EntryBody::new(dice_inputs, &salt);
-
-            // We currently only support single-blk entries.
-            let plaintext = body.as_bytes();
-            assert!(plaintext.len() + aead_ctx.aead().max_overhead() < blk.len());
-            let encrypted = aead_ctx.seal(plaintext, nonce, ad, &mut blk)?;
-            let payload_size = encrypted.len();
-            let payload_index = header_index + 1;
-            instance_img.write_block(payload_index, &blk).map_err(Error::FailedIo)?;
-
-            let header = EntryHeader::new(PvmfwEntry::UUID, payload_size);
-            header.write_to_prefix(blk.as_mut_slice()).unwrap();
-            blk[header.as_bytes().len()..].fill(0);
-            instance_img.write_block(header_index, &blk).map_err(Error::FailedIo)?;
-
-            Ok((true, salt))
-        }
+        PvmfwEntry::New { header_index } => Ok((None, instance_img, header_index)),
     }
 }
 
+pub(crate) fn record_instance_entry(
+    body: &EntryBody,
+    secret: &[u8],
+    instance_img: &mut Partition,
+    header_index: usize,
+) -> Result<()> {
+    // We currently only support single-blk entries.
+    let mut blk = [0; BLK_SIZE];
+    let plaintext = body.as_bytes();
+    let aead_ctx = aead_ctx_from_secret(secret)?;
+    assert!(plaintext.len() + aead_ctx.aead().max_overhead() < blk.len());
+    let encrypted = aead_ctx.seal(plaintext, /* nonce */ &[], /* ad */ &[], &mut blk)?;
+    let payload_size = encrypted.len();
+    let payload_index = header_index + 1;
+    instance_img.write_block(payload_index, &blk).map_err(Error::FailedIo)?;
+
+    let header = EntryHeader::new(PvmfwEntry::UUID, payload_size);
+    header.write_to_prefix(blk.as_mut_slice()).unwrap();
+    blk[header.as_bytes().len()..].fill(0);
+    instance_img.write_block(header_index, &blk).map_err(Error::FailedIo)?;
+
+    Ok(())
+}
+
 #[derive(FromZeroes, FromBytes)]
 #[repr(C, packed)]
 struct Header {
@@ -276,15 +261,15 @@
 
 #[derive(AsBytes, FromZeroes, FromBytes)]
 #[repr(C)]
-struct EntryBody {
-    code_hash: Hash,
-    auth_hash: Hash,
-    salt: Hidden,
+pub(crate) struct EntryBody {
+    pub code_hash: Hash,
+    pub auth_hash: Hash,
+    pub salt: Hidden,
     mode: u8,
 }
 
 impl EntryBody {
-    fn new(dice_inputs: &PartialInputs, salt: &Hidden) -> Self {
+    pub(crate) fn new(dice_inputs: &PartialInputs, salt: &Hidden) -> Self {
         let mode = match dice_inputs.mode {
             DiceMode::kDiceModeNotInitialized => 0,
             DiceMode::kDiceModeNormal => 1,
@@ -300,7 +285,7 @@
         }
     }
 
-    fn mode(&self) -> DiceMode {
+    pub(crate) fn mode(&self) -> DiceMode {
         match self.mode {
             1 => DiceMode::kDiceModeNormal,
             2 => DiceMode::kDiceModeDebug,
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index f80bae1..12d63d5 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -37,7 +37,9 @@
 use crate::entry::RebootReason;
 use crate::fdt::modify_for_next_stage;
 use crate::helpers::GUEST_PAGE_SIZE;
-use crate::instance::get_or_generate_instance_salt;
+use crate::instance::EntryBody;
+use crate::instance::Error as InstanceError;
+use crate::instance::{get_recorded_entry, record_instance_entry};
 use alloc::borrow::Cow;
 use alloc::boxed::Box;
 use core::ops::Range;
@@ -150,11 +152,43 @@
         error!("Failed to compute partial DICE inputs: {e:?}");
         RebootReason::InternalError
     })?;
-    let (new_instance, salt) = get_or_generate_instance_salt(&mut pci_root, &dice_inputs, cdi_seal)
-        .map_err(|e| {
-            error!("Failed to get instance.img salt: {e}");
+
+    let (recorded_entry, mut instance_img, header_index) =
+        get_recorded_entry(&mut pci_root, cdi_seal).map_err(|e| {
+            error!("Failed to get entry from instance.img: {e}");
             RebootReason::InternalError
         })?;
+    let (new_instance, salt) = if let Some(entry) = recorded_entry {
+        // The RKP VM is allowed to run if it has passed the verified boot check and
+        // contains the expected version in its AVB footer.
+        // The comparison below with the previous boot information is skipped to enable the
+        // simultaneous update of the pvmfw and RKP VM.
+        // For instance, when both the pvmfw and RKP VM are updated, the code hash of the
+        // RKP VM will differ from the one stored in the instance image. In this case, the
+        // RKP VM is still allowed to run.
+        // This ensures that the updated RKP VM will retain the same CDIs in the next stage.
+        if !dice_inputs.rkp_vm_marker {
+            ensure_dice_measurements_match_entry(&dice_inputs, &entry).map_err(|e| {
+                error!(
+                    "Dice measurements do not match recorded entry.
+                This may be because of update: {e}"
+                );
+                RebootReason::InternalError
+            })?;
+        }
+        (false, entry.salt)
+    } else {
+        let salt = rand::random_array().map_err(|e| {
+            error!("Failed to generated instance.img salt: {e}");
+            RebootReason::InternalError
+        })?;
+        let entry = EntryBody::new(&dice_inputs, &salt);
+        record_instance_entry(&entry, cdi_seal, &mut instance_img, header_index).map_err(|e| {
+            error!("Failed to get recorded entry in instance.img: {e}");
+            RebootReason::InternalError
+        })?;
+        (true, salt)
+    };
     trace!("Got salt from instance.img: {salt:x?}");
 
     let new_bcc_handover = if cfg!(dice_changes) {
@@ -207,6 +241,21 @@
     Ok(bcc_range)
 }
 
+fn ensure_dice_measurements_match_entry(
+    dice_inputs: &PartialInputs,
+    entry: &EntryBody,
+) -> Result<(), InstanceError> {
+    if entry.code_hash != dice_inputs.code_hash {
+        Err(InstanceError::RecordedCodeHashMismatch)
+    } else if entry.auth_hash != dice_inputs.auth_hash {
+        Err(InstanceError::RecordedAuthHashMismatch)
+    } else if entry.mode() != dice_inputs.mode {
+        Err(InstanceError::RecordedDiceModeMismatch)
+    } else {
+        Ok(())
+    }
+}
+
 /// Logs the given PCI error and returns the appropriate `RebootReason`.
 fn handle_pci_error(e: PciError) -> RebootReason {
     error!("{}", e);
diff --git a/service_vm/test_apk/Android.bp b/service_vm/test_apk/Android.bp
index de731f6..cd992db 100644
--- a/service_vm/test_apk/Android.bp
+++ b/service_vm/test_apk/Android.bp
@@ -57,6 +57,7 @@
     static_libs: [
         "RkpdAppTestUtil",
         "androidx.work_work-testing",
+        "bouncycastle-unbundled",
     ],
     instrumentation_for: "rkpdapp",
     // This app is a variation of rkpdapp, with additional permissions to run
diff --git a/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java b/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java
index e7061e1..2a771f3 100644
--- a/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java
+++ b/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java
@@ -44,6 +44,12 @@
 import com.android.virt.vm_attestation.testservice.IAttestationService;
 import com.android.virt.vm_attestation.testservice.IAttestationService.SigningResult;
 
+import org.bouncycastle.asn1.ASN1Boolean;
+import org.bouncycastle.asn1.ASN1Encodable;
+import org.bouncycastle.asn1.ASN1OctetString;
+import org.bouncycastle.asn1.ASN1Sequence;
+import org.bouncycastle.asn1.DEROctetString;
+import org.bouncycastle.asn1.DERUTF8String;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -215,28 +221,46 @@
         assertThat(sig.verify(signingResult.signature)).isTrue();
     }
 
-    private void checkAvfAttestationExtension(X509Certificate cert, byte[] challenge) {
-        byte[] extension = cert.getExtensionValue(AVF_ATTESTATION_EXTENSION_OID);
-        assertThat(extension).isNotNull();
-        // TODO(b/325610326): Use bouncycastle to parse the extension and check other fields.
-        assertWithMessage("The extension should contain the challenge")
-                .that(containsSubarray(extension, challenge))
-                .isTrue();
+    private void checkAvfAttestationExtension(X509Certificate cert, byte[] challenge)
+            throws Exception {
+        byte[] extensionValue = cert.getExtensionValue(AVF_ATTESTATION_EXTENSION_OID);
+        ASN1OctetString extString = ASN1OctetString.getInstance(extensionValue);
+        ASN1Sequence seq = ASN1Sequence.getInstance(extString.getOctets());
+        // AVF attestation extension should contain 3 elements in the following format:
+        //
+        //  AttestationExtension ::= SEQUENCE {
+        //     attestationChallenge       OCTET_STRING,
+        //     isVmSecure                 BOOLEAN,
+        //     vmComponents               SEQUENCE OF VmComponent,
+        //  }
+        //   VmComponent ::= SEQUENCE {
+        //     name               UTF8String,
+        //     securityVersion    INTEGER,
+        //     codeHash           OCTET STRING,
+        //     authorityHash      OCTET STRING,
+        //  }
+        assertThat(seq).hasSize(3);
+
+        ASN1OctetString expectedChallenge = new DEROctetString(challenge);
+        assertThat(seq.getObjectAt(0)).isEqualTo(expectedChallenge);
+        assertWithMessage("The VM should be unsecure as it is debuggable.")
+                .that(seq.getObjectAt(1))
+                .isEqualTo(ASN1Boolean.FALSE);
+        ASN1Sequence vmComponents = ASN1Sequence.getInstance(seq.getObjectAt(2));
+        assertExtensionContainsPayloadApk(vmComponents);
     }
 
-    private boolean containsSubarray(byte[] array, byte[] subarray) {
-        for (int i = 0; i < array.length - subarray.length + 1; i++) {
-            boolean found = true;
-            for (int j = 0; j < subarray.length; j++) {
-                if (array[i + j] != subarray[j]) {
-                    found = false;
-                    break;
-                }
-            }
-            if (found) {
-                return true;
+    private void assertExtensionContainsPayloadApk(ASN1Sequence vmComponents) throws Exception {
+        DERUTF8String payloadApkName = new DERUTF8String("apk:" + TEST_APP_PACKAGE_NAME);
+        boolean found = false;
+        for (ASN1Encodable encodable : vmComponents) {
+            ASN1Sequence vmComponent = ASN1Sequence.getInstance(encodable);
+            assertThat(vmComponent).hasSize(4);
+            if (payloadApkName.equals(vmComponent.getObjectAt(0))) {
+                assertWithMessage("Payload APK should not be found twice.").that(found).isFalse();
+                found = true;
             }
         }
-        return false;
+        assertWithMessage("vmComponents should contain the payload APK.").that(found).isTrue();
     }
 }
diff --git a/tests/pvmfw/Android.bp b/tests/pvmfw/Android.bp
index c12f67a..0483066 100644
--- a/tests/pvmfw/Android.bp
+++ b/tests/pvmfw/Android.bp
@@ -53,4 +53,5 @@
         ":test_avf_debug_policy_without_adb",
         "assets/bcc.dat",
     ],
+    data_device_bins_first: ["dtc_static"],
 }
diff --git a/tests/pvmfw/AndroidTest.xml b/tests/pvmfw/AndroidTest.xml
index 6ff7b6f..5784f26 100644
--- a/tests/pvmfw/AndroidTest.xml
+++ b/tests/pvmfw/AndroidTest.xml
@@ -22,6 +22,24 @@
         <option name="force-root" value="true"/>
     </target_preparer>
 
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="throw-if-cmd-fail" value="true" />
+        <!-- Prepare test directories. -->
+        <option name="run-command" value="mkdir -p /data/local/tmp/pvmfw" />
+        <option name="teardown-command" value="rm -rf /data/local/tmp/pvmfw" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.PushFilePreparer">
+        <option name="cleanup" value="true" />
+        <option name="abort-on-push-failure" value="true" />
+        <option name="push-file" key="dtc_static" value="/data/local/tmp/pvmfw/dtc_static" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="throw-if-cmd-fail" value="true" />
+        <option name="run-command" value="[ ! -d /proc/device-tree/avf/reference ] || /data/local/tmp/pvmfw/dtc_static -f -qqq /proc/device-tree/avf/reference -o /data/local/tmp/pvmfw/reference_dt.dtb" />
+    </target_preparer>
+
     <test class="com.android.compatibility.common.tradefed.testtype.JarHostTest" >
         <option name="jar" value="CustomPvmfwHostTestCases.jar" />
     </test>
diff --git a/tests/pvmfw/helper/Android.bp b/tests/pvmfw/helper/Android.bp
index 1b96842..90ca03e 100644
--- a/tests/pvmfw/helper/Android.bp
+++ b/tests/pvmfw/helper/Android.bp
@@ -5,5 +5,8 @@
 java_library_host {
     name: "PvmfwHostTestHelper",
     srcs: ["java/**/*.java"],
-    libs: ["androidx.annotation_annotation"],
+    libs: [
+        "androidx.annotation_annotation",
+        "truth",
+    ],
 }
diff --git a/tests/pvmfw/helper/java/com/android/pvmfw/test/host/Pvmfw.java b/tests/pvmfw/helper/java/com/android/pvmfw/test/host/Pvmfw.java
index b0c1207..a77ba40 100644
--- a/tests/pvmfw/helper/java/com/android/pvmfw/test/host/Pvmfw.java
+++ b/tests/pvmfw/helper/java/com/android/pvmfw/test/host/Pvmfw.java
@@ -16,6 +16,8 @@
 
 package com.android.pvmfw.test.host;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static java.nio.ByteOrder.LITTLE_ENDIAN;
 
 import androidx.annotation.NonNull;
@@ -34,22 +36,52 @@
     private static final int SIZE_4K = 4 << 10; // 4 KiB, PAGE_SIZE
     private static final int BUFFER_SIZE = 1024;
     private static final int HEADER_MAGIC = 0x666d7670;
-    private static final int HEADER_DEFAULT_VERSION = getVersion(1, 0);
+    private static final int HEADER_DEFAULT_VERSION = makeVersion(1, 2);
     private static final int HEADER_FLAGS = 0;
 
+    private static final int PVMFW_ENTRY_BCC = 0;
+    private static final int PVMFW_ENTRY_DP = 1;
+    private static final int PVMFW_ENTRY_VM_DTBO = 2;
+    private static final int PVMFW_ENTRY_VM_REFERENCE_DT = 3;
+    private static final int PVMFW_ENTRY_MAX = 4;
+
     @NonNull private final File mPvmfwBinFile;
-    @NonNull private final File mBccFile;
-    @Nullable private final File mDebugPolicyFile;
+    private final File[] mEntries;
+    private final int mEntryCnt;
     private final int mVersion;
 
+    public static int makeVersion(int major, int minor) {
+        return ((major & 0xFFFF) << 16) | (minor & 0xFFFF);
+    }
+
     private Pvmfw(
             @NonNull File pvmfwBinFile,
             @NonNull File bccFile,
             @Nullable File debugPolicyFile,
+            @Nullable File vmDtboFile,
+            @Nullable File vmReferenceDtFile,
             int version) {
         mPvmfwBinFile = Objects.requireNonNull(pvmfwBinFile);
-        mBccFile = Objects.requireNonNull(bccFile);
-        mDebugPolicyFile = debugPolicyFile;
+
+        if (version >= makeVersion(1, 2)) {
+            mEntryCnt = PVMFW_ENTRY_VM_REFERENCE_DT + 1;
+        } else if (version >= makeVersion(1, 1)) {
+            mEntryCnt = PVMFW_ENTRY_VM_DTBO + 1;
+        } else {
+            mEntryCnt = PVMFW_ENTRY_DP + 1;
+        }
+
+        mEntries = new File[PVMFW_ENTRY_MAX];
+        mEntries[PVMFW_ENTRY_BCC] = Objects.requireNonNull(bccFile);
+        mEntries[PVMFW_ENTRY_DP] = debugPolicyFile;
+
+        if (PVMFW_ENTRY_VM_DTBO < mEntryCnt) {
+            mEntries[PVMFW_ENTRY_VM_DTBO] = vmDtboFile;
+        }
+        if (PVMFW_ENTRY_VM_REFERENCE_DT < mEntryCnt) {
+            mEntries[PVMFW_ENTRY_VM_REFERENCE_DT] = Objects.requireNonNull(vmReferenceDtFile);
+        }
+
         mVersion = version;
     }
 
@@ -60,62 +92,54 @@
     public void serialize(@NonNull File outFile) throws IOException {
         Objects.requireNonNull(outFile);
 
-        int headerSize = alignTo(getHeaderSize(mVersion), SIZE_8B);
-        int bccOffset = headerSize;
-        int bccSize = (int) mBccFile.length();
+        int headerSize = alignTo(getHeaderSize(), SIZE_8B);
+        int[] entryOffsets = new int[mEntryCnt];
+        int[] entrySizes = new int[mEntryCnt];
 
-        int debugPolicyOffset = alignTo(bccOffset + bccSize, SIZE_8B);
-        int debugPolicySize = mDebugPolicyFile == null ? 0 : (int) mDebugPolicyFile.length();
+        entryOffsets[PVMFW_ENTRY_BCC] = headerSize;
+        entrySizes[PVMFW_ENTRY_BCC] = (int) mEntries[PVMFW_ENTRY_BCC].length();
 
-        int totalSize = debugPolicyOffset + debugPolicySize;
-        if (hasVmDtbo(mVersion)) {
-            // Add VM DTBO size as well.
-            totalSize += Integer.BYTES * 2;
+        for (int i = 1; i < mEntryCnt; i++) {
+            entryOffsets[i] = alignTo(entryOffsets[i - 1] + entrySizes[i - 1], SIZE_8B);
+            entrySizes[i] = mEntries[i] == null ? 0 : (int) mEntries[i].length();
         }
 
+        int totalSize = alignTo(entryOffsets[mEntryCnt - 1] + entrySizes[mEntryCnt - 1], SIZE_8B);
+
         ByteBuffer header = ByteBuffer.allocate(headerSize).order(LITTLE_ENDIAN);
         header.putInt(HEADER_MAGIC);
         header.putInt(mVersion);
         header.putInt(totalSize);
         header.putInt(HEADER_FLAGS);
-        header.putInt(bccOffset);
-        header.putInt(bccSize);
-        header.putInt(debugPolicyOffset);
-        header.putInt(debugPolicySize);
-
-        if (hasVmDtbo(mVersion)) {
-            // Add placeholder entry for VM DTBO.
-            // TODO(b/291191157): Add a real DTBO and test.
-            header.putInt(0);
-            header.putInt(0);
+        for (int i = 0; i < mEntryCnt; i++) {
+            header.putInt(entryOffsets[i]);
+            header.putInt(entrySizes[i]);
         }
 
         try (FileOutputStream pvmfw = new FileOutputStream(outFile)) {
             appendFile(pvmfw, mPvmfwBinFile);
             padTo(pvmfw, SIZE_4K);
+
+            int baseOffset = (int) pvmfw.getChannel().size();
             pvmfw.write(header.array());
-            padTo(pvmfw, SIZE_8B);
-            appendFile(pvmfw, mBccFile);
-            if (mDebugPolicyFile != null) {
+
+            for (int i = 0; i < mEntryCnt; i++) {
                 padTo(pvmfw, SIZE_8B);
-                appendFile(pvmfw, mDebugPolicyFile);
+                if (mEntries[i] != null) {
+                    assertThat((int) pvmfw.getChannel().size() - baseOffset)
+                            .isEqualTo(entryOffsets[i]);
+                    appendFile(pvmfw, mEntries[i]);
+                }
             }
+
             padTo(pvmfw, SIZE_4K);
         }
     }
 
     private void appendFile(@NonNull FileOutputStream out, @NonNull File inFile)
             throws IOException {
-        byte buffer[] = new byte[BUFFER_SIZE];
         try (FileInputStream in = new FileInputStream(inFile)) {
-            int size;
-            while (true) {
-                size = in.read(buffer);
-                if (size < 0) {
-                    return;
-                }
-                out.write(buffer, /* offset= */ 0, size);
-            }
+            in.transferTo(out);
         }
     }
 
@@ -126,27 +150,15 @@
         }
     }
 
-    private static int getHeaderSize(int version) {
-        if (version == getVersion(1, 0)) {
-            return Integer.BYTES * 8; // Header has 8 integers.
-        }
-        return Integer.BYTES * 10; // Default + VM DTBO (offset, size)
-    }
-
-    private static boolean hasVmDtbo(int version) {
-        int major = getMajorVersion(version);
-        int minor = getMinorVersion(version);
-        return major > 1 || (major == 1 && minor >= 1);
+    private int getHeaderSize() {
+        // Header + (entry offset, entry, size) * mEntryCnt
+        return Integer.BYTES * (4 + mEntryCnt * 2);
     }
 
     private static int alignTo(int x, int size) {
         return (x + size - 1) & ~(size - 1);
     }
 
-    private static int getVersion(int major, int minor) {
-        return ((major & 0xFFFF) << 16) | (minor & 0xFFFF);
-    }
-
     private static int getMajorVersion(int version) {
         return (version >> 16) & 0xFFFF;
     }
@@ -160,6 +172,8 @@
         @NonNull private final File mPvmfwBinFile;
         @NonNull private final File mBccFile;
         @Nullable private File mDebugPolicyFile;
+        @Nullable private File mVmDtboFile;
+        @Nullable private File mVmReferenceDtFile;
         private int mVersion;
 
         public Builder(@NonNull File pvmfwBinFile, @NonNull File bccFile) {
@@ -175,14 +189,32 @@
         }
 
         @NonNull
+        public Builder setVmDtbo(@Nullable File vmDtboFile) {
+            mVmDtboFile = vmDtboFile;
+            return this;
+        }
+
+        @NonNull
+        public Builder setVmReferenceDt(@Nullable File vmReferenceDtFile) {
+            mVmReferenceDtFile = vmReferenceDtFile;
+            return this;
+        }
+
+        @NonNull
         public Builder setVersion(int major, int minor) {
-            mVersion = getVersion(major, minor);
+            mVersion = makeVersion(major, minor);
             return this;
         }
 
         @NonNull
         public Pvmfw build() {
-            return new Pvmfw(mPvmfwBinFile, mBccFile, mDebugPolicyFile, mVersion);
+            return new Pvmfw(
+                    mPvmfwBinFile,
+                    mBccFile,
+                    mDebugPolicyFile,
+                    mVmDtboFile,
+                    mVmReferenceDtFile,
+                    mVersion);
         }
     }
 }
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java b/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java
index d9d425a..a3216c2 100644
--- a/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java
+++ b/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java
@@ -19,6 +19,7 @@
 import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
 
 import static org.junit.Assume.assumeTrue;
 
@@ -26,16 +27,18 @@
 import androidx.annotation.Nullable;
 
 import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
+import com.android.microdroid.test.host.CommandRunner;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.TestDevice;
+import com.android.tradefed.util.CommandResult;
+import com.android.tradefed.util.CommandStatus;
 import com.android.tradefed.util.FileUtil;
 
 import org.junit.After;
 import org.junit.Before;
 
 import java.io.File;
-import java.util.Objects;
 import java.util.Map;
 
 /** Base class for testing custom pvmfw */
@@ -50,6 +53,9 @@
     @NonNull
     public static final String MICRODROID_CONFIG_PATH = "assets/microdroid/vm_config_apex.json";
 
+    @NonNull
+    public static final String VM_REFERENCE_DT_PATH = "/data/local/tmp/pvmfw/reference_dt.dtb";
+
     @NonNull public static final String MICRODROID_LOG_PATH = TEST_ROOT + "log.txt";
     public static final int BOOT_COMPLETE_TIMEOUT_MS = 30000; // 30 seconds
     public static final int BOOT_FAILURE_WAIT_TIME_MS = 10000; // 10 seconds
@@ -60,17 +66,28 @@
     @NonNull public static final String CUSTOM_PVMFW_IMG_PATH = TEST_ROOT + PVMFW_FILE_NAME;
     @NonNull public static final String CUSTOM_PVMFW_IMG_PATH_PROP = "hypervisor.pvmfw.path";
 
-    @Nullable private static File mPvmfwBinFileOnHost;
-    @Nullable private static File mBccFileOnHost;
+    @NonNull private static final String DUMPSYS = "dumpsys";
 
-    @Nullable private TestDevice mAndroidDevice;
+    @NonNull
+    private static final String DUMPSYS_MISSING_SERVICE_MSG_PREFIX = "Can't find service: ";
+
+    @NonNull
+    private static final String SECRET_KEEPER_AIDL =
+            "android.hardware.security.secretkeeper.ISecretkeeper/default";
+
+    @Nullable private File mPvmfwBinFileOnHost;
+    @Nullable private File mBccFileOnHost;
+    @Nullable private File mVmReferenceDtFile;
+    private boolean mSecretKeeperSupported;
+
+    @NonNull private TestDevice mAndroidDevice;
     @Nullable private ITestDevice mMicrodroidDevice;
 
     @Nullable private File mCustomPvmfwFileOnHost;
 
     @Before
     public void setUp() throws Exception {
-        mAndroidDevice = (TestDevice) Objects.requireNonNull(getDevice());
+        mAndroidDevice = (TestDevice) getDevice();
 
         // Check device capabilities
         assumeDeviceIsCapable(mAndroidDevice);
@@ -78,12 +95,29 @@
                 "Skip if protected VMs are not supported",
                 mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true));
 
-        // tradefed copies the test artifacts under /tmp when running tests,
-        // so we should *find* the artifacts with the file name.
-        mPvmfwBinFileOnHost =
-                getTestInformation().getDependencyFile(PVMFW_FILE_NAME, /* targetFirst= */ false);
-        mBccFileOnHost =
-                getTestInformation().getDependencyFile(BCC_FILE_NAME, /* targetFirst= */ false);
+        mPvmfwBinFileOnHost = findTestFile(PVMFW_FILE_NAME);
+        mBccFileOnHost = findTestFile(BCC_FILE_NAME);
+
+        // This is prepared by AndroidTest.xml
+        mVmReferenceDtFile = mAndroidDevice.pullFile(VM_REFERENCE_DT_PATH);
+
+        CommandRunner runner = new CommandRunner(mAndroidDevice);
+        CommandResult result = runner.runForResult(DUMPSYS, SECRET_KEEPER_AIDL);
+
+        // dumpsys prints 'Can't find service: ~' to stderr if secret keeper HAL is missing,
+        // but it doesn't return any error code for it.
+        // Read stderr to know whether secret keeper is supported, and stop test for any other case.
+        assertWithMessage("Failed to run " + DUMPSYS + ", result=" + result)
+                .that(result.getStatus() == CommandStatus.SUCCESS && result.getExitCode() == 0)
+                .isTrue();
+        if (result.getStderr() != null && !result.getStderr().trim().isEmpty()) {
+            assertWithMessage(
+                            "Unexpected stderr from " + DUMPSYS + ", stderr=" + result.getStderr())
+                    .that(result.getStderr().trim().startsWith(DUMPSYS_MISSING_SERVICE_MSG_PREFIX))
+                    .isTrue();
+        } else {
+            mSecretKeeperSupported = true;
+        }
 
         // Prepare for system properties for custom pvmfw.img.
         // File will be prepared later in individual test and then pushed to device
@@ -100,15 +134,12 @@
 
     @After
     public void shutdown() throws Exception {
-        if (!mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true)) {
-            return;
-        }
-        if (mMicrodroidDevice != null) {
-            mAndroidDevice.shutdownMicrodroid(mMicrodroidDevice);
-            mMicrodroidDevice = null;
-        }
+        shutdownMicrodroid();
+
         mAndroidDevice.uninstallPackage(PACKAGE_NAME);
 
+        FileUtil.deleteFile(mVmReferenceDtFile);
+
         // Cleanup for custom pvmfw.img
         setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, "");
         FileUtil.deleteFile(mCustomPvmfwFileOnHost);
@@ -116,16 +147,30 @@
         cleanUpVirtualizationTestSetup(mAndroidDevice);
     }
 
+    /** Returns android device */
+    @NonNull
+    public TestDevice getAndroidDevice() {
+        return mAndroidDevice;
+    }
+
     /** Returns pvmfw.bin file on host for building custom pvmfw with */
+    @NonNull
     public File getPvmfwBinFile() {
         return mPvmfwBinFileOnHost;
     }
 
     /** Returns BCC file on host for building custom pvmfw with */
+    @NonNull
     public File getBccFile() {
         return mBccFileOnHost;
     }
 
+    /** Returns VM reference DT, generated from DUT, on host for building custom pvmfw with. */
+    @Nullable
+    public File getVmReferenceDtFile() {
+        return mVmReferenceDtFile;
+    }
+
     /**
      * Returns a custom pvmfw file.
      *
@@ -133,11 +178,22 @@
      * calling {@link #launchProtectedVmAndWaitForBootCompleted}, so virtualization manager can read
      * the file path from sysprop and boot pVM with it.
      */
+    @NonNull
     public File getCustomPvmfwFile() {
         return mCustomPvmfwFileOnHost;
     }
 
     /**
+     * Returns whether a secretkeeper is supported.
+     *
+     * <p>If {@code true}, then VM reference DT must exist. (i.e. {@link #getVmReferenceDtFile} must
+     * exist {@code null}).
+     */
+    public boolean isSecretKeeperSupported() {
+        return mSecretKeeperSupported;
+    }
+
+    /**
      * Launches protected VM with custom pvmfw ({@link #getCustomPvmfwFile}) and wait for boot
      * completed. Throws exception when boot failed.
      */
@@ -162,4 +218,12 @@
         assertThat(mMicrodroidDevice.enableAdbRoot()).isTrue();
         return mMicrodroidDevice;
     }
+
+    /** Shuts down microdroid if it's running */
+    public void shutdownMicrodroid() throws Exception {
+        if (mMicrodroidDevice != null) {
+            mAndroidDevice.shutdownMicrodroid(mMicrodroidDevice);
+            mMicrodroidDevice = null;
+        }
+    }
 }
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java b/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
index 803405d..223f93f 100644
--- a/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
+++ b/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
@@ -192,10 +192,15 @@
                 getTestInformation()
                         .getDependencyFile(debugPolicyFileName, /* targetFirst= */ false);
 
-        Pvmfw pvmfw =
+        Pvmfw.Builder builder =
                 new Pvmfw.Builder(getPvmfwBinFile(), getBccFile())
-                        .setDebugPolicyOverlay(mCustomDebugPolicyFileOnHost)
-                        .build();
+                        .setDebugPolicyOverlay(mCustomDebugPolicyFileOnHost);
+        if (isSecretKeeperSupported()) {
+            builder.setVmReferenceDt(getVmReferenceDtFile());
+        } else {
+            builder.setVersion(1, 1);
+        }
+        Pvmfw pvmfw = builder.build();
         pvmfw.serialize(getCustomPvmfwFile());
     }
 
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
index b68316d..19334d6 100644
--- a/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
+++ b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
@@ -35,19 +35,31 @@
 @RunWith(DeviceJUnit4ClassRunner.class)
 public class PvmfwImgTest extends CustomPvmfwHostTestCaseBase {
     @Test
-    public void testConfigVersion1_0_boots() throws Exception {
-        Pvmfw pvmfw = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 0).build();
-        pvmfw.serialize(getCustomPvmfwFile());
+    public void testPvmfw_beforeVmReferenceDt_whenSecretKeeperExists() throws Exception {
+        // VM reference DT is added since version 1.2
+        List<int[]> earlyVersions = Arrays.asList(new int[] {1, 0}, new int[] {1, 1});
+        Pvmfw.Builder builder = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile());
 
-        launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
-    }
+        for (int[] pair : earlyVersions) {
+            int major = pair[0];
+            int minor = pair[1];
+            String version = "v" + major + "." + minor;
 
-    @Test
-    public void testConfigVersion1_1_boots() throws Exception {
-        Pvmfw pvmfw = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 1).build();
-        pvmfw.serialize(getCustomPvmfwFile());
+            // Pvmfw config before v1.2 can't have secret keeper key in VM reference DT.
+            Pvmfw pvmfw = builder.setVersion(major, minor).build();
+            pvmfw.serialize(getCustomPvmfwFile());
 
-        launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
+            if (isSecretKeeperSupported()) {
+                // If secret keeper is supported, we can't boot with early version
+                assertThrows(
+                        "pvmfw shouldn't boot without VM reference DT, version=" + version,
+                        DeviceRuntimeException.class,
+                        () -> launchProtectedVmAndWaitForBootCompleted(BOOT_FAILURE_WAIT_TIME_MS));
+            } else {
+                launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
+                shutdownMicrodroid();
+            }
+        }
     }
 
     @Test
@@ -65,13 +77,21 @@
                         new int[] {0xFFFF, 1},
                         new int[] {0xFFFF, 0xFFFF});
 
-        Pvmfw.Builder builder = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile());
+        Pvmfw.Builder builder =
+                new Pvmfw.Builder(getPvmfwBinFile(), getBccFile())
+                        .setVmReferenceDt(getVmReferenceDtFile());
 
         for (int[] pair : invalid_versions) {
             int major = pair[0];
             int minor = pair[1];
             String version = "v" + major + "." + minor;
 
+            if (Pvmfw.makeVersion(major, minor) >= Pvmfw.makeVersion(1, 2)
+                    && getVmReferenceDtFile() == null) {
+                // VM reference DT is unavailable, so we can't even build Pvmfw.
+                continue;
+            }
+
             Pvmfw pvmfw = builder.setVersion(major, minor).build();
             pvmfw.serialize(getCustomPvmfwFile());
 
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/PvmfwTest.java b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwTest.java
new file mode 100644
index 0000000..bea72eb
--- /dev/null
+++ b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwTest.java
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.pvmfw.test;
+
+import static org.junit.Assert.assertThrows;
+
+import com.android.pvmfw.test.host.Pvmfw;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/** Test test helper */
+@RunWith(DeviceJUnit4ClassRunner.class)
+public class PvmfwTest extends CustomPvmfwHostTestCaseBase {
+    @Test
+    public void testPvmfw_withConfig1_2_requiresReferenceDt() {
+        assertThrows(
+                "pvmfw config 1.2 must require VM reference DT",
+                NullPointerException.class,
+                () -> {
+                    new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 2).build();
+                });
+    }
+
+    @Test
+    public void testPvmfw_before1_2_doesNotRequiresReferenceDt() {
+        new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 1).build();
+    }
+}
diff --git a/tests/pvmfw/tools/PvmfwTool.java b/tests/pvmfw/tools/PvmfwTool.java
index 62c641b..e150ec4 100644
--- a/tests/pvmfw/tools/PvmfwTool.java
+++ b/tests/pvmfw/tools/PvmfwTool.java
@@ -25,28 +25,41 @@
 public class PvmfwTool {
     public static void printUsage() {
         System.out.println("pvmfw-tool: Appends pvmfw.bin and config payloads.");
-        System.out.println("Requires BCC and optional debug policy dtbo files");
-        System.out.println("");
-        System.out.println("Usage: pvmfw-tool <out> <pvmfw.bin> <bcc.dat> [<dp.dtbo>]");
+        System.out.println("            Requires BCC and VM reference DT.");
+        System.out.println("            VM DTBO and Debug policy can optionally be specified");
+        System.out.println(
+                "Usage: pvmfw-tool <out> <pvmfw.bin> <bcc.dat> <VM reference DT> [VM DTBO] [debug"
+                        + " policy]");
     }
 
     public static void main(String[] args) {
-        if (args.length != 4 && args.length != 3) {
+        if (args.length < 3 || args.length > 6) {
             printUsage();
             System.exit(1);
         }
 
         File out = new File(args[0]);
-        File pvmfw_bin = new File(args[1]);
-        File bcc_dat = new File(args[2]);
+        File pvmfwBin = new File(args[1]);
+        File bccData = new File(args[2]);
+        File vmReferenceDt = new File(args[3]);
+
+        File vmDtbo = null;
+        File dp = null;
+        if (args.length > 4) {
+            vmDtbo = new File(args[4]);
+        }
+        if (args.length > 5) {
+            dp = new File(args[5]);
+        }
 
         try {
-            Pvmfw.Builder builder = new Pvmfw.Builder(pvmfw_bin, bcc_dat);
-            if (args.length == 4) {
-                File dtbo = new File(args[3]);
-                builder.setDebugPolicyOverlay(dtbo);
-            }
-            builder.build().serialize(out);
+            Pvmfw pvmfw =
+                    new Pvmfw.Builder(pvmfwBin, bccData)
+                            .setVmReferenceDt(vmReferenceDt)
+                            .setDebugPolicyOverlay(dp)
+                            .setVmDtbo(vmDtbo)
+                            .build();
+            pvmfw.serialize(out);
         } catch (IOException e) {
             e.printStackTrace();
             printUsage();
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 22bea58..6be219e 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -49,7 +49,7 @@
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
         BnVirtualMachineService, IVirtualMachineService,
 };
-use android_hardware_security_secretkeeper::aidl::android::hardware::security::secretkeeper::ISecretkeeper::ISecretkeeper;
+use android_hardware_security_secretkeeper::aidl::android::hardware::security::secretkeeper::ISecretkeeper::{BnSecretkeeper, ISecretkeeper};
 use android_hardware_security_secretkeeper::aidl::android::hardware::security::secretkeeper::SecretId::SecretId;
 use android_hardware_security_authgraph::aidl::android::hardware::security::authgraph::{
     Arc::Arc as AuthgraphArc, IAuthGraphKeyExchange::IAuthGraphKeyExchange,
@@ -1501,11 +1501,12 @@
         }
     }
 
-    fn getSecretkeeper(&self) -> binder::Result<Option<Strong<dyn ISecretkeeper>>> {
-        // TODO(b/327526008): Session establishment wth secretkeeper is failing.
-        // Re-enable this when fixed.
-        let _sk_supported = is_secretkeeper_supported();
-        Ok(None)
+    fn getSecretkeeper(&self) -> binder::Result<Strong<dyn ISecretkeeper>> {
+        if !is_secretkeeper_supported() {
+            return Err(StatusCode::NAME_NOT_FOUND)?;
+        }
+        let sk = binder::wait_for_interface(SECRETKEEPER_IDENTIFIER)?;
+        Ok(BnSecretkeeper::new_binder(SecretkeeperProxy(sk), BinderFeatures::default()))
     }
 
     fn requestAttestation(&self, csr: &[u8], test_mode: bool) -> binder::Result<Vec<Certificate>> {
@@ -1514,8 +1515,11 @@
 }
 
 fn is_secretkeeper_supported() -> bool {
-    binder::is_declared(SECRETKEEPER_IDENTIFIER)
-        .expect("Could not check for declared Secretkeeper interface")
+    // TODO(b/327526008): Session establishment wth secretkeeper is failing.
+    // Re-enable this when fixed.
+    let _sk_supported = binder::is_declared(SECRETKEEPER_IDENTIFIER)
+        .expect("Could not check for declared Secretkeeper interface");
+    false
 }
 
 impl VirtualMachineService {
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index fc7fcd2..0c39501 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -77,6 +77,9 @@
         "virtualizationservice_defaults",
     ],
     test_suites: ["general-tests"],
+    rustlibs: [
+        "libtempfile",
+    ],
     data: [
         ":test_rkp_cert_chain",
     ],
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
index 6806a5c..662c8f1 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -58,9 +58,9 @@
     Certificate[] requestAttestation(in byte[] csr, in boolean testMode);
 
     /**
-     * Request connection to Secretkeeper. This is used by pVM to store Anti-Rollback protected
-     * secrets. Note that the return value is nullable to reflect that Secretkeeper HAL may not be
-     * present.
+     * Request connection to Secretkeeper. This is used by pVM to store rollback protected secrets.
+     * Note that this returns error if Secretkeeper is not supported on device. Guest should check
+     * that Secretkeeper is supported from Linux device tree before calling this.
      */
-    @nullable ISecretkeeper getSecretkeeper();
+    ISecretkeeper getSecretkeeper();
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index bbfb220..2fe14c0 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -39,7 +39,10 @@
 use openssl::x509::X509;
 use rand::Fill;
 use rkpd_client::get_rkpd_attestation_key;
-use rustutils::system_properties;
+use rustutils::{
+    system_properties,
+    users::{multiuser_get_app_id, multiuser_get_user_id},
+};
 use serde::Deserialize;
 use service_vm_comm::Response;
 use std::collections::{HashMap, HashSet};
@@ -385,7 +388,6 @@
         Ok(ParcelFileDescriptor::new(file))
     }
 
-    // TODO(b/294177871) Persist this Id, along with client uuid.
     fn allocateInstanceId(&self) -> binder::Result<[u8; 64]> {
         let mut id = [0u8; 64];
         id.try_fill(&mut rand::thread_rng())
@@ -393,6 +395,16 @@
             .or_service_specific_exception(-1)?;
         let uid = get_calling_uid();
         info!("Allocated a VM's instance_id: {:?}, for uid: {:?}", hex::encode(id), uid);
+        let state = &mut *self.state.lock().unwrap();
+        if let Some(sk_state) = &mut state.sk_state {
+            let user_id = multiuser_get_user_id(uid);
+            let app_id = multiuser_get_app_id(uid);
+            info!("Recording potential existence of state for (user_id={user_id}, app_id={app_id}");
+            if let Err(e) = sk_state.add_id(&id, user_id, app_id) {
+                error!("Failed to record the instance_id: {e:?}");
+            }
+        }
+
         Ok(id)
     }
 
diff --git a/virtualizationservice/src/maintenance.rs b/virtualizationservice/src/maintenance.rs
index 7fc2f37..0a367c5 100644
--- a/virtualizationservice/src/maintenance.rs
+++ b/virtualizationservice/src/maintenance.rs
@@ -15,7 +15,7 @@
 use android_hardware_security_secretkeeper::aidl::android::hardware::security::secretkeeper::{
     ISecretkeeper::ISecretkeeper, SecretId::SecretId,
 };
-use anyhow::Result;
+use anyhow::{Context, Result};
 use log::{error, info, warn};
 
 mod vmdb;
@@ -88,6 +88,13 @@
         }
     }
 
+    /// Record a new VM ID.
+    pub fn add_id(&mut self, vm_id: &VmId, user_id: u32, app_id: u32) -> Result<()> {
+        let user_id: i32 = user_id.try_into().context(format!("user_id {user_id} out of range"))?;
+        let app_id: i32 = app_id.try_into().context(format!("app_id {app_id} out of range"))?;
+        self.vm_id_db.add_vm_id(vm_id, user_id, app_id)
+    }
+
     /// Delete the VM IDs associated with Android user ID `user_id`.
     pub fn delete_ids_for_user(&mut self, user_id: i32) -> Result<()> {
         let vm_ids = self.vm_id_db.vm_ids_for_user(user_id)?;
diff --git a/virtualizationservice/src/maintenance/vmdb.rs b/virtualizationservice/src/maintenance/vmdb.rs
index bdff034..ce1e1e7 100644
--- a/virtualizationservice/src/maintenance/vmdb.rs
+++ b/virtualizationservice/src/maintenance/vmdb.rs
@@ -14,7 +14,7 @@
 
 //! Database of VM IDs.
 
-use anyhow::{Context, Result};
+use anyhow::{anyhow, Context, Result};
 use log::{debug, error, info, warn};
 use rusqlite::{params, params_from_iter, Connection, OpenFlags, Rows};
 use std::path::PathBuf;
@@ -29,6 +29,15 @@
 /// (Default value of `SQLITE_LIMIT_VARIABLE_NUMBER` for <= 3.32.0)
 const MAX_VARIABLES: usize = 999;
 
+/// Return the current time as milliseconds since epoch.
+fn db_now() -> u64 {
+    let now = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap_or(std::time::Duration::ZERO)
+        .as_millis();
+    now.try_into().unwrap_or(u64::MAX)
+}
+
 /// Identifier for a VM and its corresponding secret.
 pub type VmId = [u8; 64];
 
@@ -37,6 +46,8 @@
     conn: Connection,
 }
 
+struct RetryOnFailure(bool);
+
 impl VmIdDb {
     /// Connect to the VM ID database file held in the given directory, creating it if necessary.
     /// The second return value indicates whether a new database file was created.
@@ -49,8 +60,11 @@
             std::fs::create_dir(&db_path).context("failed to create {db_path:?}")?;
             info!("created persistent db dir {db_path:?}");
         }
-
         db_path.push(DB_FILENAME);
+        Self::new_at_path(db_path, RetryOnFailure(true))
+    }
+
+    fn new_at_path(db_path: PathBuf, retry: RetryOnFailure) -> Result<(Self, bool)> {
         let (flags, created) = if db_path.exists() {
             debug!("connecting to existing database {db_path:?}");
             (
@@ -69,15 +83,42 @@
                 true,
             )
         };
-        let mut result = Self {
-            conn: Connection::open_with_flags(db_path, flags)
+        let mut db = Self {
+            conn: Connection::open_with_flags(&db_path, flags)
                 .context(format!("failed to open/create DB with {flags:?}"))?,
         };
 
         if created {
-            result.init_tables().context("failed to create tables")?;
+            db.init_tables().context("failed to create tables")?;
+        } else {
+            // An existing .sqlite file may have an earlier schema.
+            match db.schema_version() {
+                Err(e) => {
+                    // Couldn't determine a schema version, so wipe and try again.
+                    error!("failed to determine VM DB schema: {e:?}");
+                    if retry.0 {
+                        // This is the first attempt, so wipe and retry.
+                        error!("resetting database file {db_path:?}");
+                        let _ = std::fs::remove_file(&db_path);
+                        return Self::new_at_path(db_path, RetryOnFailure(false));
+                    } else {
+                        // An earlier attempt at wiping/retrying has failed, so give up.
+                        return Err(anyhow!("failed to reset database file {db_path:?}"));
+                    }
+                }
+                Ok(0) => db.upgrade_tables_v0_v1().context("failed to upgrade schema v0 -> v1")?,
+                Ok(1) => {
+                    // Current version, no action needed.
+                }
+                Ok(version) => {
+                    // If the database looks like it's from a future version, leave it alone and
+                    // fail to connect to it.
+                    error!("database from the future (v{version})");
+                    return Err(anyhow!("database from the future (v{version})"));
+                }
+            }
         }
-        Ok((result, created))
+        Ok((db, created))
     }
 
     /// Delete the associated database file.
@@ -94,8 +135,63 @@
         }
     }
 
-    /// Create the database table and indices.
+    fn schema_version(&mut self) -> Result<i32> {
+        let version: i32 = self
+            .conn
+            .query_row("PRAGMA main.user_version", (), |row| row.get(0))
+            .context("failed to read pragma")?;
+        Ok(version)
+    }
+
+    /// Create the database table and indices using the current schema.
     fn init_tables(&mut self) -> Result<()> {
+        self.init_tables_v1()
+    }
+
+    /// Create the database table and indices using the v1 schema.
+    fn init_tables_v1(&mut self) -> Result<()> {
+        info!("creating v1 database schema");
+        self.conn
+            .execute(
+                "CREATE TABLE IF NOT EXISTS main.vmids (
+                     vm_id BLOB PRIMARY KEY,
+                     user_id INTEGER,
+                     app_id INTEGER,
+                     created INTEGER
+                 ) WITHOUT ROWID;",
+                (),
+            )
+            .context("failed to create table")?;
+        self.conn
+            .execute("CREATE INDEX IF NOT EXISTS main.vmids_user_index ON vmids(user_id);", [])
+            .context("Failed to create user index")?;
+        self.conn
+            .execute(
+                "CREATE INDEX IF NOT EXISTS main.vmids_app_index ON vmids(user_id, app_id);",
+                [],
+            )
+            .context("Failed to create app index")?;
+        self.conn
+            .execute("PRAGMA main.user_version = 1;", ())
+            .context("failed to declare version")?;
+        Ok(())
+    }
+
+    fn upgrade_tables_v0_v1(&mut self) -> Result<()> {
+        let _rows = self
+            .conn
+            .execute("ALTER TABLE main.vmids ADD COLUMN created INTEGER;", ())
+            .context("failed to alter table v0->v1")?;
+        self.conn
+            .execute("PRAGMA main.user_version = 1;", ())
+            .context("failed to set schema version")?;
+        Ok(())
+    }
+
+    /// Create the database table and indices using the v0 schema.
+    #[cfg(test)]
+    fn init_tables_v0(&mut self) -> Result<()> {
+        info!("creating v0 database schema");
         self.conn
             .execute(
                 "CREATE TABLE IF NOT EXISTS main.vmids (
@@ -119,13 +215,13 @@
     }
 
     /// Add the given VM ID into the database.
-    #[allow(dead_code)] // TODO(b/294177871): connect this up
     pub fn add_vm_id(&mut self, vm_id: &VmId, user_id: i32, app_id: i32) -> Result<()> {
+        let now = db_now();
         let _rows = self
             .conn
             .execute(
-                "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
-                params![vm_id, &user_id, &app_id],
+                "REPLACE INTO main.vmids (vm_id, user_id, app_id, created) VALUES (?1, ?2, ?3, ?4);",
+                params![vm_id, &user_id, &app_id, &now],
             )
             .context("failed to add VM ID")?;
         Ok(())
@@ -177,16 +273,21 @@
     }
 }
 
+/// Current schema version.
+#[cfg(test)]
+const SCHEMA_VERSION: usize = 1;
+
+/// Create a new in-memory database for testing.
 #[cfg(test)]
 pub fn new_test_db() -> VmIdDb {
-    let mut db = VmIdDb { conn: Connection::open_in_memory().unwrap() };
-    db.init_tables().unwrap();
-    db
+    tests::new_test_db_version(SCHEMA_VERSION)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::io::Write;
+
     const VM_ID1: VmId = [1u8; 64];
     const VM_ID2: VmId = [2u8; 64];
     const VM_ID3: VmId = [3u8; 64];
@@ -201,6 +302,113 @@
     const APP_C: i32 = 70;
     const APP_UNKNOWN: i32 = 99;
 
+    pub fn new_test_db_version(version: usize) -> VmIdDb {
+        let mut db = VmIdDb { conn: Connection::open_in_memory().unwrap() };
+        match version {
+            0 => db.init_tables_v0().unwrap(),
+            1 => db.init_tables_v1().unwrap(),
+            _ => panic!("unexpected version {version}"),
+        }
+        db
+    }
+
+    fn show_contents(db: &VmIdDb) {
+        let mut stmt = db.conn.prepare("SELECT * FROM main.vmids;").unwrap();
+        let mut rows = stmt.query(()).unwrap();
+        while let Some(row) = rows.next().unwrap() {
+            println!("  {row:?}");
+        }
+    }
+
+    #[test]
+    fn test_schema_version0() {
+        let mut db0 = VmIdDb { conn: Connection::open_in_memory().unwrap() };
+        db0.init_tables_v0().unwrap();
+        let version = db0.schema_version().unwrap();
+        assert_eq!(0, version);
+    }
+
+    #[test]
+    fn test_schema_version1() {
+        let mut db1 = VmIdDb { conn: Connection::open_in_memory().unwrap() };
+        db1.init_tables_v1().unwrap();
+        let version = db1.schema_version().unwrap();
+        assert_eq!(1, version);
+    }
+
+    #[test]
+    fn test_schema_upgrade_v0_v1() {
+        let mut db = new_test_db_version(0);
+        let version = db.schema_version().unwrap();
+        assert_eq!(0, version);
+
+        // Manually insert a row before upgrade.
+        db.conn
+            .execute(
+                "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
+                params![&VM_ID1, &USER1, APP_A],
+            )
+            .unwrap();
+
+        db.upgrade_tables_v0_v1().unwrap();
+        let version = db.schema_version().unwrap();
+        assert_eq!(1, version);
+
+        assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
+        show_contents(&db);
+    }
+
+    #[test]
+    fn test_corrupt_database_file() {
+        let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
+        let mut db_path = db_dir.path().to_owned();
+        db_path.push(DB_FILENAME);
+        {
+            let mut file = std::fs::File::create(db_path).unwrap();
+            let _ = file.write_all(b"This is not an SQLite file!");
+        }
+
+        // Non-DB file should be wiped and start over.
+        let (mut db, created) =
+            VmIdDb::new(&db_dir.path().to_string_lossy()).expect("failed to replace bogus DB");
+        assert!(created);
+        db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
+        assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
+    }
+
+    #[test]
+    fn test_non_upgradable_database_file() {
+        let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
+        let mut db_path = db_dir.path().to_owned();
+        db_path.push(DB_FILENAME);
+        {
+            // Create an unrelated database that happens to apparently have a schema version of 0.
+            let (db, created) = VmIdDb::new(&db_dir.path().to_string_lossy()).unwrap();
+            assert!(created);
+            db.conn.execute("DROP TABLE main.vmids", ()).unwrap();
+            db.conn.execute("PRAGMA main.user_version = 0;", ()).unwrap();
+        }
+
+        // Should fail to open a database because the upgrade fails.
+        let result = VmIdDb::new(&db_dir.path().to_string_lossy());
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_database_from_the_future() {
+        let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
+        {
+            let (mut db, created) = VmIdDb::new(&db_dir.path().to_string_lossy()).unwrap();
+            assert!(created);
+            db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
+            // Make the database look like it's from a future version.
+            db.conn.execute("PRAGMA main.user_version = 99;", ()).unwrap();
+        }
+        // Should fail to open a database from the future.
+        let result = VmIdDb::new(&db_dir.path().to_string_lossy());
+        assert!(result.is_err());
+    }
+
     #[test]
     fn test_add_remove() {
         let mut db = new_test_db();
@@ -239,6 +447,7 @@
         assert_eq!(vec![VM_ID5], db.vm_ids_for_user(USER3).unwrap());
         assert_eq!(empty, db.vm_ids_for_user(USER_UNKNOWN).unwrap());
         assert_eq!(empty, db.vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
+        show_contents(&db);
     }
 
     #[test]
@@ -254,12 +463,13 @@
         // Manually insert a row with a VM ID that's the wrong size.
         db.conn
             .execute(
-                "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
-                params![&[99u8; 60], &USER1, APP_A],
+                "REPLACE INTO main.vmids (vm_id, user_id, app_id, created) VALUES (?1, ?2, ?3, ?4);",
+                params![&[99u8; 60], &USER1, APP_A, &db_now()],
             )
             .unwrap();
 
         // Invalid row is skipped and remainder returned.
         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());
+        show_contents(&db);
     }
 }