examples: secure_storage-rs: avoid persistent reads into shared memory

crates/optee-utee/src/object/persistent_object.rs

@@ -15,10 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use alloc::vec::Vec;
 use optee_utee_sys as raw;
 
 use super::{DataFlag, GenericObject, ObjectHandle, ObjectStorageConstants, Whence};
-use crate::{Error, Result};
+use crate::{Error, ErrorKind, Result};
 
 /// An object identified by an Object Identifier and including a Data Stream.
 ///
@@ -329,7 +330,9 @@ impl PersistentObject {
     ///
     /// # Parameters
     ///
-    /// 1) `buffer`: A pre-allocated buffer for saving the object's data stream.
+    /// 1) `buffer`: A pre-allocated TA-private buffer for saving the object's
+    ///    data stream. Client-shared memory must not be passed here; use
+    ///    [`Self::read_remaining_to_vec`] and copy the result to the client instead.
     /// 2) `count`: The returned value contains the number of bytes read.
     /// # Example
     ///
@@ -382,6 +385,32 @@ impl PersistentObject {
         }
     }
 
+    /// Reads the remaining data from the current cursor into a new TA-private vector.
+    ///
+    /// This method does not rewind the cursor. To read the entire object, first
+    /// call `seek(0, Whence::DataSeekSet)`.
+    ///
+    /// The allocation is sized using the object's current metadata. The returned
+    /// vector contains only the bytes actually read, and the data position advances
+    /// accordingly. At or beyond the end of the object, returns an empty vector.
+    /// Concurrent growth after the metadata query is not included in this read.
+    ///
+    /// Returns errors from [`GenericObject::info`] or [`Self::read`], or
+    /// [`ErrorKind::OutOfMemory`] if the allocation fails.
+    pub fn read_remaining_to_vec(&mut self) -> Result<Vec<u8>> {
+        let info = self.info()?;
+        let size = info.data_size().saturating_sub(info.raw.dataPosition);
+        let mut data = Vec::new();
+        if size != 0 {
+            data.try_reserve_exact(size)
+                .map_err(|_| ErrorKind::OutOfMemory)?;
+            data.resize(size, 0);
+            let count = self.read(&mut data)?;
+            data.truncate(count as usize);
+        }
+        Ok(data)
+    }
+
     /// Write the passed in buffer data into from the data stream associate with
     /// the object.
     ///
@@ -549,6 +578,87 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn test_read_remaining_to_vec() {
+        let _lock = SERIAL_TEST_LOCK.lock().expect("should get the lock");
+        // Full read, nonzero position, short read, empty object, EOF, and beyond EOF.
+        for (size, position, bytes) in [
+            (4, 0, &b"data"[..]),
+            (4, 2, &b"ta"[..]),
+            (4, 0, &b"da"[..]),
+            (0, 0, &b""[..]),
+            (4, 4, &b""[..]),
+            (4, 6, &b""[..]),
+        ] {
+            let mut raw_handle = MockHandle::new();
+            let handle = raw_handle.as_handle();
+            let info = mock_api::TEE_GetObjectInfo1_context();
+            let read = mock_api::TEE_ReadObjectData_context();
+            let close = mock_api::TEE_CloseObject_context();
+            info.expect().times(1).return_once_st(move |_, info| {
+                unsafe {
+                    (*info).dataSize = size;
+                    (*info).dataPosition = position;
+                }
+                raw::TEE_SUCCESS
+            });
+            if position < size {
+                read.expect()
+                    .times(1)
+                    .return_once_st(move |_, buf, len, count| {
+                        assert_eq!(len, size - position);
+                        unsafe {
+                            core::ptr::copy_nonoverlapping(bytes.as_ptr(), buf.cast(), bytes.len());
+                            *count = bytes.len();
+                        }
+                        raw::TEE_SUCCESS
+                    });
+            } else {
+                read.expect().never();
+            }
+            close.expect().times(1).return_once_st(|_| ());
+            let mut object = PersistentObject(ObjectHandle::from_raw(handle).unwrap());
+            assert_eq!(object.read_remaining_to_vec().unwrap(), bytes);
+        }
+    }
+
+    #[test]
+    fn test_read_remaining_to_vec_errors() {
+        let _lock = SERIAL_TEST_LOCK.lock().expect("should get the lock");
+        for (info_error, size, expected) in [
+            (
+                raw::TEE_ERROR_STORAGE_NOT_AVAILABLE,
+                4,
+                raw::TEE_ERROR_STORAGE_NOT_AVAILABLE,
+            ),
+            (raw::TEE_SUCCESS, 4, raw::TEE_ERROR_CORRUPT_OBJECT),
+            (raw::TEE_SUCCESS, usize::MAX, raw::TEE_ERROR_OUT_OF_MEMORY),
+        ] {
+            let mut raw_handle = MockHandle::new();
+            let handle = raw_handle.as_handle();
+            let info = mock_api::TEE_GetObjectInfo1_context();
+            let read = mock_api::TEE_ReadObjectData_context();
+            let close = mock_api::TEE_CloseObject_context();
+            info.expect().times(1).return_once_st(move |_, info| {
+                unsafe { (*info).dataSize = size };
+                info_error
+            });
+            if info_error == raw::TEE_SUCCESS && size == 4 {
+                read.expect()
+                    .times(1)
+                    .return_once_st(|_, _, _, _| raw::TEE_ERROR_CORRUPT_OBJECT);
+            } else {
+                read.expect().never();
+            }
+            close.expect().times(1).return_once_st(|_| ());
+            let mut object = PersistentObject(ObjectHandle::from_raw(handle).unwrap());
+            assert_eq!(
+                object.read_remaining_to_vec().unwrap_err().raw_code(),
+                expected
+            );
+        }
+    }
+
     #[test]
     // If a persistent object is successfully created, TEE_CloseObject will be
     // called when it is dropped.

examples/ta/secure_storage-rs/src/main.rs

@@ -115,13 +115,17 @@ pub fn read_raw_object((p0, p1, _, _): &mut ParametersAny<'_>) -> Result<()> {
         DataFlag::ACCESS_READ | DataFlag::SHARE_READ,
     )?;
     let obj_info = object.info()?;
+    let data_size = obj_info.data_size();
+    if p1.get_capacity() < data_size {
+        return Err(ErrorKind::ShortBuffer.into());
+    }
 
-    let read_bytes = object.read(unsafe { p1.get_buffer_mut() })?;
-    if read_bytes != obj_info.data_size() as u32 {
+    let data = object.read_remaining_to_vec()?;
+    if data.len() != data_size {
         return Err(ErrorKind::ExcessData.into());
     }
 
-    p1.set_updated_size(read_bytes as usize)?;
+    p1.set_output(&data)?;
 
     Ok(())
 }