This is an automated email from the ASF dual-hosted git repository.

ivila pushed a commit to branch main
in repository 
https://gitbox.apache.org/repos/asf/incubator-teaclave-trustzone-sdk.git

commit d1d406541efce361cdeeeac801dea12ab57e1f52
Author: ivila <[email protected]>
AuthorDate: Wed Jul 9 09:55:00 2025 +0800

    optee-utee: object: refactor Object
    
    * move Drop(called with `TEE_CloseObject`) to `ObjectHandle` and
      remove Drop in both `PersistentObject` and `TransientObject`.
    * rename `ObjHandle` trait to `GenericObject` and move generic
      functions into it.
    * PersistentObject: get rid of calling `core::mem::forget` after
      making a successful called to `close_and_delete`.
    
    Signed-off-by: Zehui Chen <[email protected]>
    Reviewed-by: Yuan Zhuang <[email protected]>
    Reviewed-by: Zhaofeng Chen <[email protected]>
---
 optee-utee/src/crypto_op.rs                |  31 ++--
 optee-utee/src/object/enum_handle.rs       |   3 +-
 optee-utee/src/object/generic_object.rs    | 174 +++++++++++++++++++
 optee-utee/src/object/mod.rs               |   4 +-
 optee-utee/src/object/object_handle.rs     | 158 +++++++++++------
 optee-utee/src/object/persistent_object.rs | 270 +++++++++++++++++------------
 optee-utee/src/object/transient_object.rs  | 242 +++++++-------------------
 7 files changed, 522 insertions(+), 360 deletions(-)

diff --git a/optee-utee/src/crypto_op.rs b/optee-utee/src/crypto_op.rs
index f5b8f3f..17cafad 100644
--- a/optee-utee/src/crypto_op.rs
+++ b/optee-utee/src/crypto_op.rs
@@ -15,13 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::{Attribute, Error, ObjHandle, Result, TransientObject};
-use optee_utee_sys as raw;
+use alloc::{boxed::Box, vec::Vec};
 use core::{mem, ptr};
-#[cfg(not(target_os = "optee"))]
-use alloc::boxed::Box;
-#[cfg(not(target_os = "optee"))]
-use alloc::vec::Vec;
+
+use optee_utee_sys as raw;
+
+use crate::{Attribute, Error, GenericObject, Result, TransientObject};
 
 /// Specify one of the available cryptographic operations.
 #[repr(u32)]
@@ -216,7 +215,7 @@ impl OperationHandle {
         }
     }
 
-    fn set_key<T: ObjHandle>(&self, object: &T) -> Result<()> {
+    fn set_key<T: GenericObject>(&self, object: &T) -> Result<()> {
         match unsafe { raw::TEE_SetOperationKey(self.handle(), 
object.handle()) } {
             raw::TEE_SUCCESS => return Ok(()),
             code => Err(Error::from_raw_error(code)),
@@ -713,7 +712,7 @@ impl Cipher {
     /// 6) If operation is not in initial state.
     /// 7) Hardware or cryptographic algorithm failure.
     /// 8) If the Implementation detects any other error.
-    pub fn set_key<T: ObjHandle>(&self, object: &T) -> Result<()> {
+    pub fn set_key<T: GenericObject>(&self, object: &T) -> Result<()> {
         self.0.set_key(object)
     }
 
@@ -743,7 +742,7 @@ impl Cipher {
     /// 6) If operation is not in initial state.
     /// 7) Hardware or cryptographic algorithm failure.
     /// 8) If the Implementation detects any other error.
-    pub fn set_key_2<T: ObjHandle, D: ObjHandle>(&self, object1: &T, object2: 
&D) -> Result<()> {
+    pub fn set_key_2<T: GenericObject, D: GenericObject>(&self, object1: &T, 
object2: &D) -> Result<()> {
         match unsafe {
             raw::TEE_SetOperationKey2(self.handle(), object1.handle(), 
object2.handle())
         } {
@@ -948,7 +947,7 @@ impl Mac {
     }
 
     /// Function usage is similar to [Cipher::set_key](Cipher::set_key).
-    pub fn set_key<T: ObjHandle>(&self, object: &T) -> Result<()> {
+    pub fn set_key<T: GenericObject>(&self, object: &T) -> Result<()> {
         self.0.set_key(object)
     }
 
@@ -1031,9 +1030,7 @@ impl AE {
     /// 5) Hardware or cryptographic algorithm failure.
     /// 6) If the Implementation detects any other error.
     pub fn update_aad(&self, aad_data: &[u8]) {
-        unsafe {
-            raw::TEE_AEUpdateAAD(self.handle(), aad_data.as_ptr() as _, 
aad_data.len())
-        };
+        unsafe { raw::TEE_AEUpdateAAD(self.handle(), aad_data.as_ptr() as _, 
aad_data.len()) };
     }
 
     /// Accumulate data for an Authentication Encryption operation.
@@ -1237,7 +1234,7 @@ impl AE {
     }
 
     /// Function usage is similar to [Cipher::set_key](Cipher::set_key).
-    pub fn set_key<T: ObjHandle>(&self, object: &T) -> Result<()> {
+    pub fn set_key<T: GenericObject>(&self, object: &T) -> Result<()> {
         self.0.set_key(object)
     }
 
@@ -1495,7 +1492,7 @@ impl Asymmetric {
     }
 
     /// Function usage is similar to [Cipher::set_key](Cipher::set_key).
-    pub fn set_key<T: ObjHandle>(&self, object: &T) -> Result<()> {
+    pub fn set_key<T: GenericObject>(&self, object: &T) -> Result<()> {
         self.0.set_key(object)
     }
 
@@ -1529,7 +1526,7 @@ impl DeriveKey {
     ///
     /// ``` rust,no_run
     /// # use optee_utee::{AttributeMemref, AttributeId, TransientObject, 
TransientObjectType};
-    /// # use optee_utee::{DeriveKey, AlgorithmId};
+    /// # use optee_utee::{DeriveKey, AlgorithmId, GenericObject};
     /// # fn example1() -> optee_utee::Result<()> {
     ///
     /// let attr_prime = AttributeMemref::from_ref(AttributeId::DhPrime, 
&[23u8]);
@@ -1620,7 +1617,7 @@ impl DeriveKey {
     }
 
     /// Function usage is similar to [Cipher::set_key](Cipher::set_key).
-    pub fn set_key<T: ObjHandle>(&self, object: &T) -> Result<()> {
+    pub fn set_key<T: GenericObject>(&self, object: &T) -> Result<()> {
         self.0.set_key(object)
     }
 
diff --git a/optee-utee/src/object/enum_handle.rs 
b/optee-utee/src/object/enum_handle.rs
index 17d0d81..160a052 100644
--- a/optee-utee/src/object/enum_handle.rs
+++ b/optee-utee/src/object/enum_handle.rs
@@ -16,7 +16,6 @@
 // under the License.
 
 use alloc::boxed::Box;
-use core::ptr;
 
 use optee_utee_sys as raw;
 
@@ -34,7 +33,7 @@ impl ObjectEnumHandle {
     /// Allocate an object enumerator.
     /// Once an object enumerator has been allocated, it can be reused for 
multiple enumerations.
     pub fn allocate() -> Result<Self> {
-        let raw_handle: *mut raw::TEE_ObjectEnumHandle = 
Box::into_raw(Box::new(ptr::null_mut()));
+        let raw_handle: *mut raw::TEE_ObjectEnumHandle = 
Box::into_raw(Box::new(core::ptr::null_mut()));
         match unsafe { raw::TEE_AllocatePersistentObjectEnumerator(raw_handle) 
} {
             raw::TEE_SUCCESS => Ok(Self { raw: raw_handle }),
             code => {
diff --git a/optee-utee/src/object/generic_object.rs 
b/optee-utee/src/object/generic_object.rs
new file mode 100644
index 0000000..5d9bfa5
--- /dev/null
+++ b/optee-utee/src/object/generic_object.rs
@@ -0,0 +1,174 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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.
+
+use core::mem;
+
+use super::{AttributeId, ObjectInfo, UsageFlag};
+use crate::{Error, Result};
+
+use optee_utee_sys as raw;
+
+/// A generic trait for an object (transient or persistent).
+pub trait GenericObject {
+    /// Return the handle of an object.
+    fn handle(&self) -> raw::TEE_ObjectHandle;
+
+    /// Return the characteristics of an object.
+    ///
+    /// # Errors
+    ///
+    /// For [PersistentObject](crate::PersistentObject):
+    ///
+    /// * `CorruptObject`: If the persistent object is corrupt. The object
+    ///    handle SHALL behave based on the
+    ///    `gpd.ta.doesNotCloseHandleOnCorruptObject` property.
+    /// * `StorageNotAvailable`: If the persistent object is stored in a
+    ///    storage area which is currently inaccessible.
+    ///
+    /// # Panics
+    ///
+    /// * If object is not a valid opened object handle.
+    /// * If the implementation detects any other error associated with this
+    ///   function that is not explicitly associated with a defined return code
+    ///   for this function.
+    fn info(&self) -> Result<ObjectInfo> {
+        let mut raw_info: raw::TEE_ObjectInfo = unsafe { mem::zeroed() };
+        match unsafe { raw::TEE_GetObjectInfo1(self.handle(), &mut raw_info) } 
{
+            raw::TEE_SUCCESS => Ok(ObjectInfo::from_raw(raw_info)),
+            code => Err(Error::from_raw_error(code)),
+        }
+    }
+
+    /// Restrict the object usage flags of an object handle to contain at most
+    /// the flags passed in the obj_usage parameter.
+    ///
+    /// # Errors
+    ///
+    /// For [PersistentObject](crate::PersistentObject):
+    ///
+    /// * `CorruptObject`: If the persistent object is corrupt. The object
+    ///    handle SHALL behave based on the
+    ///    `gpd.ta.doesNotCloseHandleOnCorruptObject` property.
+    /// * `StorageNotAvailable`: If the persistent object is stored in a
+    ///    storage area which is currently inaccessible.
+    ///
+    /// # Panics
+    ///
+    /// * If object is not a valid opened object handle.
+    /// * If the implementation detects any other error associated with this
+    ///   function that is not explicitly associated with a defined return code
+    ///   for this function.
+    fn restrict_usage(&mut self, obj_usage: UsageFlag) -> Result<()> {
+        match unsafe { raw::TEE_RestrictObjectUsage1(self.handle(), 
obj_usage.bits()) } {
+            raw::TEE_SUCCESS => Ok(()),
+            code => Err(Error::from_raw_error(code)),
+        }
+    }
+
+    /// Extract one buffer attribute from an object. The attribute is
+    /// identified by the argument id.
+    ///
+    /// # Errors
+    ///
+    /// For all:
+    ///
+    /// * `ItemNotFound`: If the attribute is not found on this object.
+    /// * `SHORT_BUFFER`: If buffer is NULL or too small to contain the key
+    ///    part.
+    ///
+    /// For [PersistentObject](crate::PersistentObject):
+    ///
+    /// * `CorruptObject`: If the persistent object is corrupt. The object
+    ///    handle SHALL behave based on the
+    ///    `gpd.ta.doesNotCloseHandleOnCorruptObject` property.
+    /// * `StorageNotAvailable`: If the persistent object is stored in a
+    ///    storage area which is currently inaccessible.
+    ///
+    /// # Panics
+    ///
+    /// * If object is not a valid opened object handle.
+    /// * If the object is not initialized.
+    /// * If Bit [29] of attributeID is not set to 0, so the attribute is not a
+    ///   buffer attribute.
+    /// * If Bit [28] of attributeID is set to 0, denoting a protected
+    ///   attribute, and the object usage does not contain the
+    ///   TEE_USAGE_EXTRACTABLE flag.
+    /// * If the implementation detects any other error associated with this
+    ///   function that is not explicitly associated with a defined return code
+    ///   for this function.
+    fn ref_attribute(&self, id: AttributeId, buffer: &mut [u8]) -> 
Result<usize> {
+        let mut size = buffer.len();
+        match unsafe {
+            raw::TEE_GetObjectBufferAttribute(
+                self.handle(),
+                id as u32,
+                buffer as *mut _ as _,
+                &mut size,
+            )
+        } {
+            raw::TEE_SUCCESS => Ok(size),
+            code => Err(Error::from_raw_error(code)),
+        }
+    }
+
+    /// Extract one value attribute from an object. The attribute is identified
+    /// by the argument id.
+    ///
+    /// # Errors
+    ///
+    /// For all:
+    ///
+    /// * `ItemNotFound`: If the attribute is not found on this object.
+    /// * `SHORT_BUFFER`: If buffer is NULL or too small to contain the key
+    ///    part.
+    ///
+    /// For [PersistentObject](crate::PersistentObject):
+    ///
+    /// * `CorruptObject`: If the persistent object is corrupt. The object
+    ///    handle SHALL behave based on the
+    ///    `gpd.ta.doesNotCloseHandleOnCorruptObject` property.
+    /// * `StorageNotAvailable`: If the persistent object is stored in a
+    ///    storage area which is currently inaccessible.
+    ///
+    /// # Panics
+    ///
+    /// * If object is not a valid opened object handle.
+    /// * If the object is not initialized.
+    /// * If Bit [29] of attributeID is not set to 0, so the attribute is not a
+    ///   buffer attribute.
+    /// * If Bit [28] of attributeID is set to 0, denoting a protected
+    ///   attribute, and the object usage does not contain the
+    ///   TEE_USAGE_EXTRACTABLE flag.
+    /// * If the implementation detects any other error associated with this
+    ///   function that is not explicitly associated with a defined return code
+    ///   for this function.
+    fn value_attribute(&self, id: u32) -> Result<(u32, u32)> {
+        let mut value_a: u32 = 0;
+        let mut value_b: u32 = 0;
+        match unsafe {
+            raw::TEE_GetObjectValueAttribute(
+                self.handle(),
+                id,
+                &mut value_a as *mut _,
+                &mut value_b as *mut _,
+            )
+        } {
+            raw::TEE_SUCCESS => Ok((value_a, value_b)),
+            code => Err(Error::from_raw_error(code)),
+        }
+    }
+}
diff --git a/optee-utee/src/object/mod.rs b/optee-utee/src/object/mod.rs
index 9c98922..875fb0d 100644
--- a/optee-utee/src/object/mod.rs
+++ b/optee-utee/src/object/mod.rs
@@ -17,6 +17,7 @@
 
 mod attribute;
 mod enum_handle;
+mod generic_object;
 mod object_define;
 mod object_handle;
 mod object_info;
@@ -25,8 +26,9 @@ mod transient_object;
 
 pub use attribute::*;
 pub use enum_handle::ObjectEnumHandle;
+pub use generic_object::GenericObject;
 pub use object_define::*;
-pub use object_handle::{ObjHandle, ObjectHandle};
+pub use object_handle::ObjectHandle;
 pub use object_info::ObjectInfo;
 pub use persistent_object::PersistentObject;
 pub use transient_object::{TransientObject, TransientObjectType};
diff --git a/optee-utee/src/object/object_handle.rs 
b/optee-utee/src/object/object_handle.rs
index edb3dc2..1666e24 100644
--- a/optee-utee/src/object/object_handle.rs
+++ b/optee-utee/src/object/object_handle.rs
@@ -15,80 +15,130 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use core::mem;
-
 use optee_utee_sys as raw;
 
-use super::{AttributeId, ObjectInfo, UsageFlag};
-use crate::{Error, Result};
+use super::GenericObject;
 
 /// An opaque handle on an object.
-pub struct ObjectHandle {
-    pub(crate) raw: *mut raw::TEE_ObjectHandle,
-}
+#[derive(Debug)]
+pub struct ObjectHandle(raw::TEE_ObjectHandle);
 
 impl ObjectHandle {
-    pub fn handle(&self) -> raw::TEE_ObjectHandle {
-        unsafe { *(self.raw) }
+    pub fn from_raw(raw: raw::TEE_ObjectHandle) -> crate::Result<ObjectHandle> 
{
+        if raw.is_null() {
+            return Err(crate::ErrorKind::BadParameters.into());
+        }
+        Ok(Self(raw))
     }
 
-    pub fn is_null(&self) -> bool {
-        self.raw.is_null()
+    pub fn handle(&self) -> raw::TEE_ObjectHandle {
+        self.0
     }
 
-    pub fn from_raw(raw: *mut raw::TEE_ObjectHandle) -> ObjectHandle {
-        Self { raw }
+    /// Forget the inner handle to prevent a double-free, this function would 
be
+    /// called when the inner handle is(or will be) freed externally.
+    ///
+    /// Example:
+    /// ``` rust,no_run
+    /// # use optee_utee::ObjectHandle;
+    /// # use optee_utee_sys as raw;
+    /// # let external_handle: raw::TEE_ObjectHandle = core::ptr::null_mut();
+    /// # fn main() -> optee_utee::Result<()> {
+    /// # let external_handle = core::ptr::null_mut();
+    /// // `external_handle` is a handle that is constructed and controlled
+    /// // externally.
+    /// // `handle` is valid, and will call TEE_CloseObject on
+    /// // `external_handle` when it is dropping, which is not allowed
+    /// // as the `external_handle` is externally controlled.
+    /// let mut handle = ObjectHandle::from_raw(external_handle)?;
+    /// // ... Some operation
+    /// // forget the inner handle, so it won't call TEE_CloseObject on
+    /// // `external_handle`
+    /// handle.forget();
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn forget(mut self) {
+        self.0 = core::ptr::null_mut();
     }
+}
 
-    pub fn info(&self) -> Result<ObjectInfo> {
-        let mut raw_info: raw::TEE_ObjectInfo = unsafe { mem::zeroed() };
-        match unsafe { raw::TEE_GetObjectInfo1(self.handle(), &mut raw_info) } 
{
-            raw::TEE_SUCCESS => Ok(ObjectInfo::from_raw(raw_info)),
-            code => Err(Error::from_raw_error(code)),
-        }
+// functions for internal usage
+impl ObjectHandle {
+    pub(crate) fn new_null() -> Self {
+        Self(core::ptr::null_mut())
     }
 
-    pub fn restrict_usage(&mut self, obj_usage: UsageFlag) -> Result<()> {
-        match unsafe { raw::TEE_RestrictObjectUsage1(self.handle(), 
obj_usage.bits()) } {
-            raw::TEE_SUCCESS => Ok(()),
-            code => Err(Error::from_raw_error(code)),
-        }
+    pub(crate) fn is_null(&self) -> bool {
+        self.0.is_null()
     }
+}
 
-    pub fn ref_attribute(&self, id: AttributeId, buffer: &mut [u8]) -> 
Result<usize> {
-        let mut size = buffer.len();
-        match unsafe {
-            raw::TEE_GetObjectBufferAttribute(
-                self.handle(),
-                id as u32,
-                buffer as *mut _ as _,
-                &mut size,
-            )
-        } {
-            raw::TEE_SUCCESS => Ok(size),
-            code => Err(Error::from_raw_error(code)),
+impl Drop for ObjectHandle {
+    fn drop(&mut self) {
+        if !self.is_null() {
+            unsafe { raw::TEE_CloseObject(self.handle()) }
         }
     }
+}
 
-    pub fn value_attribute(&self, id: u32) -> Result<(u32, u32)> {
-        let mut value_a: u32 = 0;
-        let mut value_b: u32 = 0;
-        match unsafe {
-            raw::TEE_GetObjectValueAttribute(
-                self.handle(),
-                id,
-                &mut value_a as *mut _,
-                &mut value_b as *mut _,
-            )
-        } {
-            raw::TEE_SUCCESS => Ok((value_a, value_b)),
-            code => Err(Error::from_raw_error(code)),
-        }
+impl GenericObject for ObjectHandle {
+    fn handle(&self) -> raw::TEE_ObjectHandle {
+        self.handle()
     }
 }
 
-/// A trait for an object (transient or persistent) to return its handle.
-pub trait ObjHandle {
-    /// Return the handle of an object.
-    fn handle(&self) -> raw::TEE_ObjectHandle;
+#[cfg(test)]
+mod tests {
+    extern crate std;
+
+    use optee_utee_mock::object::{set_global_object_mock, 
MockObjectController, SERIAL_TEST_LOCK};
+
+    use super::*;
+
+    /// Ensures `ObjectHandle` can be safely constructed from a raw handle
+    /// and automatically calls `TEE_CloseObject` when dropped.
+    #[test]
+    fn test_from_raw() {
+        let _lock = SERIAL_TEST_LOCK.lock();
+
+        let mut mock = MockObjectController::new();
+        let mut handle_struct = 
MockObjectController::new_valid_test_handle_struct();
+        let handle = MockObjectController::new_valid_test_handle(&mut 
handle_struct);
+
+        mock.expect_TEE_CloseObject_once(handle.clone());
+        set_global_object_mock(mock);
+
+        let obj = ObjectHandle::from_raw(unsafe { *handle.get() }).expect("it 
should be ok");
+        assert_eq!(obj.handle(), unsafe { *handle.get() });
+    }
+
+    /// Ensures `ObjectHandle` can call `forget` to prevent automatically
+    /// calls `TEE_CloseObject` when dropped.
+    #[test]
+    fn test_forget() {
+        let _lock = SERIAL_TEST_LOCK.lock();
+
+        let mut handle_struct = 
MockObjectController::new_valid_test_handle_struct();
+        let handle = MockObjectController::new_valid_test_handle(&mut 
handle_struct);
+        // making sure nothing would be called(includes TEE_CloseObject)
+        let mock = MockObjectController::new();
+        set_global_object_mock(mock);
+
+        let obj = ObjectHandle::from_raw(unsafe { *handle.get() }).expect("it 
should be ok");
+        assert_eq!(obj.handle(), unsafe { *handle.get() });
+
+        obj.forget();
+    }
+
+    #[test]
+    fn test_new_null() {
+        let _lock = SERIAL_TEST_LOCK.lock();
+        // making sure nothing would be called(includes TEE_CloseObject)
+        let mock = MockObjectController::new();
+        set_global_object_mock(mock);
+
+        let obj = ObjectHandle::new_null();
+        assert!(obj.is_null());
+    }
 }
diff --git a/optee-utee/src/object/persistent_object.rs 
b/optee-utee/src/object/persistent_object.rs
index dd6cfde..f34c650 100644
--- a/optee-utee/src/object/persistent_object.rs
+++ b/optee-utee/src/object/persistent_object.rs
@@ -15,20 +15,15 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use alloc::boxed::Box;
-use core::ptr;
-
 use optee_utee_sys as raw;
 
-use super::{
-    AttributeId, DataFlag, ObjHandle, ObjectHandle, ObjectInfo, 
ObjectStorageConstants, UsageFlag,
-    Whence,
-};
+use super::{DataFlag, GenericObject, ObjectHandle, ObjectStorageConstants, 
Whence};
 use crate::{Error, Result};
 
 /// An object identified by an Object Identifier and including a Data Stream.
 ///
 /// Contrast [TransientObject](TransientObject).
+#[derive(Debug)]
 pub struct PersistentObject(ObjectHandle);
 
 impl PersistentObject {
@@ -82,26 +77,21 @@ impl PersistentObject {
         object_id: &[u8],
         flags: DataFlag,
     ) -> Result<Self> {
-        let raw_handle: *mut raw::TEE_ObjectHandle = 
Box::into_raw(Box::new(ptr::null_mut()));
+        let mut handle: raw::TEE_ObjectHandle = core::ptr::null_mut();
+        // Move as much code as possible out of unsafe blocks to maximize 
Rust’s
+        // safety checks.
+        let handle_mut = &mut handle;
         match unsafe {
             raw::TEE_OpenPersistentObject(
                 storage_id as u32,
                 object_id.as_ptr() as _,
                 object_id.len(),
                 flags.bits(),
-                raw_handle as *mut _,
+                handle_mut,
             )
         } {
-            raw::TEE_SUCCESS => {
-                let handle = ObjectHandle::from_raw(raw_handle);
-                Ok(Self(handle))
-            }
-            code => {
-                unsafe {
-                    drop(Box::from_raw(raw_handle));
-                }
-                Err(Error::from_raw_error(code))
-            }
+            raw::TEE_SUCCESS => Ok(Self(ObjectHandle::from_raw(handle)?)),
+            code => Err(Error::from_raw_error(code)),
         }
     }
 
@@ -166,10 +156,13 @@ impl PersistentObject {
         attributes: Option<ObjectHandle>,
         initial_data: &[u8],
     ) -> Result<Self> {
-        let raw_handle: *mut raw::TEE_ObjectHandle = 
Box::into_raw(Box::new(ptr::null_mut()));
+        let mut handle: raw::TEE_ObjectHandle = core::ptr::null_mut();
+        // Move as much code as possible out of unsafe blocks to maximize 
Rust’s
+        // safety checks.
+        let handle_mut = &mut handle;
         let attributes = match attributes {
             Some(a) => a.handle(),
-            None => ptr::null_mut(),
+            None => core::ptr::null_mut(),
         };
         match unsafe {
             raw::TEE_CreatePersistentObject(
@@ -180,19 +173,11 @@ impl PersistentObject {
                 attributes,
                 initial_data.as_ptr() as _,
                 initial_data.len(),
-                raw_handle as *mut _,
+                handle_mut,
             )
         } {
-            raw::TEE_SUCCESS => {
-                let handle = ObjectHandle::from_raw(raw_handle);
-                Ok(Self(handle))
-            }
-            code => {
-                unsafe {
-                    drop(Box::from_raw(raw_handle));
-                }
-                Err(Error::from_raw_error(code))
-            }
+            raw::TEE_SUCCESS => Ok(Self(ObjectHandle::from_raw(handle)?)),
+            code => Err(Error::from_raw_error(code)),
         }
     }
 
@@ -211,7 +196,6 @@ impl PersistentObject {
     ///     Ok(mut object) =>
     ///     {
     ///         object.close_and_delete()?;
-    ///         std::mem::forget(object);
     ///         Ok(())
     ///     }
     ///     Err(e) => Err(e),
@@ -221,26 +205,47 @@ impl PersistentObject {
     ///
     /// # Errors
     ///
-    /// 1) `StorageNotAvailable`: If the [PersistentObject](PersistentObject) 
is stored in a storage area which is
-    ///    currently inaccessible.
+    /// 1) `StorageNotAvailable`: If the [PersistentObject](PersistentObject) 
is stored in a storage
+    ///    area which is currently inaccessible.
     ///
     /// # Panics
     ///
     /// 1) If object is not a valid opened object.
     /// 2) If the Implementation detects any other error associated with this 
function which is not
     ///    explicitly associated with a defined return code for this function.
-    // this function is conflicted with Drop implementation, when use this one 
to avoid panic:
-    // Call `mem::forget` for this structure to avoid double drop the object
-    pub fn close_and_delete(&mut self) -> Result<()> {
-        match unsafe { 
raw::TEE_CloseAndDeletePersistentObject1(self.0.handle()) } {
-            raw::TEE_SUCCESS => {
-                unsafe {
-                    drop(Box::from_raw(self.0.raw));
-                }
-                return Ok(());
-            }
+    ///
+    /// # Breaking Changes
+    ///
+    /// Now we no longer need to call `core::mem::forget` after successfully 
calling
+    /// `close_and_delete`, and code like this will now produce a compilation 
error.
+    /// ``` rust,compile_fail
+    /// # use optee_utee::{PersistentObject, ObjectStorageConstants, DataFlag};
+    /// # fn main() -> optee_utee::Result<()> {
+    /// # let obj_id = [0_u8];
+    /// let mut obj = PersistentObject::open (
+    ///         ObjectStorageConstants::Private,
+    ///         &obj_id,
+    ///         DataFlag::ACCESS_READ,
+    /// )?;
+    /// obj.close_and_delete()?;
+    /// core::mem::forget(obj); // will get compilation error in this line
+    /// //                ^^^ value used here after move
+    /// # Ok(())
+    /// # }
+    pub fn close_and_delete(self) -> Result<()> {
+        let result = match unsafe { 
raw::TEE_CloseAndDeletePersistentObject1(self.0.handle()) } {
+            raw::TEE_SUCCESS => Ok(()),
             code => Err(Error::from_raw_error(code)),
-        }
+        };
+        // According to `GPD_TEE_Internal_Core_API_Specification_v1.3.1`:
+        // At 5.7.4 TEE_CloseAndDeletePersistentObject1:
+        // Deleting an object is atomic; once this function returns, the object
+        // is definitely deleted and no more open handles for the object exist.
+        //
+        // So we must forget the raw_handle to prevent calling TEE_CloseObject
+        // on it (no matter the result of TEE_CloseAndDeletePersistentObject1).
+        self.0.forget();
+        return result;
     }
 
     /// Changes the identifier of an object.
@@ -294,53 +299,6 @@ impl PersistentObject {
             code => Err(Error::from_raw_error(code)),
         }
     }
-    /// Return the characteristics of an object.
-    /// Function is similar to [TransientObject::info](TransientObject::info) 
besides extra errors.
-    ///
-    /// # Errors
-    ///
-    /// 1) `CorruptObject`: If the [PersistentObject](PersistentObject) is 
corrupt. The object handle is closed.
-    /// 2) `StorageNotAvailable`: If the [PersistentObject](PersistentObject) 
is stored in a storage area which is
-    ///    currently inaccessible.
-    pub fn info(&self) -> Result<ObjectInfo> {
-        self.0.info()
-    }
-
-    /// Restrict the object usage flags of an object handle to contain at most 
the flags passed in the obj_usage parameter.
-    /// Function is similar to 
[TransientObject::restrict_usage](TransientObject::restrict_usage) besides 
extra errors.
-    ///
-    /// # Errors
-    ///
-    /// 1) `CorruptObject`: If the [PersistentObject](PersistentObject) is 
corrupt. The object handle is closed.
-    /// 2) `StorageNotAvailable`: If the [PersistentObject](PersistentObject) 
is stored in a storage area which is
-    ///    currently inaccessible.
-    pub fn restrict_usage(&mut self, obj_usage: UsageFlag) -> Result<()> {
-        self.0.restrict_usage(obj_usage)
-    }
-
-    /// Extract one buffer attribute from an object. The attribute is 
identified by the argument id.
-    /// Function is similar to 
[TransientObject::ref_attribute](TransientObject::ref_attribute) besides extra 
errors.
-    ///
-    /// # Errors
-    ///
-    /// 1) `CorruptObject`: If the [PersistentObject](PersistentObject) is 
corrupt. The object handle is closed.
-    /// 2) `StorageNotAvailable`: If the [PersistentObject](PersistentObject) 
is stored in a storage area which is
-    ///    currently inaccessible.
-    pub fn ref_attribute(&self, id: AttributeId, buffer: &mut [u8]) -> 
Result<usize> {
-        self.0.ref_attribute(id, buffer)
-    }
-
-    /// Extract one value attribute from an object. The attribute is 
identified by the argument id.
-    /// Function is similar to 
[TransientObject::value_attribute](TransientObject::value_attribute) besides 
extra errors.
-    ///
-    /// # Errors
-    ///
-    /// 1) `CorruptObject`: If the [PersistentObject](PersistentObject) is 
corrupt. The object handle is closed.
-    /// 2) `StorageNotAvailable`: If the [PersistentObject](PersistentObject) 
is stored in a storage area which is
-    ///    currently inaccessible.
-    pub fn value_attribute(&self, id: u32) -> Result<(u32, u32)> {
-        self.0.value_attribute(id)
-    }
 
     /// Read requested size from the data stream associate with the object 
into the buffer.
     ///
@@ -529,26 +487,120 @@ impl PersistentObject {
     }
 }
 
-impl ObjHandle for PersistentObject {
+impl GenericObject for PersistentObject {
     fn handle(&self) -> raw::TEE_ObjectHandle {
         self.0.handle()
     }
 }
 
-impl Drop for PersistentObject {
-    /// Close an opened [PersistentObject](PersistentObject).
-    ///
-    /// # Panics
-    ///
-    /// 1) If object is not a valid opened object.
-    /// 2) If the Implementation detects any other error associated with this 
function which is not
-    ///    explicitly associated with a defined return code for this function.
-    fn drop(&mut self) {
-        unsafe {
-            if self.0.raw != Box::into_raw(Box::new(ptr::null_mut())) {
-                raw::TEE_CloseObject(self.0.handle());
-            }
-            drop(Box::from_raw(self.0.raw));
-        }
+#[cfg(test)]
+mod tests {
+    use optee_utee_mock::{
+        object::{set_global_object_mock, MockObjectController, 
SERIAL_TEST_LOCK},
+        raw,
+    };
+
+    use super::*;
+
+    #[test]
+    // If a persistent object is successfully created, TEE_CloseObject will be
+    // called when it is dropped.
+    fn test_create_and_drop() {
+        let _lock = SERIAL_TEST_LOCK.lock();
+
+        let mut mock = MockObjectController::new();
+        let mut handle_struct = 
MockObjectController::new_valid_test_handle_struct();
+        let handle = MockObjectController::new_valid_test_handle(&mut 
handle_struct);
+
+        mock.expect_TEE_CreatePersistentObject_success_once(handle.clone());
+        mock.expect_TEE_CloseObject_once(handle);
+
+        set_global_object_mock(mock);
+
+        let _obj = PersistentObject::create(
+            ObjectStorageConstants::Private,
+            &[],
+            DataFlag::ACCESS_WRITE,
+            None,
+            &[],
+        )
+        .expect("it should be ok");
+    }
+
+    #[test]
+    fn test_create_failed() {
+        let _lock = SERIAL_TEST_LOCK.lock();
+
+        static RETURN_CODE: raw::TEE_Result = raw::TEE_ERROR_BAD_STATE;
+
+        let mut mock = MockObjectController::new();
+        mock.expect_TEE_CreatePersistentObject_fail_once(RETURN_CODE);
+
+        set_global_object_mock(mock);
+
+        let err = PersistentObject::create(
+            ObjectStorageConstants::Private,
+            &[],
+            DataFlag::ACCESS_WRITE,
+            None,
+            &[],
+        )
+        .expect_err("it should be err");
+
+        assert_eq!(err.raw_code(), RETURN_CODE);
+    }
+
+    #[test]
+    // If a persistent object successfully `close_and_delete`, it should not
+    // call `TEE_CloseObject` anymore.
+    fn test_create_and_successfully_close_delete() {
+        let _lock = SERIAL_TEST_LOCK.lock();
+
+        let mut mock = MockObjectController::new();
+        let mut handle_struct = 
MockObjectController::new_valid_test_handle_struct();
+        let handle = MockObjectController::new_valid_test_handle(&mut 
handle_struct);
+
+        mock.expect_TEE_CreatePersistentObject_success_once(handle.clone());
+        mock.expect_TEE_CloseAndDeletePersistentObject1_success_once(handle);
+
+        set_global_object_mock(mock);
+
+        let obj = PersistentObject::create(
+            ObjectStorageConstants::Private,
+            &[],
+            DataFlag::ACCESS_WRITE,
+            None,
+            &[],
+        )
+        .expect("it should be ok");
+
+        obj.close_and_delete().expect("it should be ok");
+    }
+
+    #[test]
+    // Even a persistent object failed at `close_and_delete`, `TEE_CloseObject`
+    // should not be called.
+    fn test_create_and_failed_close_delete() {
+        let _lock = SERIAL_TEST_LOCK.lock();
+
+        let mut mock = MockObjectController::new();
+        let mut handle_struct = 
MockObjectController::new_valid_test_handle_struct();
+        let handle = MockObjectController::new_valid_test_handle(&mut 
handle_struct);
+
+        mock.expect_TEE_CreatePersistentObject_success_once(handle.clone());
+        
mock.expect_TEE_CloseAndDeletePersistentObject1_fail_once(raw::TEE_ERROR_BAD_STATE);
+
+        set_global_object_mock(mock);
+
+        let obj = PersistentObject::create(
+            ObjectStorageConstants::Private,
+            &[],
+            DataFlag::ACCESS_WRITE,
+            None,
+            &[],
+        )
+        .expect("it should be ok");
+
+        obj.close_and_delete().expect_err("it should be err");
     }
 }
diff --git a/optee-utee/src/object/transient_object.rs 
b/optee-utee/src/object/transient_object.rs
index 1a5095d..b0fb198 100644
--- a/optee-utee/src/object/transient_object.rs
+++ b/optee-utee/src/object/transient_object.rs
@@ -15,12 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use alloc::{boxed::Box, vec::Vec};
-use core::ptr;
+use alloc::vec::Vec;
 
 use optee_utee_sys as raw;
 
-use super::{Attribute, AttributeId, ObjHandle, ObjectHandle, ObjectInfo, 
UsageFlag};
+use super::{Attribute, GenericObject, ObjectHandle};
 use crate::{Error, Result};
 
 /// Define types of [TransientObject](TransientObject) with predefined maximum 
sizes.
@@ -102,12 +101,17 @@ pub enum TransientObjectType {
 /// Transient objects are used to hold a cryptographic object (key or 
key-pair).
 ///
 /// Contrast [PersistentObject](PersistentObject).
+#[derive(Debug)]
 pub struct TransientObject(ObjectHandle);
 
 impl TransientObject {
     /// Create a [TransientObject](TransientObject) with a null handle which 
points to nothing.
+    //
+    // TODO: This function is only used in examples and should be removed when
+    // TransientObject is fully refactored in the future. Keep it for now and
+    // leave it for future PR.
     pub fn null_object() -> Self {
-        Self(ObjectHandle::from_raw(ptr::null_mut()))
+        Self(ObjectHandle::new_null())
     }
 
     /// Check if current object is created with null handle.
@@ -158,14 +162,14 @@ impl TransientObject {
     /// 1) If the Implementation detects any error associated with this 
function which is not
     ///    explicitly associated with a defined return code for this function.
     pub fn allocate(object_type: TransientObjectType, max_object_size: usize) 
-> Result<Self> {
-        let raw_handle: *mut raw::TEE_ObjectHandle = 
Box::into_raw(Box::new(ptr::null_mut()));
+        let mut handle: raw::TEE_ObjectHandle = core::ptr::null_mut();
+        // Move as much code as possible out of unsafe blocks to maximize 
Rust’s
+        // safety checks.
+        let handle_mut = &mut handle;
         match unsafe {
-            raw::TEE_AllocateTransientObject(object_type as u32, 
max_object_size as u32, raw_handle)
+            raw::TEE_AllocateTransientObject(object_type as u32, 
max_object_size as u32, handle_mut)
         } {
-            raw::TEE_SUCCESS => {
-                let handle = ObjectHandle::from_raw(raw_handle);
-                Ok(Self(handle))
-            }
+            raw::TEE_SUCCESS => Ok(Self(ObjectHandle::from_raw(handle)?)),
             code => Err(Error::from_raw_error(code)),
         }
     }
@@ -233,154 +237,6 @@ impl TransientObject {
         }
     }
 
-    /// Return the characteristics of an object.
-    ///
-    /// # Example
-    ///
-    /// ``` rust,no_run
-    /// # use optee_utee::{TransientObject, TransientObjectType};
-    /// # fn main() -> optee_utee::Result<()> {
-    /// match TransientObject::allocate(TransientObjectType::Aes, 128) {
-    ///     Ok(object) => {
-    ///         match object.info() {
-    ///             Ok(info) => {
-    ///                 // ...
-    ///                 Ok(())
-    ///             }
-    ///             Err(e) => Err(e),
-    ///         }
-    ///     }
-    ///     Err(e) => Err(e),
-    /// }
-    /// # }
-    /// ```
-    ///
-    /// # Panics
-    ///
-    /// 1) If object is not a valid opened object.
-    /// 2) If the Implementation detects any other error associated with this 
function which is not
-    ///    explicitly associated with a defined return code for this function.
-    pub fn info(&self) -> Result<ObjectInfo> {
-        self.0.info()
-    }
-
-    /// Restrict the object usage flags of an object handle to contain at most 
the flags passed in the obj_usage parameter.
-    ///
-    /// The initial value of the key usage contains all usage flags.
-    ///
-    /// # Parameters
-    ///
-    /// 1) `obj_usage`: New object usage, an OR combination of one or more of 
the [UsageFlag](UsageFlag).
-    ///
-    /// # Example
-    ///
-    /// ``` rust,no_run
-    /// # use optee_utee::{TransientObject, TransientObjectType, UsageFlag};
-    /// # fn main() -> optee_utee::Result<()> {
-    /// match TransientObject::allocate(TransientObjectType::Aes, 128) {
-    ///     Ok(mut object) =>
-    ///     {
-    ///         object.restrict_usage(UsageFlag::ENCRYPT)?;
-    ///         Ok(())
-    ///     }
-    ///     Err(e) => Err(e),
-    /// }
-    /// # }
-    /// ```
-    ///
-    /// # Panics
-    ///
-    /// 1) If object is not a valid opened object.
-    /// 2) If the Implementation detects any other error associated with this 
function which is not
-    ///    explicitly associated with a defined return code for this function.
-    pub fn restrict_usage(&mut self, obj_usage: UsageFlag) -> Result<()> {
-        self.0.restrict_usage(obj_usage)
-    }
-
-    /// Extract one buffer attribute from an object. The attribute is 
identified by the argument id.
-    ///
-    /// # Parameters
-    ///
-    /// 1) `id`: Identifier of the attribute to retrieve.
-    /// 2) `ref_attr`: Output buffer to get the content of the attribute.
-    ///
-    /// # Example
-    ///
-    /// ``` rust,no_run
-    /// # use optee_utee::{TransientObject, TransientObjectType, AttributeId};
-    /// # fn main() -> optee_utee::Result<()> {
-    /// # let id = AttributeId::SecretValue;
-    /// match TransientObject::allocate(TransientObjectType::Aes, 128) {
-    ///     Ok(object) => {
-    ///         let mut attr = [0u8; 16];
-    ///         match object.ref_attribute(id, &mut attr) {
-    ///             Ok(size) => {
-    ///                 // ...
-    ///                 Ok(())
-    ///             }
-    ///             Err(e) => Err(e),
-    ///         }
-    ///     }
-    ///     Err(e) => Err(e),
-    /// }
-    /// # }
-    /// ```
-    ///
-    /// # Errors
-    ///
-    /// 1) `ItemNotFound`: If the attribute is not found on this object.
-    /// 2) `ShortBuffer`: If buffer is NULL or too small to contain the key 
part.
-    ///
-    /// # Panics
-    ///
-    /// 1) If object is not a valid opened object.
-    /// 2) If the object is not initialized.
-    /// 3) If the Attribute is not a buffer attribute.
-    pub fn ref_attribute(&self, id: AttributeId, buffer: &mut [u8]) -> 
Result<usize> {
-        self.0.ref_attribute(id, buffer)
-    }
-
-    /// Extract one value attribute from an object. The attribute is 
identified by the argument id.
-    ///
-    /// # Parameters
-    ///
-    /// 1) `id`: Identifier of the attribute to retrieve.
-    /// 2) `value_attr`: Two value placeholders to get the content of the 
attribute.
-    ///
-    /// # Example
-    ///
-    /// ``` rust,no_run
-    /// # use optee_utee::{TransientObject, TransientObjectType};
-    /// # fn main() -> optee_utee::Result<()> {
-    /// # let id = 0_u32;
-    /// match TransientObject::allocate(TransientObjectType::Aes, 128) {
-    ///     Ok(object) => {
-    ///         match object.value_attribute(id) {
-    ///             Ok((a,b)) => {
-    ///                 // ...
-    ///                 Ok(())
-    ///             }
-    ///             Err(e) => Err(e),
-    ///         }
-    ///     }
-    ///     Err(e) => Err(e),
-    /// }
-    /// # }
-    /// ```
-    ///
-    /// # Errors
-    ///
-    /// 1) `ItemNotFound`: If the attribute is not found on this object.
-    ///
-    /// # Panics
-    ///
-    /// 1) If object is not a valid opened object.
-    /// 2) If the object is not initialized.
-    /// 3) If the Attribute is not a value attribute.
-    pub fn value_attribute(&self, id: u32) -> Result<(u32, u32)> {
-        self.0.value_attribute(id)
-    }
-
     /// Populates an uninitialized object handle with the attributes of 
another object handle;
     /// that is, it populates the attributes of this handle with the 
attributes of src_handle.
     /// It is most useful in the following situations:
@@ -425,7 +281,7 @@ impl TransientObject {
     /// 3) If the type and size of src_object and self are not compatible.
     /// 4) If the Implementation detects any other error associated with this 
function which is not
     ///    explicitly associated with a defined return code for this function.
-    pub fn copy_attribute_from<T: ObjHandle>(&mut self, src_object: &T) -> 
Result<()> {
+    pub fn copy_attribute_from<T: GenericObject>(&mut self, src_object: &T) -> 
Result<()> {
         match unsafe { raw::TEE_CopyObjectAttributes1(self.handle(), 
src_object.handle()) } {
             raw::TEE_SUCCESS => Ok(()),
             code => Err(Error::from_raw_error(code)),
@@ -485,28 +341,60 @@ impl TransientObject {
     }
 }
 
-impl ObjHandle for TransientObject {
+impl GenericObject for TransientObject {
     fn handle(&self) -> raw::TEE_ObjectHandle {
         self.0.handle()
     }
 }
 
-impl Drop for TransientObject {
-    /// Deallocates a [TransientObject](TransientObject) previously allocated.
-    /// After this function has been called, the object handle is no longer 
valid and all resources
-    /// associated with the [TransientObject](TransientObject) SHALL have been 
reclaimed.
-    ///
-    /// # Panics
-    ///
-    /// 1) If object is not a valid opened object.
-    /// 2) If the Implementation detects any other error associated with this 
function which is not
-    ///    explicitly associated with a defined return code for this function.
-    fn drop(&mut self) {
-        unsafe {
-            if self.0.raw != ptr::null_mut() {
-                raw::TEE_FreeTransientObject(self.0.handle());
-            }
-            drop(Box::from_raw(self.0.raw));
-        }
+#[cfg(test)]
+mod tests {
+    use optee_utee_mock::{
+        object::{set_global_object_mock, MockObjectController, 
SERIAL_TEST_LOCK},
+        raw,
+    };
+
+    use super::*;
+
+    #[test]
+    // If a transient object is successfully allocated, TEE_CloseObject will be
+    // called when it is dropped.
+    //
+    // According to `GPD_TEE_Internal_Core_API_Specification_v1.3.1`:
+    // At 5.5.5 TEE_CloseObject:
+    // The `TEE_CloseObject` function closes an opened object handle. The 
object
+    // can be persistent or transient.
+    // For transient objects, `TEE_CloseObject` is equivalent to
+    // `TEE_FreeTransientObject`.
+    fn test_allocate_and_drop() {
+        let _lock = SERIAL_TEST_LOCK.lock();
+
+        let mut mock = MockObjectController::new();
+        let mut handle_struct = 
MockObjectController::new_valid_test_handle_struct();
+        let handle = MockObjectController::new_valid_test_handle(&mut 
handle_struct);
+
+        mock.expect_TEE_AllocateTransientObject_success_once(handle.clone());
+        mock.expect_TEE_CloseObject_once(handle);
+
+        set_global_object_mock(mock);
+
+        let _obj =
+            TransientObject::allocate(TransientObjectType::Aes, 
128).expect("it should be ok");
+    }
+
+    #[test]
+    fn test_allocate_fail() {
+        let _lock = SERIAL_TEST_LOCK.lock();
+
+        let mut mock = MockObjectController::new();
+        static RETURN_CODE: raw::TEE_Result = raw::TEE_ERROR_BAD_STATE;
+
+        mock.expect_TEE_AllocateTransientObject_fail_once(RETURN_CODE);
+        set_global_object_mock(mock);
+
+        let err =
+            TransientObject::allocate(TransientObjectType::Aes, 
128).expect_err("it should be err");
+
+        assert_eq!(err.raw_code(), RETURN_CODE);
     }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]


Reply via email to