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 47280afb5cdedafb91e5d672d78fb34ae3f67b7a Author: ivila <[email protected]> AuthorDate: Tue Mar 25 11:17:20 2025 +0800 optee-teec: refactor Context and Session to improve usability. * make as_raw_mut_ptr unvisible outside the crate to prevent misuse. * eliminate lifetime constraint of Session. * implement Send + Sync for both Context and Session. * eliminate lifetime parameter in mnist-rs example. Signed-off-by: Zehui Chen <[email protected]> Reviewed-by: Yuan Zhuang <[email protected]> --- examples/mnist-rs/host/src/commands/serve.rs | 8 +- examples/mnist-rs/host/src/tee.rs | 18 ++-- optee-teec/Cargo.toml | 1 - optee-teec/src/context.rs | 118 +++++++++++++++------------ optee-teec/src/extension.rs | 8 +- optee-teec/src/lib.rs | 6 +- optee-teec/src/operation.rs | 4 +- optee-teec/src/session.rs | 78 ++++++++++-------- 8 files changed, 132 insertions(+), 109 deletions(-) diff --git a/examples/mnist-rs/host/src/commands/serve.rs b/examples/mnist-rs/host/src/commands/serve.rs index 6f8b3e1..99ef605 100644 --- a/examples/mnist-rs/host/src/commands/serve.rs +++ b/examples/mnist-rs/host/src/commands/serve.rs @@ -59,7 +59,7 @@ pub fn execute(args: &Args) -> anyhow::Result<()> { } fn handle( - caller: &mut crate::tee::InferenceTaConnector<'_>, + caller: &mut crate::tee::InferenceTaConnector, request: &mut Request, ) -> anyhow::Result<HttpResponse> { if request.method().ne(&tiny_http::Method::Post) { @@ -74,7 +74,7 @@ fn handle( } fn handle_image( - caller: &mut crate::tee::InferenceTaConnector<'_>, + caller: &mut crate::tee::InferenceTaConnector, request: &mut Request, ) -> anyhow::Result<HttpResponse> { let mut data = Vec::with_capacity(IMAGE_SIZE); @@ -93,7 +93,7 @@ fn handle_image( } fn handle_binary( - caller: &mut crate::tee::InferenceTaConnector<'_>, + caller: &mut crate::tee::InferenceTaConnector, request: &mut Request, ) -> anyhow::Result<HttpResponse> { let mut data = Vec::with_capacity(IMAGE_SIZE); @@ -108,7 +108,7 @@ fn handle_binary( } fn handle_infer( - caller: &mut crate::tee::InferenceTaConnector<'_>, + caller: &mut crate::tee::InferenceTaConnector, image: &[u8], ) -> anyhow::Result<u8> { let result = caller.infer_batch(bytemuck::cast_slice(image))?; diff --git a/examples/mnist-rs/host/src/tee.rs b/examples/mnist-rs/host/src/tee.rs index d59dbc9..a3ff51c 100644 --- a/examples/mnist-rs/host/src/tee.rs +++ b/examples/mnist-rs/host/src/tee.rs @@ -21,12 +21,12 @@ use proto::{inference, train, Image}; const MAX_OUTPUT_SERIALIZE_SIZE: usize = 1 * 1024; const MAX_MODEL_RECORD_SIZE: usize = 10 * 1024 * 1024; -pub struct TrainerTaConnector<'a> { - sess: Session<'a>, +pub struct TrainerTaConnector { + sess: Session, } -impl<'a> TrainerTaConnector<'a> { - pub fn new(ctx: &'a mut Context, learning_rate: f64) -> optee_teec::Result<Self> { +impl TrainerTaConnector { + pub fn new(ctx: &mut Context, learning_rate: f64) -> optee_teec::Result<Self> { let bytes = learning_rate.to_le_bytes(); let uuid = Uuid::parse_str(train::UUID).map_err(|err| { println!("parse uuid \"{}\" failed due to: {:?}", train::UUID, err); @@ -106,14 +106,12 @@ impl<'a> TrainerTaConnector<'a> { } } -pub struct InferenceTaConnector<'a> { - sess: Session<'a>, +pub struct InferenceTaConnector { + sess: Session, } -unsafe impl Send for InferenceTaConnector<'_> {} - -impl<'a> InferenceTaConnector<'a> { - pub fn new(ctx: &'a mut Context, record: &[u8]) -> optee_teec::Result<Self> { +impl InferenceTaConnector { + pub fn new(ctx: &mut Context, record: &[u8]) -> optee_teec::Result<Self> { let uuid = Uuid::parse_str(inference::UUID).map_err(|err| { println!( "parse uuid \"{}\" failed due to: {:?}", diff --git a/optee-teec/Cargo.toml b/optee-teec/Cargo.toml index 8a9e0d5..e77c1d2 100644 --- a/optee-teec/Cargo.toml +++ b/optee-teec/Cargo.toml @@ -27,7 +27,6 @@ edition = "2018" [dependencies] optee-teec-sys = { version = "0.4.0", path = "optee-teec-sys" } optee-teec-macros = { version = "0.4.0", path = "macros" } -libc = "0.2" uuid = "0.7" hex = "0.3" num_enum = "0.7.3" diff --git a/optee-teec/src/context.rs b/optee-teec/src/context.rs index 4f37a10..78fd9c2 100644 --- a/optee-teec/src/context.rs +++ b/optee-teec/src/context.rs @@ -16,70 +16,78 @@ // under the License. use crate::{raw, Error, Operation, Param, ParamNone, Result, Session, Uuid}; -use libc; -use std::ptr; +use std::{cell::RefCell, ptr, rc::Rc}; + +pub struct InnerContext(pub raw::TEEC_Context); + +impl Drop for InnerContext { + fn drop(&mut self) { + unsafe { + raw::TEEC_FinalizeContext(&mut self.0); + } + } +} /// An abstraction of the logical connection between a client application and a /// TEE. pub struct Context { - raw: raw::TEEC_Context, + // Use Rc to share it with Session, eliminating the lifetime constraint. + // Use RefCell to allow conversion into a raw mutable pointer. + // As RefCell is not Send + Sync, there is no need to use Arc. + raw: Rc<RefCell<InnerContext>>, } +// Since RefCell is used for Context, Rust does not automatically implement +// Send and Sync for it. We need to manually implement them and ensure that +// InnerContext is used correctly. +unsafe impl Send for Context {} +unsafe impl Sync for Context {} + impl Context { /// Creates a TEE client context object. /// /// # Examples /// - /// ``` - /// let ctx = Context::new().unwrap(); + /// ``` no_run + /// # use optee_teec::Context; + /// # fn main() -> optee_teec::Result<()> { + /// let mut ctx = Context::new()?; + /// # Ok(()) + /// # } /// ``` pub fn new() -> Result<Context> { - Context::new_raw(0, true, false).map(|raw| Context { raw }) - } - - /// Creates a raw TEE client context with implementation defined parameters. - /// - /// # Examples - /// - /// ``` - /// let raw_ctx: optee_teec::raw::TEEC_Context = Context::new_raw(0, true).unwrap(); - /// ``` - pub fn new_raw(fd: libc::c_int, reg_mem: bool, memref_null: bool) -> Result<raw::TEEC_Context> { + // define an empty TEEC_Context let mut raw_ctx = raw::TEEC_Context { - fd, - reg_mem, - memref_null, + fd: 0, + reg_mem: false, + memref_null: false, }; - unsafe { - match raw::TEEC_InitializeContext(ptr::null_mut() as *mut libc::c_char, &mut raw_ctx) { - raw::TEEC_SUCCESS => Ok(raw_ctx), - code => Err(Error::from_raw_error(code)), - } + match unsafe { raw::TEEC_InitializeContext(ptr::null_mut(), &mut raw_ctx) } { + raw::TEEC_SUCCESS => Ok(Self { + raw: Rc::new(RefCell::new(InnerContext(raw_ctx))), + }), + code => Err(Error::from_raw_error(code)), } } - /// Converts a TEE client context to a raw pointer. - /// - /// # Examples - /// - /// ``` - /// let mut ctx = Context::new().unwrap(); - /// let mut raw_ptr: *mut optee_teec::raw::TEEC_Context = ctx.as_mut_raw_ptr(); - /// ``` - pub fn as_mut_raw_ptr(&mut self) -> *mut raw::TEEC_Context { - &mut self.raw - } - /// Opens a new session with the specified trusted application. /// /// The target trusted application is specified by `uuid`. /// /// # Examples /// - /// ``` - /// let mut ctx = Context::new().unwrap(); - /// let uuid = Uuid::parse_str("8abcf200-2450-11e4-abe2-0002a5d5c51b").unwrap(); - /// let session = ctx.open_session(uuid).unwrap(); + /// ``` no_run + /// use optee_teec::{Context, ErrorKind, Uuid}; + /// + /// fn main() -> optee_teec::Result<()> { + /// let mut ctx = Context::new()?; + /// let uuid = Uuid::parse_str("8abcf200-2450-11e4-abe2-0002a5d5c51b").map_err(|err| { + /// println!("bad uuid: {:?}", err); + /// ErrorKind::BadParameters + /// })?; + /// let session = ctx.open_session(uuid)?; + /// Ok(()) + /// } /// ``` pub fn open_session(&mut self, uuid: Uuid) -> Result<Session> { Session::new( @@ -96,12 +104,20 @@ impl Context { /// /// # Examples /// - /// ``` - /// let mut ctx = Context::new().unwrap(); - /// let uuid = Uuid::parse_str("8abcf200-2450-11e4-abe2-0002a5d5c51b").unwrap(); - /// let p0 = ParamValue(42, 0, ParamType::ValueInout); - /// let mut operation = Operation::new(0, p0, ParamNone, ParamNone, ParamNone); - /// let session = ctx.open_session_with_operation(uuid, operation).unwrap(); + /// ``` no_run + /// use optee_teec::{Context, ErrorKind, Operation, ParamNone, ParamType, ParamValue, Uuid}; + /// + /// fn main() -> optee_teec::Result<()> { + /// let mut ctx = Context::new()?; + /// let uuid = Uuid::parse_str("8abcf200-2450-11e4-abe2-0002a5d5c51b").map_err(|err| { + /// println!("bad uuid: {:?}", err); + /// ErrorKind::BadParameters + /// })?; + /// let p0 = ParamValue::new(42, 0, ParamType::ValueInout); + /// let mut operation = Operation::new(0, p0, ParamNone, ParamNone, ParamNone); + /// let session = ctx.open_session_with_operation(uuid, &mut operation)?; + /// Ok(()) + /// } /// ``` pub fn open_session_with_operation<A: Param, B: Param, C: Param, D: Param>( &mut self, @@ -112,10 +128,10 @@ impl Context { } } -impl Drop for Context { - fn drop(&mut self) { - unsafe { - raw::TEEC_FinalizeContext(&mut self.raw); - } +// Internal usage only +impl Context { + // anyone who wants to access the inner_context must take this as mut. + pub(crate) fn inner_context(&mut self) -> Rc<RefCell<InnerContext>> { + self.raw.clone() } } diff --git a/optee-teec/src/extension.rs b/optee-teec/src/extension.rs index b479ac6..206373a 100644 --- a/optee-teec/src/extension.rs +++ b/optee-teec/src/extension.rs @@ -17,7 +17,7 @@ use crate::raw; use crate::{Error, ErrorKind, Result}; -use libc::c_char; +use core::ffi::c_char; #[repr(C)] pub struct PluginMethod { @@ -48,9 +48,9 @@ pub struct PluginParameters<'a> { impl<'a> PluginParameters<'a> { pub fn new(cmd: u32, sub_cmd: u32, inout: &'a mut [u8]) -> Self { Self { - cmd: cmd, - sub_cmd: sub_cmd, - inout: inout, + cmd, + sub_cmd, + inout, outlen: 0 as usize, } } diff --git a/optee-teec/src/lib.rs b/optee-teec/src/lib.rs index b1e3aed..5220ce1 100644 --- a/optee-teec/src/lib.rs +++ b/optee-teec/src/lib.rs @@ -16,12 +16,12 @@ // under the License. pub use self::context::Context; -pub use self::error::{Error, ErrorKind, Result}; +pub use self::error::{Error, ErrorKind, ErrorOrigin, Result}; +pub use self::extension::*; pub use self::operation::Operation; pub use self::parameter::{Param, ParamNone, ParamTmpRef, ParamType, ParamTypes, ParamValue}; pub use self::session::{ConnectionMethods, Session}; pub use self::uuid::Uuid; -pub use self::extension::*; pub use optee_teec_macros::{plugin_init, plugin_invoke}; // Re-export optee_teec_sys so developers don't have to add it to their cargo // dependencies. @@ -29,8 +29,8 @@ pub use optee_teec_sys as raw; mod context; mod error; +mod extension; mod operation; mod parameter; mod session; mod uuid; -mod extension; diff --git a/optee-teec/src/operation.rs b/optee-teec/src/operation.rs index a7be290..985b954 100644 --- a/optee-teec/src/operation.rs +++ b/optee-teec/src/operation.rs @@ -22,7 +22,7 @@ use std::{marker::PhantomData, mem}; /// invoke command operation. It is also used for cancellation of operations, /// which may be desirable even if no payload is passed. pub struct Operation<A, B, C, D> { - pub raw: raw::TEEC_Operation, + raw: raw::TEEC_Operation, phantom0: PhantomData<A>, phantom1: PhantomData<B>, phantom2: PhantomData<C>, @@ -50,7 +50,7 @@ impl<A: Param, B: Param, C: Param, D: Param> Operation<A, B, C, D> { } } - pub fn as_mut_raw_ptr(&mut self) -> *mut raw::TEEC_Operation { + pub(crate) fn as_mut_raw_ptr(&mut self) -> *mut raw::TEEC_Operation { &mut self.raw } diff --git a/optee-teec/src/session.rs b/optee-teec/src/session.rs index fb55d68..8a24bc5 100644 --- a/optee-teec/src/session.rs +++ b/optee-teec/src/session.rs @@ -15,11 +15,9 @@ // specific language governing permissions and limitations // under the License. -use libc; -use std::marker; -use std::ptr; - +use super::context::InnerContext; use crate::{raw, Context, Error, Operation, Param, Result, Uuid}; +use std::{cell::RefCell, ptr, rc::Rc}; /// Session login methods. #[derive(Copy, Clone)] @@ -39,51 +37,63 @@ pub enum ConnectionMethods { } /// Represents a connection between a client application and a trusted application. -pub struct Session<'ctx> { +pub struct Session { raw: raw::TEEC_Session, - _marker: marker::PhantomData<&'ctx mut Context>, + + // Just a holder to ensure InnerContext is not dropped and to eliminate the + // lifetime constraint, never use it. + _ctx: Rc<RefCell<InnerContext>>, } -impl<'ctx> Session<'ctx> { +// Since raw::TEEC_Session contains a raw pointer, Rust does not automatically +// implement Send and Sync for it. We need to manually implement them and ensure +// that raw::TEEC_Session is used safely. +unsafe impl Send for Session {} +unsafe impl Sync for Session {} + +impl Session { /// Initializes a TEE session object with specified context and uuid. pub fn new<A: Param, B: Param, C: Param, D: Param>( - context: &'ctx mut Context, + context: &mut Context, uuid: Uuid, operation: Option<&mut Operation<A, B, C, D>>, ) -> Result<Self> { + // define an empty TEEC_Session let mut raw_session = raw::TEEC_Session { - ctx: context.as_mut_raw_ptr(), + ctx: ptr::null_mut(), session_id: 0, }; + // define all parameters for raw::TEEC_OpenSession outside of the unsafe + // block to maximize Rust's safety checks and leverage the compiler's + // validation. let mut err_origin: u32 = 0; let raw_operation = match operation { Some(o) => o.as_mut_raw_ptr(), - None => ptr::null_mut() as *mut raw::TEEC_Operation, + None => ptr::null_mut(), }; - unsafe { - match raw::TEEC_OpenSession( - context.as_mut_raw_ptr(), + let inner_ctx = context.inner_context(); + let raw_ctx = &mut inner_ctx.borrow_mut().0; + let raw_uuid = uuid.as_raw_ptr(); + + match unsafe { + raw::TEEC_OpenSession( + raw_ctx, &mut raw_session, - uuid.as_raw_ptr(), + raw_uuid, ConnectionMethods::LoginPublic as u32, - ptr::null() as *const libc::c_void, + ptr::null(), raw_operation, &mut err_origin, - ) { - raw::TEEC_SUCCESS => Ok(Self { - raw: raw_session, - _marker: marker::PhantomData, - }), - code => Err(Error::from_raw_error(code)), - } + ) + } { + raw::TEEC_SUCCESS => Ok(Self { + raw: raw_session, + _ctx: context.inner_context(), + }), + code => Err(Error::from_raw_error(code).with_origin(err_origin.into())), } } - /// Converts a TEE client context to a raw pointer. - pub fn as_mut_raw_ptr(&mut self) -> *mut raw::TEEC_Session { - &mut self.raw - } - /// Invokes a command with an operation with this session. pub fn invoke_command<A: Param, B: Param, C: Param, D: Param>( &mut self, @@ -91,21 +101,21 @@ impl<'ctx> Session<'ctx> { operation: &mut Operation<A, B, C, D>, ) -> Result<()> { let mut err_origin: u32 = 0; - unsafe { - match raw::TEEC_InvokeCommand( + match unsafe { + raw::TEEC_InvokeCommand( &mut self.raw, command_id, operation.as_mut_raw_ptr(), &mut err_origin, - ) { - raw::TEEC_SUCCESS => Ok(()), - code => Err(Error::from_raw_error(code)), - } + ) + } { + raw::TEEC_SUCCESS => Ok(()), + code => Err(Error::from_raw_error(code).with_origin(err_origin.into())), } } } -impl<'ctx> Drop for Session<'ctx> { +impl Drop for Session { fn drop(&mut self) { unsafe { raw::TEEC_CloseSession(&mut self.raw); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
