libavb_rs: make lifetimes more flexible

Previously we required that the Ops object and the preloaded partition
data it wrapped had the same lifetime. This created a bit of awkwardness
because the preloaded data gets passed into the verification result,
which means the verification result also carried an (unnecessary) borrow
of the entire Ops object.

This CL splits the preloaded data into its own lifetime. This allows the
verification result to not care about the Ops object, and only borrow
the preloaded data, giving the caller much more flexibility.

Bug: b/330335215
Test: atest libavb_rs_test libavb_rs_uuid_test libavb_rs_unittest libavb_rs_uuid_unittest
Test: `mm` with the compile-fail test included fails to build as expected
Change-Id: If5a1f3f0ca7235532b511143e39abf4b551fc0f9

4 files changed, 236 insertions, 82 deletions

diff --git a/rust/src/ops.rs b/rust/src/ops.rs
index e030dda..4b67b0f 100644
--- a/rust/src/ops.rs
+++ b/rust/src/ops.rs
@@ -33,7 +33,19 @@ use uuid::Uuid;
 /// Base implementation-provided callbacks for verification.
 ///
 /// See libavb `AvbOps` for more complete documentation.
-pub trait Ops {
+///
+/// # Lifetimes
+/// The trait lifetime `'a` indicates the lifetime of any preloaded partition data.
+///
+/// Preloading partitions is an optional feature which allows libavb to use data already loaded to
+/// RAM rather than allocating memory itself and loading data from disk. Preloading changes the
+/// data ownership model so that the verification result borrows this existing data rather than
+/// allocating and owning the data itself. Because of this borrow, we need the lifetime here to
+/// ensure that the underlying data outlives the verification result object.
+///
+/// If `get_preloaded_partition()` is left unimplemented, all data is loaded and owned by the
+/// verification result rather than borrowed, and this trait lifetime can be `'static`.
+pub trait Ops<'a> {
     /// Reads data from the requested partition on disk.
     ///
     /// # Arguments
@@ -68,7 +80,7 @@ pub trait Ops {
     /// * `Err<IoError::NotImplemented>` if the requested partition has not been preloaded;
     ///   verification will next attempt to load the partition via `read_from_partition()`.
     /// * Any other `Err<IoError>` if an error occurred; verification will exit immediately.
-    fn get_preloaded_partition(&mut self, _partition: &CStr) -> IoResult<&[u8]> {
+    fn get_preloaded_partition(&mut self, _partition: &CStr) -> IoResult<&'a [u8]> {
         Err(IoError::NotImplemented)
     }
 
@@ -276,25 +288,29 @@ pub struct PublicKeyForPartitionInfo {
 /// perform `Ops` (Rust)
 /// callback
 /// ```
-pub(crate) struct UserData<'a>(&'a mut dyn Ops);
+///
+/// # Lifetimes
+/// * `'o`: lifetime of the `Ops` object
+/// * `'p`: lifetime of any preloaded data provided by `Ops`
+pub(crate) struct UserData<'o, 'p>(&'o mut dyn Ops<'p>);
 
-impl<'a> UserData<'a> {
-    pub(crate) fn new(ops: &'a mut dyn Ops) -> Self {
+impl<'o, 'p> UserData<'o, 'p> {
+    pub(crate) fn new(ops: &'o mut dyn Ops<'p>) -> Self {
         Self(ops)
     }
 }
 
 /// Wraps the C `AvbOps` struct with lifetime information for the compiler.
-pub(crate) struct ScopedAvbOps<'a> {
+pub(crate) struct ScopedAvbOps<'o, 'p> {
     /// `AvbOps` holds a raw pointer to `UserData` with no lifetime information.
     avb_ops: AvbOps,
     /// This provides the necessary lifetime information so the compiler can make sure that
     /// the `UserData` stays alive at least as long as we do.
-    _user_data: PhantomData<UserData<'a>>,
+    _user_data: PhantomData<UserData<'o, 'p>>,
 }
 
-impl<'a> ScopedAvbOps<'a> {
-    pub(crate) fn new(user_data: &'a mut UserData<'a>) -> Self {
+impl<'o, 'p> ScopedAvbOps<'o, 'p> {
+    pub(crate) fn new(user_data: &mut UserData<'o, 'p>) -> Self {
         Self {
             avb_ops: AvbOps {
                 // Rust won't transitively cast so we need to cast twice manually, but the compiler
@@ -320,7 +336,7 @@ impl<'a> ScopedAvbOps<'a> {
     }
 }
 
-impl<'a> AsMut<AvbOps> for ScopedAvbOps<'a> {
+impl<'o, 'p> AsMut<AvbOps> for ScopedAvbOps<'o, 'p> {
     fn as_mut(&mut self) -> &mut AvbOps {
         &mut self.avb_ops
     }
@@ -328,6 +344,9 @@ impl<'a> AsMut<AvbOps> for ScopedAvbOps<'a> {
 
 /// Extracts the user-provided `Ops` from a raw `AvbOps`.
 ///
+/// This function is used in libavb callbacks to bridge libavb's raw C `AvbOps` struct to our Rust
+/// implementation.
+///
 /// # Arguments
 /// * `avb_ops`: The raw `AvbOps` pointer used by libavb.
 ///
@@ -335,17 +354,33 @@ impl<'a> AsMut<AvbOps> for ScopedAvbOps<'a> {
 /// The Rust `Ops` extracted from `avb_ops.user_data`.
 ///
 /// # Safety
-/// Only call this function on an `AvbOps` created via `ScopedAvbOps`.
+/// * only call this function on an `AvbOps` created via `ScopedAvbOps`
+/// * drop all references to the returned `Ops` and preloaded data before returning control to
+///   libavb or calling this function again
+///
+/// In practice, these conditions are met since we call this at most once in each callback
+/// to extract the `Ops`, and drop the references at callback completion.
 ///
-/// Additionally, this should be considered a mutable borrow of the contained `Ops`:
-/// * do not return back to libavb while still holding the returned reference, or it will result
-///   in a dangling reference
-/// * do not call this again until the previous `Ops` goes out of scope, or it will violate Rust's
-///   mutable borrowing rules
+/// # Lifetimes
+/// * `'o`: lifetime of the `Ops` object
+/// * `'p`: lifetime of any preloaded data provided by `Ops`
 ///
-/// In practice, these conditions are met since we call this exactly once in each callback
-/// to extract the `Ops`, and drop it at callback completion.
-unsafe fn as_ops<'a>(avb_ops: *mut AvbOps) -> IoResult<&'a mut dyn Ops> {
+/// It's difficult to accurately provide the lifetimes when calling this function, since we are in
+/// a C callback which provides no lifetime information in the args. We solve this in the safety
+/// requirements by requiring the caller to drop both references before returning, which is always
+/// a subset of the actual object lifetimes as the objects must remain valid while libavb is
+/// actively using them:
+///
+/// ```ignore
+/// ops/preloaded lifetime {  // Actual 'o/'p start
+///   call into libavb {
+///     libavb callbacks {
+///       as_ops()            // as_ops() 'o/'p start
+///     }                     // as_ops() 'o/'p end
+///   }
+/// }                         // Actual 'o/'p end
+/// ```
+unsafe fn as_ops<'o, 'p>(avb_ops: *mut AvbOps) -> IoResult<&'o mut dyn Ops<'p>> {
     // SAFETY: we created this AvbOps object and passed it to libavb so we know it meets all
     // the criteria for `as_mut()`.
     let avb_ops = unsafe { avb_ops.as_mut() }.ok_or(IoError::Io)?;
diff --git a/rust/src/verify.rs b/rust/src/verify.rs
index e972226..0098644 100644
--- a/rust/src/verify.rs
+++ b/rust/src/verify.rs
@@ -34,6 +34,7 @@ use avb_bindgen::{
 use core::{
     ffi::{c_char, CStr},
     fmt,
+    marker::PhantomData,
     ptr::{self, null, null_mut, NonNull},
     slice,
 };
@@ -200,13 +201,25 @@ impl fmt::Debug for PartitionData {
 /// Wraps a raw C `AvbSlotVerifyData` struct.
 ///
 /// This provides a Rust safe view over the raw data; no copies are made.
+///
+/// # Lifetimes
+/// * `'a`: the lifetime of any preloaded partition data borrowed from an `Ops<'a>` object.
+///
+/// If the `Ops` doesn't provide any preloaded data, `SlotVerifyData` doesn't borrow anything
+/// and instead allocates and owns all data internally, freeing it accordingly on `Drop`. In this
+/// case, `'a` can be `'static` which imposes no lifetime restrictions on `SlotVerifyData`.
 pub struct SlotVerifyData<'a> {
     /// Internally owns the underlying data and deletes it on drop.
     raw_data: NonNull<AvbSlotVerifyData>,
 
-    /// This provides the necessary lifetime borrow so the compiler can make sure that the `Ops`
-    /// stays alive at least as long as we do, since it owns any preloaded partition data.
-    _ops: &'a dyn Ops,
+    /// This provides the necessary lifetimes so the compiler can make sure that the preloaded
+    /// partition data stays alive at least as long as we do, since the underlying
+    /// `AvbSlotVerifyData` may wrap this data rather than making a copy.
+    //
+    // We do not want to actually borrow an `Ops` here, since in some cases `Ops` is just a
+    // temporary object and may go out of scope before us. The only shared data is the preloaded
+    // partition contents, not the entire `Ops` object.
+    _preloaded: PhantomData<&'a [u8]>,
 }
 
 // Useful so that `SlotVerifyError`, which may hold a `SlotVerifyData`, can derive `PartialEq`.
@@ -225,24 +238,28 @@ impl<'a> SlotVerifyData<'a> {
     /// The returned `SlotVerifyData` will take ownership of the given `AvbSlotVerifyData` and
     /// properly release the allocated memory when it drops.
     ///
+    /// If `ops` provided any preloaded data, the returned `SlotVerifyData` also borrows the data to
+    /// account for the underlying `AvbSlotVerifyData` holding a pointer to it. If there was no
+    /// preloaded data, then `SlotVerifyData` owns all its data.
+    ///
     /// # Arguments
-    /// * `data`: a `AvbSlotVerifyData` object created by libavb.
-    /// * `ops`: the user-provided callback ops; borrowing this here ensures that any preloaded
-    ///          partition data stays unmodified while `data` is wrapping it.
+    /// * `data`: a `AvbSlotVerifyData` object created by libavb using `ops`.
+    /// * `ops`: the user-provided `Ops` object that was used for verification; only used here to
+    ///          grab the preloaded data lifetime.
     ///
     /// # Returns
     /// The new object, or `Err(SlotVerifyError::Internal)` if the data looks invalid.
     ///
     /// # Safety
-    /// * `data` must be a valid `AvbSlotVerifyData` object created by libavb
+    /// * `data` must be a valid `AvbSlotVerifyData` object created by libavb using `ops`.
     /// * after calling this function, do not access `data` except through the returned object
     unsafe fn new(
         data: *mut AvbSlotVerifyData,
-        ops: &'a mut dyn Ops,
+        _ops: &dyn Ops<'a>,
     ) -> SlotVerifyNoDataResult<Self> {
         let ret = Self {
             raw_data: NonNull::new(data).ok_or(SlotVerifyError::Internal)?,
-            _ops: ops,
+            _preloaded: PhantomData,
         };
 
         // Validate all the contained data here so accessors will never fail.
@@ -319,6 +336,7 @@ impl<'a> SlotVerifyData<'a> {
     }
 }
 
+/// Frees any internally-allocated and owned data.
 impl<'a> Drop for SlotVerifyData<'a> {
     fn drop(&mut self) {
         // SAFETY:
@@ -371,11 +389,11 @@ impl<'a> fmt::Debug for SlotVerifyData<'a> {
 /// 2. if `AllowVerificationError` is given in `flags`, it will also be returned on verification
 ///    failure
 ///
-/// If a `SlotVerifyData` is returned, it will borrow the provided `ops`. This is to ensure that
-/// any data shared by `SlotVerifyData` and `ops` - in particular preloaded partition contents -
-/// is not modified until `SlotVerifyData` is dropped.
+/// A returned `SlotVerifyData` will also borrow any preloaded data provided by `ops`. The `ops`
+/// object itself can go out of scope, but any preloaded data it could provide must outlive the
+/// returned object.
 pub fn slot_verify<'a>(
-    ops: &'a mut dyn Ops,
+    ops: &mut dyn Ops<'a>,
     requested_partitions: &[&CStr],
     ab_suffix: Option<&CStr>,
     flags: SlotVerifyFlags,
diff --git a/rust/tests/test_ops.rs b/rust/tests/test_ops.rs
index 1c62c4d..d752650 100644
--- a/rust/tests/test_ops.rs
+++ b/rust/tests/test_ops.rs
@@ -19,20 +19,53 @@ use std::{cmp::min, collections::HashMap, ffi::CStr};
 #[cfg(feature = "uuid")]
 use uuid::Uuid;
 
-/// Represents a single fake partition.
-#[derive(Default)]
-pub struct FakePartition {
-    /// Partition contents.
-    pub contents: Vec<u8>,
+/// Where the fake partition contents come from.
+pub enum PartitionContents<'a> {
+    /// Read on-demand from disk.
+    FromDisk(Vec<u8>),
+    /// Preloaded and passed in.
+    Preloaded(&'a [u8]),
+}
+
+impl<'a> PartitionContents<'a> {
+    /// Returns the partition data.
+    pub fn as_slice(&self) -> &[u8] {
+        match self {
+            Self::FromDisk(v) => v,
+            Self::Preloaded(c) => c,
+        }
+    }
+
+    /// Returns a mutable reference to the `FromDisk` data for test modification. Panicks if the
+    /// data is actually `Preloaded` instead.
+    pub fn as_mut_vec(&mut self) -> &mut Vec<u8> {
+        match self {
+            Self::FromDisk(v) => v,
+            Self::Preloaded(_) => panic!("Cannot mutate preloaded partition data"),
+        }
+    }
+}
 
-    /// Whether the partition should report as preloaded or not.
-    pub preloaded: bool,
+/// Represents a single fake partition.
+pub struct FakePartition<'a> {
+    /// Partition contents, either preloaded or read on-demand.
+    pub contents: PartitionContents<'a>,
 
     /// Partition UUID.
     #[cfg(feature = "uuid")]
     pub uuid: Uuid,
 }
 
+impl<'a> FakePartition<'a> {
+    fn new(contents: PartitionContents<'a>) -> Self {
+        Self {
+            contents,
+            #[cfg(feature = "uuid")]
+            uuid: Default::default(),
+        }
+    }
+}
+
 /// Fake vbmeta key state.
 pub struct FakeVbmetaKeyState {
     /// Key trust & rollback index info.
@@ -47,9 +80,9 @@ pub struct FakeVbmetaKeyState {
 ///
 /// The user is expected to set up the internal values to the desired device state - disk contents,
 /// rollback indices, etc. This class then uses this state to implement the avb callback operations.
-pub struct TestOps {
+pub struct TestOps<'a> {
     /// Partitions to provide to libavb callbacks.
-    pub partitions: HashMap<&'static str, FakePartition>,
+    pub partitions: HashMap<&'static str, FakePartition<'a>>,
 
     /// Vbmeta public keys as a map of {(key, metadata): state}. Querying unknown keys will
     /// return `IoError::Io`.
@@ -69,27 +102,40 @@ pub struct TestOps {
     pub persistent_values: HashMap<String, IoResult<Vec<u8>>>,
 }
 
-impl TestOps {
-    /// Adds a partition with the given contents.
+impl<'a> TestOps<'a> {
+    /// Adds a fake on-disk partition with the given contents.
     ///
     /// Reduces boilerplate a bit by taking in a raw array and returning a &mut so tests can
     /// do something like this:
     ///
     /// ```
     /// test_ops.add_partition("foo", [1, 2, 3, 4]);
-    /// test_ops.add_partition("bar", [0, 0]).preloaded = true;
+    /// test_ops.add_partition("bar", [0, 0]).uuid = uuid!(...);
     /// ```
     pub fn add_partition<T: Into<Vec<u8>>>(
         &mut self,
         name: &'static str,
         contents: T,
-    ) -> &mut FakePartition {
+    ) -> &mut FakePartition<'a> {
         self.partitions.insert(
             name,
-            FakePartition {
-                contents: contents.into(),
-                ..Default::default()
-            },
+            FakePartition::new(PartitionContents::FromDisk(contents.into())),
+        );
+        self.partitions.get_mut(name).unwrap()
+    }
+
+    /// Adds a preloaded partition with the given contents.
+    ///
+    /// Same a `add_partition()` except that the preloaded data is not owned by
+    /// the `TestOps` but passed in, which means it can outlive `TestOps`.
+    pub fn add_preloaded_partition(
+        &mut self,
+        name: &'static str,
+        contents: &'a [u8],
+    ) -> &mut FakePartition<'a> {
+        self.partitions.insert(
+            name,
+            FakePartition::new(PartitionContents::Preloaded(contents)),
         );
         self.partitions.get_mut(name).unwrap()
     }
@@ -145,7 +191,7 @@ impl TestOps {
     }
 }
 
-impl Default for TestOps {
+impl Default for TestOps<'_> {
     fn default() -> Self {
         Self {
             partitions: HashMap::new(),
@@ -157,23 +203,19 @@ impl Default for TestOps {
     }
 }
 
-impl Ops for TestOps {
+impl<'a> Ops<'a> for TestOps<'a> {
     fn read_from_partition(
         &mut self,
         partition: &CStr,
         offset: i64,
         buffer: &mut [u8],
     ) -> IoResult<usize> {
-        let partition = self
+        let contents = self
             .partitions
             .get(partition.to_str()?)
-            .ok_or(IoError::NoSuchPartition)?;
-
-        // We should never be trying to read a preloaded partition from disk since we already
-        // have it available in memory.
-        assert!(!partition.preloaded);
-
-        let contents = &partition.contents;
+            .ok_or(IoError::NoSuchPartition)?
+            .contents
+            .as_slice();
 
         // Negative offset means count backwards from the end.
         let offset = {
@@ -202,13 +244,12 @@ impl Ops for TestOps {
         Ok(bytes_read)
     }
 
-    fn get_preloaded_partition(&mut self, partition: &CStr) -> IoResult<&[u8]> {
+    fn get_preloaded_partition(&mut self, partition: &CStr) -> IoResult<&'a [u8]> {
         match self.partitions.get(partition.to_str()?) {
             Some(FakePartition {
-                contents,
-                preloaded: true,
+                contents: PartitionContents::Preloaded(preloaded),
                 ..
-            }) => Ok(&contents[..]),
+            }) => Ok(&preloaded[..]),
             _ => Err(IoError::NotImplemented),
         }
     }
@@ -251,7 +292,7 @@ impl Ops for TestOps {
     fn get_size_of_partition(&mut self, partition: &CStr) -> IoResult<u64> {
         self.partitions
             .get(partition.to_str()?)
-            .map(|p| u64::try_from(p.contents.len()).unwrap())
+            .map(|p| u64::try_from(p.contents.as_slice().len()).unwrap())
             .ok_or(IoError::NoSuchPartition)
     }
 
diff --git a/rust/tests/verify_tests.rs b/rust/tests/verify_tests.rs
index 015bc9f..2e4aaad 100644
--- a/rust/tests/verify_tests.rs
+++ b/rust/tests/verify_tests.rs
@@ -66,7 +66,7 @@ const TEST_HASHTREE_DIGEST_HEX: &str = "5373fc4ee3dd898325eeeffb5a1dbb041900c5f1
 const TEST_HASHTREE_ALGORITHM: &str = "sha1";
 
 /// Initializes a `TestOps` object such that verification will succeed on `TEST_PARTITION_NAME`.
-fn test_ops_one_image_one_vbmeta() -> TestOps {
+fn test_ops_one_image_one_vbmeta<'a>() -> TestOps<'a> {
     let mut ops = TestOps::default();
     ops.add_partition(TEST_PARTITION_NAME, fs::read(TEST_IMAGE_PATH).unwrap());
     ops.add_partition("vbmeta", fs::read(TEST_VBMETA_PATH).unwrap());
@@ -77,7 +77,9 @@ fn test_ops_one_image_one_vbmeta() -> TestOps {
 }
 
 /// Calls `slot_verify()` using standard args for `test_ops_one_image_one_vbmeta()` setup.
-fn verify_one_image_one_vbmeta(ops: &mut TestOps) -> SlotVerifyResult<SlotVerifyData> {
+fn verify_one_image_one_vbmeta<'a>(
+    ops: &mut TestOps<'a>,
+) -> SlotVerifyResult<'a, SlotVerifyData<'a>> {
     slot_verify(
         ops,
         &[&CString::new(TEST_PARTITION_NAME).unwrap()],
@@ -89,7 +91,7 @@ fn verify_one_image_one_vbmeta(ops: &mut TestOps) -> SlotVerifyResult<SlotVerify
 
 /// Initializes a `TestOps` object such that verification will succeed on `TEST_PARTITION_NAME` and
 /// `TEST_PARTITION_2_NAME`.
-fn test_ops_two_images_one_vbmeta() -> TestOps {
+fn test_ops_two_images_one_vbmeta<'a>() -> TestOps<'a> {
     let mut ops = test_ops_one_image_one_vbmeta();
     // Add in the contents of the second partition and overwrite the vbmeta partition to
     // include both partition descriptors.
@@ -99,7 +101,7 @@ fn test_ops_two_images_one_vbmeta() -> TestOps {
 }
 
 /// Calls `slot_verify()` for both test partitions.
-fn verify_two_images(ops: &mut TestOps) -> SlotVerifyResult<SlotVerifyData> {
+fn verify_two_images<'a>(ops: &mut TestOps<'a>) -> SlotVerifyResult<'a, SlotVerifyData<'a>> {
     slot_verify(
         ops,
         &[
@@ -114,7 +116,7 @@ fn verify_two_images(ops: &mut TestOps) -> SlotVerifyResult<SlotVerifyData> {
 
 /// Initializes a `TestOps` object such that verification will succeed on the `boot` partition with
 /// a combined image + vbmeta.
-fn test_ops_boot_partition() -> TestOps {
+fn test_ops_boot_partition<'a>() -> TestOps<'a> {
     let mut ops = test_ops_one_image_one_vbmeta();
     ops.partitions.clear();
     ops.add_partition(
@@ -125,7 +127,7 @@ fn test_ops_boot_partition() -> TestOps {
 }
 
 /// Calls `slot_verify()` using standard args for `test_ops_boot_partition()` setup.
-fn verify_boot_partition(ops: &mut TestOps) -> SlotVerifyResult<SlotVerifyData> {
+fn verify_boot_partition<'a>(ops: &mut TestOps<'a>) -> SlotVerifyResult<'a, SlotVerifyData<'a>> {
     slot_verify(
         ops,
         &[&CString::new("boot").unwrap()],
@@ -140,7 +142,7 @@ fn verify_boot_partition(ops: &mut TestOps) -> SlotVerifyResult<SlotVerifyData>
 
 /// Initializes a `TestOps` object such that verification will succeed on
 /// `TEST_PARTITION_PERSISTENT_DIGEST_NAME`.
-fn test_ops_persistent_digest(image: Vec<u8>) -> TestOps {
+fn test_ops_persistent_digest<'a>(image: Vec<u8>) -> TestOps<'a> {
     let mut ops = test_ops_one_image_one_vbmeta();
     ops.partitions.clear();
     // Use the vbmeta image with the persistent digest descriptor.
@@ -154,7 +156,7 @@ fn test_ops_persistent_digest(image: Vec<u8>) -> TestOps {
 }
 
 /// Calls `slot_verify()` using standard args for `test_ops_persistent_digest()` setup.
-fn verify_persistent_digest(ops: &mut TestOps) -> SlotVerifyResult<SlotVerifyData> {
+fn verify_persistent_digest<'a>(ops: &mut TestOps<'a>) -> SlotVerifyResult<'a, SlotVerifyData<'a>> {
     slot_verify(
         ops,
         &[&CString::new(TEST_PARTITION_PERSISTENT_DIGEST_NAME).unwrap()],
@@ -166,7 +168,11 @@ fn verify_persistent_digest(ops: &mut TestOps) -> SlotVerifyResult<SlotVerifyDat
 
 /// Modifies the partition contents by flipping a bit.
 fn modify_partition_contents(ops: &mut TestOps, partition: &str) {
-    ops.partitions.get_mut(partition).unwrap().contents[0] ^= 0x01;
+    ops.partitions
+        .get_mut(partition)
+        .unwrap()
+        .contents
+        .as_mut_vec()[0] ^= 0x01;
 }
 
 /// Returns the persistent value name for `TEST_PARTITION_PERSISTENT_DIGEST_NAME`.
@@ -220,11 +226,9 @@ fn one_image_one_vbmeta_passes_verification_with_correct_data() {
 #[test]
 fn preloaded_image_passes_verification() {
     let mut ops = test_ops_one_image_one_vbmeta();
-    // Mark the image partition to be preloaded.
-    ops.partitions
-        .get_mut(TEST_PARTITION_NAME)
-        .unwrap()
-        .preloaded = true;
+    // Use preloaded data instead for the test partition.
+    let preloaded = fs::read(TEST_IMAGE_PATH).unwrap();
+    ops.add_preloaded_partition(TEST_PARTITION_NAME, &preloaded);
 
     let result = verify_one_image_one_vbmeta(&mut ops);
 
@@ -233,6 +237,63 @@ fn preloaded_image_passes_verification() {
     assert!(partition_data.preloaded());
 }
 
+// When all images are loaded from disk (rather than preloaded), libavb allocates memory itself for
+// the data, so there is no shared ownership; the returned verification data owns the image data
+// and can hold onto it even after the `ops` goes away.
+#[test]
+fn verification_data_from_disk_can_outlive_ops() {
+    let result = {
+        let mut ops = test_ops_one_image_one_vbmeta();
+        verify_one_image_one_vbmeta(&mut ops)
+    };
+
+    let data = result.unwrap();
+
+    // The verification data owns the images and we can still access them.
+    assert_eq!(
+        data.partition_data()[0].data(),
+        fs::read(TEST_IMAGE_PATH).unwrap()
+    );
+}
+
+// When preloaded data is passed into ops but outlives it, we can also continue to access it from
+// the verification data after the ops goes away. The ops was only borrowing it, and now the
+// verification data continues to borrow it.
+#[test]
+fn verification_data_preloaded_can_outlive_ops() {
+    let preloaded = fs::read(TEST_IMAGE_PATH).unwrap();
+
+    let result = {
+        let mut ops = test_ops_one_image_one_vbmeta();
+        ops.add_preloaded_partition(TEST_PARTITION_NAME, &preloaded);
+        verify_one_image_one_vbmeta(&mut ops)
+    };
+
+    let data = result.unwrap();
+
+    // The verification data is borrowing the preloaded images and we can still access them.
+    assert_eq!(data.partition_data()[0].data(), preloaded);
+}
+
+// When preloaded data is passed into ops but also goes out of scope, the verification data loses
+// access to it, violating lifetime rules.
+//
+// Our lifetimes *must* be configured such that this does not compile, since `result` is borrowing
+// `preloaded` which has gone out of scope.
+//
+// TODO: figure out how to make a compile-fail test; for now we just have to manually test by
+// un-commenting the code.
+// #[test]
+// fn verification_data_preloaded_cannot_outlive_result() {
+//     let result = {
+//         let preloaded = fs::read(TEST_IMAGE_PATH).unwrap();
+//         let mut ops = test_ops_one_image_one_vbmeta();
+//         ops.add_preloaded_partition(TEST_PARTITION_NAME, &preloaded);
+//         verify_one_image_one_vbmeta(&mut ops)
+//     };
+//     result.unwrap();
+// }
+
 #[test]
 fn slotted_partition_passes_verification() {
     let mut ops = test_ops_one_image_one_vbmeta();
@@ -427,6 +488,7 @@ fn undersized_partition_fails_verification() {
         .get_mut(TEST_PARTITION_NAME)
         .unwrap()
         .contents
+        .as_mut_vec()
         .pop();
 
     let result = verify_one_image_one_vbmeta(&mut ops);
@@ -578,10 +640,8 @@ fn one_image_one_vbmeta_verification_data_display() {
 #[test]
 fn preloaded_image_verification_data_display() {
     let mut ops = test_ops_one_image_one_vbmeta();
-    ops.partitions
-        .get_mut(TEST_PARTITION_NAME)
-        .unwrap()
-        .preloaded = true;
+    let preloaded = fs::read(TEST_IMAGE_PATH).unwrap();
+    ops.add_preloaded_partition(TEST_PARTITION_NAME, &preloaded);
 
     let result = verify_one_image_one_vbmeta(&mut ops);