[Draft] Initiate LOAD_TA RPC if TA not found

litebox_common_optee/src/lib.rs

@@ -679,6 +679,14 @@
         Self::from_bytes(bytes)
     }
 
+    pub fn to_u64_array(self) -> [u64; 2] {
+        let bytes = self.to_bytes();
+        [
+            u64::from_le_bytes(bytes[0..8].try_into().unwrap()),
+            u64::from_le_bytes(bytes[8..16].try_into().unwrap()),
+        ]
+    }
+
     /// Converts the UUID to a 16-byte array with little-endian encoding.
     pub fn to_le_bytes(self) -> [u8; 16] {
         let mut bytes = [0u8; 16];
@@ -688,6 +696,16 @@
         bytes[8..16].copy_from_slice(&self.clock_seq_and_node);
         bytes
     }
+
+    /// Converts the UUID to a 16-byte array with big-endian encoding (RFC 4122 format).
+    pub fn to_bytes(self) -> [u8; 16] {
+        let mut bytes = [0u8; 16];
+        bytes[0..4].copy_from_slice(&self.time_low.to_be_bytes());
+        bytes[4..6].copy_from_slice(&self.time_mid.to_be_bytes());
+        bytes[6..8].copy_from_slice(&self.time_hi_and_version.to_be_bytes());
+        bytes[8..16].copy_from_slice(&self.clock_seq_and_node);
+        bytes
+    }
 }
 
 /// TA flags from `optee_os/lib/libutee/include/user_ta_header.h`.
@@ -2117,6 +2135,34 @@
         }
     }
 
+    /// Set a parameter's attribute type by index with bounds checking against `num_params`.
+    pub fn set_param_attr_type(
+        &mut self,
+        index: usize,
+        attr_type: OpteeMsgAttrType,
+    ) -> Result<(), OpteeSmcReturnCode> {
+        if index >= self.num_params as usize {
+            Err(OpteeSmcReturnCode::ENotAvail)
+        } else {
+            self.params[index].attr = OpteeMsgAttr(attr_type as u64);
+            Ok(())
+        }
+    }
+
+    /// Set an rmem parameter by index with bounds checking against `num_params`.
+    pub fn set_param_rmem(
+        &mut self,
+        index: usize,
+        rmem: OpteeMsgParamRmem,
+    ) -> Result<(), OpteeSmcReturnCode> {
+        if index >= self.num_params as usize {
+            Err(OpteeSmcReturnCode::ENotAvail)
+        } else {
+            self.params[index].data.copy_from_slice(rmem.as_bytes());
+            Ok(())
+        }
+    }
+
     /// Set a tmem parameter by index with bounds checking against `num_params`.
     pub fn set_param_tmem(
         &mut self,
@@ -2130,10 +2176,52 @@
             Ok(())
         }
     }
+}
 
-    // Note: RPC does not use rmem params. Rmem requires pre-registered shared memory
-    // references from the normal-world driver, which is a main-messaging-path concept.
-    // RPC uses tmem for buffer references since OP-TEE provides physical addresses directly.
+/// Prepare a LOAD_TA RPC request to be sent to normal world.
+///
+/// When `memref` is `None`, the request asks normal world to return the TA size.
+/// When it is `Some`, the request provides registered memory for the TA binary.
+pub fn prepare_load_ta_rpc(
+    rpc_msg_args: &mut OpteeRpcArgs,
+    ta_uuid: TeeUuid,
+    memref_size: u64,
+    memref: Option<OpteeMsgParamRmem>,
+) -> Result<(), OpteeSmcReturnCode> {
+    rpc_msg_args.cmd = OpteeRpcCommand::LoadTa;
+    rpc_msg_args.num_params = 2;
+
+    rpc_msg_args.set_param_attr_type(0, OpteeMsgAttrType::ValueInput)?;
+    let uuid_bytes = ta_uuid.to_u64_array();
+    rpc_msg_args.set_param_value(
+        0,
+        OpteeMsgParamValue {
+            a: uuid_bytes[0],
+            b: uuid_bytes[1],
+            c: 0,
+        },
+    )?;
+
+    if memref.is_none() {
+        // First call of LOAD_TA protocol: normal world returns the TA size in memref_size.
+        rpc_msg_args.set_param_attr_type(1, OpteeMsgAttrType::TmemOutput)?;
+        rpc_msg_args.set_param_tmem(
+            1,
+            OpteeMsgParamTmem {
+                buf_ptr: 0,
+                size: memref_size,
+                shm_ref: 0,
+            },
+        )?;
+    } else {
+        // Second call of LOAD_TA protocol: secure world provides a memref for the TA binary.
+        rpc_msg_args.set_param_attr_type(1, OpteeMsgAttrType::RmemOutput)?;
+        if let Some(rmem) = memref {
+            rpc_msg_args.set_param_rmem(1, rmem)?;
+        }
+    }
+
+    Ok(())
 }
 
 /// Serialize the params portion as raw bytes into `buf`.
@@ -2517,6 +2605,11 @@
             uuid.clock_seq_and_node,
             [0xaf, 0x63, 0x00, 0x02, 0xa5, 0xd5, 0xc5, 0x1b]
         );
+        assert_eq!(
+            uuid.to_u64_array(),
+            [0xe311f8e7_e0b34f38, 0x1bc5d5a5_020063af]
+        );
+        assert_eq!(TeeUuid::from_u64_array(uuid.to_u64_array()), uuid);
     }
 
     #[test]
@@ -2616,6 +2709,58 @@
         assert_eq!(header_out.num_params, 2);
     }
 
+    #[test]
+    fn test_optee_rpc_args_attr_and_rmem_setters() {
+        let header = OpteeMsgArgsHeader {
+            cmd: OpteeRpcCommand::LoadTa as u32,
+            func: 0,
+            session: 0,
+            cancel_id: 0,
+            pad: 0,
+            ret: 0,
+            ret_origin: 0,
+            num_params: 1,
+        };
+        let raw_params = [0u8; size_of::<OpteeMsgParam>()];
+        let mut rpc_args = OpteeRpcArgs::from_header_and_raw_params(&header, &raw_params)
+            .expect("should parse RPC args");
+
+        rpc_args.params[0].attr = OpteeMsgAttr::META_VALUE_INPUT;
+        rpc_args
+            .set_param_attr_type(0, OpteeMsgAttrType::RmemOutput)
+            .expect("attribute index should be available");
+        assert_eq!(
+            rpc_args.params[0].attr.attr_type(),
+            OpteeMsgAttrType::RmemOutput as u8
+        );
+        assert!(!rpc_args.params[0].attr.meta());
+        assert!(!rpc_args.params[0].attr.noncontig());
+
+        let rmem = OpteeMsgParamRmem {
+            offs: 0x0102_0304_0506_0708,
+            size: 0x1112_1314_1516_1718,
+            shm_ref: 0x2122_2324_2526_2728,
+        };
+        rpc_args
+            .set_param_rmem(0, rmem)
+            .expect("rmem index should be available");
+        assert_eq!(&rpc_args.params[0].data[0..8], &rmem.offs.to_le_bytes());
+        assert_eq!(&rpc_args.params[0].data[8..16], &rmem.size.to_le_bytes());
+        assert_eq!(
+            &rpc_args.params[0].data[16..24],
+            &rmem.shm_ref.to_le_bytes()
+        );
+
+        assert_eq!(
+            rpc_args.set_param_attr_type(1, OpteeMsgAttrType::RmemOutput),
+            Err(OpteeSmcReturnCode::ENotAvail)
+        );
+        assert_eq!(
+            rpc_args.set_param_rmem(1, rmem),
+            Err(OpteeSmcReturnCode::ENotAvail)
+        );
+    }
+
     #[test]
     fn test_rpc_args_rejects_main_cmd() {
         // Pick a cmd value that lies in the gap between Plugin (12) and I2C Transfer (21),

litebox_runner_lvbs/src/lib.rs

@@ -17,6 +17,7 @@ use litebox_common_lvbs::{NUM_VTLCALL_PARAMS, VsmError, VsmFunction};
 use litebox_common_optee::{
     OpteeMessageCommand, OpteeMsgArgs, OpteeRpcArgs, OpteeSmcArgs, OpteeSmcResult,
     OpteeSmcReturnCode, TeeOrigin, TeeResult, UteeEntryFunc, UteeParams, optee_msg_args_total_size,
+    prepare_load_ta_rpc,
 };
 use litebox_platform_lvbs::mshv::vsm::{LvbsVtl0Gate, LvbsVtl0PrivilegedWriter, LvbsVtl1Gate};
 use litebox_platform_lvbs::{
@@ -520,14 +521,14 @@ fn optee_smc_handler(smc_args_addr: usize) -> OpteeSmcArgs {
     };
     if let OpteeSmcResult::CallWithArg {
         msg_args,
-        rpc_args: _,
+        mut rpc_args,
         msg_args_phys_addr,
     } = smc_result
     {
         let mut msg_args = *msg_args;
         debug_serial_println!("OP-TEE SMC with MsgArgs Command: {:?}", msg_args.cmd);
         let result = match msg_args.cmd {
-            OpenSession => handle_open_session(&mut msg_args, msg_args_phys_addr),
+            OpenSession => handle_open_session(&mut msg_args, &mut rpc_args, msg_args_phys_addr),
             InvokeCommand => handle_invoke_command(&mut msg_args, msg_args_phys_addr),
             CloseSession => handle_close_session(&mut msg_args, msg_args_phys_addr),
             _ => {
@@ -548,7 +549,23 @@ fn optee_smc_handler(smc_args_addr: usize) -> OpteeSmcArgs {
         unsafe { switch_to_base_page_table() };
 
         if let Err(e) = result {
-            smc_args.set_return_code(e);
+            if e == OpteeSmcReturnCode::RpcCmd {
+                debug_serial_println!("OP-TEE SMC returning RPC command to normal world");
+                let Some(rpc_args_ref) = rpc_args.as_ref() else {
+                    smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd);
+                    return *smc_args;
+                };
+                if let Err(e) =
+                    write_rpc_args_to_normal_world(&msg_args, msg_args_phys_addr, rpc_args_ref)
+                {
+                    smc_args.set_return_code(e);
+                } else {
+                    smc_args.set_return_code(OpteeSmcReturnCode::RpcCmd);
+                }
+            } else {
+                debug_serial_println!("OP-TEE SMC returning error code: {:?}", e);
+                smc_args.set_return_code(e);
+            }
         } else {
             smc_args.set_return_code(OpteeSmcReturnCode::Ok);
         }
@@ -569,18 +586,44 @@ fn optee_smc_handler(smc_args_addr: usize) -> OpteeSmcArgs {
 /// instance cleanup for TARGET_DEAD on single-instance TAs).
 fn handle_open_session(
     msg_args: &mut OpteeMsgArgs,
+    rpc_args: &mut Option<Box<OpteeRpcArgs>>,
     msg_args_phys_addr: u64,
 ) -> Result<(), OpteeSmcReturnCode> {
     let ta_req_info = decode_ta_request(msg_args).map_err(|_| OpteeSmcReturnCode::EBadCmd)?;
     if ta_req_info.entry_func != UteeEntryFunc::OpenSession {
         return Err(OpteeSmcReturnCode::EBadCmd);
     }
+    let shim: litebox_shim_optee::OpteeShim = litebox_shim_optee::OpteeShimBuilder::new().build();
 
     let ta_uuid = ta_req_info.uuid.ok_or(OpteeSmcReturnCode::EBadCmd)?;
+    if !shim.contains_ta_bin(&ta_uuid) {
+        debug_serial_println!(
+            "TA binary not found for uuid={:?}, requesting load from normal world",
+            ta_uuid
+        );
+
+        let Some(rpc) = rpc_args.as_deref_mut() else {
+            debug_serial_println!(
+                "RPC args not present in incoming request, cannot request LOAD_TA from normal world"
+            );
+            msg_args.session = 0;
+            msg_args.ret = TeeResult::ItemNotFound;
+            msg_args.ret_origin = TeeOrigin::Tee;
+            write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?;
+            return Ok(());
+        };
+        // LOAD_TA is a two-call protocol. The first call uses a zero-sized
+        // memref (a NULL buffer) so normal world returns the TA size.
+        prepare_load_ta_rpc(rpc, ta_uuid, 0, None)?;
+        return Err(OpteeSmcReturnCode::RpcCmd);
+    }
+
     let client_identity = ta_req_info.client_identity;
     let params = &ta_req_info.params;
 
     session_manager().with_ta(&ta_uuid, |target| match target {
+        // A sibling points to an already-loaded single-instance TA, so no
+        // binary cache lookup or LOAD_TA request is needed on this path.
         OpenSessionTarget::Sibling(instance) => open_session_single_instance(
             msg_args,
             msg_args_phys_addr,
@@ -784,13 +827,6 @@ fn open_session_new_instance(
     ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo<PAGE_SIZE>,
 ) -> Result<(), OpteeSmcReturnCode> {
     let shim = litebox_shim_optee::OpteeShimBuilder::new().build();
-    if shim.get_ta_bin(&ta_uuid).is_none() {
-        msg_args.session = 0;
-        msg_args.ret = TeeResult::ItemNotFound;
-        msg_args.ret_origin = TeeOrigin::Tee;
-        write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?;
-        return Ok(());
-    }
 
     // Token is declared before `task_pt_guard` so it drops AFTER it.
     // Marker only releases once CR3 is back to base. See
@@ -1318,7 +1354,6 @@ fn write_non_ta_msg_args_to_normal_world(
 /// Unlike [`write_msg_args_to_normal_world`], this function does not access TA userspace
 /// memory and can be called from the base page table context. It simply serializes the
 /// rpc_args and writes it to the normal world physical address.
-#[expect(dead_code)]
 #[inline]
 fn write_rpc_args_to_normal_world(
     msg_args: &OpteeMsgArgs,

litebox_shim_optee/src/lib.rs

@@ -190,17 +190,14 @@ impl GlobalState {
         self.ta_uuid_map.insert(*ta_uuid, ta_bin.into())
     }
 
-    /// Get the TA binary associated with the given TA UUID.
+    /// Get the cached TA binary associated with the given TA UUID.
     pub(crate) fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option<Arc<[u8]>> {
-        if let Some(ta_bin) = self.ta_uuid_map.get(ta_uuid) {
-            Some(ta_bin)
-        } else {
-            let ta_bin = Self::rpc_get_ta_bin(ta_uuid)?;
-            if !self.store_ta_bin(ta_uuid, &ta_bin) {
-                return None;
-            }
-            Some(ta_bin)
-        }
+        self.ta_uuid_map.get(ta_uuid)
+    }
+
+    /// Return whether a TA binary is cached for the given UUID.
+    pub(crate) fn contains_ta_bin(&self, ta_uuid: &TeeUuid) -> bool {
+        self.ta_uuid_map.contains(ta_uuid)
     }
 
     /// Get the TA flags associated with the given TA UUID.
@@ -227,11 +224,6 @@ impl GlobalState {
     pub(crate) fn remove_ta_bin(&self, ta_uuid: &TeeUuid) {
         let _ = self.ta_uuid_map.remove(ta_uuid);
     }
-
-    /// RPC to get the TA binary associated with the given TA UUID. Placeholder for now.
-    fn rpc_get_ta_bin(_ta_uuid: &TeeUuid) -> Option<Arc<[u8]>> {
-        None
-    }
 }
 
 type UserMutPtr<T> = <Platform as litebox::platform::RawPointerProvider>::RawMutPointer<T>;
@@ -310,11 +302,16 @@ impl OpteeShim {
         self.0.store_ta_bin(ta_uuid, ta_bin)
     }
 
-    /// Get the TA binary associated with the given TA UUID.
+    /// Get the cached TA binary associated with the given TA UUID.
     pub fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option<Arc<[u8]>> {
         self.0.get_ta_bin(ta_uuid)
     }
 
+    /// Return whether a TA binary is cached for the given UUID.
+    pub fn contains_ta_bin(&self, ta_uuid: &TeeUuid) -> bool {
+        self.0.contains_ta_bin(ta_uuid)
+    }
+
     /// Release all user-space memory mappings owned by this shim instance.
     ///
     /// This must be called before switching to the base page table and deleting
@@ -1362,6 +1359,10 @@ impl TaUuidMap {
         self.inner.read().get(uuid).map(|info| info.binary.clone())
     }
 
+    pub(crate) fn contains(&self, uuid: &TeeUuid) -> bool {
+        self.inner.read().contains_key(uuid)
+    }
+
     /// Get the TA flags for a given UUID.
     pub(crate) fn get_flags(&self, uuid: &TeeUuid) -> Option<TaFlags> {
         self.inner.read().get(uuid).map(|info| info.flags)

litebox_shim_optee/src/syscalls/ldelf.rs

@@ -252,7 +252,7 @@ impl Task {
             "sys_open_bin"
         );
 
-        if self.global.get_ta_bin(&ta_uuid).is_none() {
+        if !self.global.contains_ta_bin(&ta_uuid) {
             return Err(TeeResult::ItemNotFound);
         }
         let new_handle = self.ta_handle_map.insert(ta_uuid);