[Draft] Dynamic TA Support by praveen-pk · Pull Request #1213 · microsoft/litebox

Praveen K Paladugu (praveen-pk) · 2026-08-20T18:07:59Z

If TA is not found within the TA uuid map, initiate an RPC to VTL0 with appropriate args.

Sangho Lee (sangho2)

[AI review]

Review of the LOAD_TA RPC initiation path. Draft-stage, so I focused on wire-protocol correctness and the VTL0-facing edges rather than polish. Comments only — no blocking verdict.

Correctness

1. UUID is byte-swapped on the wire (`litebox_runner_lvbs/src/lib.rs:595-603`)

TeeUuid::to_le_bytes() is not the inverse of TeeUuid::from_bytes() / from_u64_array().

The incoming UUID is decoded in litebox_shim_optee/src/msg_handler.rs:453 via TeeUuid::from_u64_array([a, b]), which lays the two LE u64s down as 16 octets and then reads time_low/time_mid/time_hi_and_version as big-endian (RFC 4122 octet order, litebox_common_optee/src/lib.rs:655-658). to_le_bytes() re-emits those same three fields little-endian, so each is byte-swapped relative to the octets we received.

Using the existing test vector at litebox_common_optee/src/lib.rs:2549:

driver sent:  params[0].u.value.a = 0xe311f8e7_e0b34f38
this code sends:                    0x11e3e7f8_384fb3e0

tee-supplicant does uuid_from_octets(&uuid, (void *)&params[0].a) (RFC 4122 octets), matching optee_os's tee_uuid_to_octets() in rpc_load(). So normal world will look up the wrong TA.

Suggest adding a to_bytes() / to_u64_array() on TeeUuid that mirrors from_bytes / from_u64_array, with a round-trip unit test, and using that here.

(Note to_le_bytes()'s only other caller, syscalls/pta.rs:309, is a HUK KDF input where only self-consistency matters — that's why this has gone unnoticed.)

2. Stale rmem fields leak into the RPC (`prepare_load_ta_rpc`, `lib.rs:583-616`)

rpc_args is parsed directly out of normal-world memory (read_optee_msg_args_from_phys, msg_handler.rs:182-204), so params[1].data can contain whatever the driver/previous RPC left there. The function sets attr and the size field but never clears offs (data[0..8]) or shm_ref (data[16..24]).

optee_os emits an all-zero rmem for a NULL memref (get_rpc_arg() in core/kernel/thread.c), and the Linux driver will try to resolve a non-zero shm_ref cookie. Please zero params[..num_params] (or at least param 1) before populating.

3. `set_param_memref_size` and `set_param_rmem` clobber each other (`lib.rs:609-614`)

set_param_rmem does data.copy_from_slice(rmem.as_bytes()) over all 24 bytes, so when memref is Some, the memref_size written on the previous line is silently discarded. Either drop the memref_size parameter when a full rmem is supplied, or set rmem.size = memref_size before writing. As written, the stage-2 call path is already broken.

4. `unwrap()` on a normal-world-driven path (`lib.rs:554`)

let rpc_args_ref = rpc_args.as_ref().unwrap();

A panic here is a VTL1 kernel panic (#[panic_handler] → raise_vtl0_gp_fault). Today it's guaranteed Some because handle_open_session is the only RpcCmd producer and it bails out earlier if rpc_args is None — but that's an implicit cross-function invariant that the next RpcCmd producer will break. Prefer:

let Some(rpc_args_ref) = rpc_args.as_ref() else {
    smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd);
    return *smc_args;
};

5. Behavior regression when `rpc_args` is `None` (`lib.rs:642`)

For a plain OpteeSmcFunction::CallWithArg (driver without RPC_ARG), rpc_args is None (msg_handler.rs:224-228), so a cache miss now returns Err(EBadCmd) — the driver fails the whole SMC. Previously (and still, in open_session_new_instance at lib.rs:857-864) a missing TA produced a clean TeeResult::ItemNotFound in msg_args with Ok(()). Suggest falling through to the existing ItemNotFound path instead of EBadCmd when RPC isn't available.

Design / follow-up

6. Nothing can resume the RPC yet

OpteeSmcFunction (litebox_common_optee/src/lib.rs:2276-2288) has no OPTEE_SMC_FUNCID_RETURN_FROM_RPC variant, so func_id() (lib.rs:2226) returns EBadCmd when the driver re-enters after servicing LOAD_TA, and no pending open-session state is persisted anywhere. Net effect of this PR standalone: a cached-TA miss goes from "clean ItemNotFound" to "failed OpenSession". Worth stating explicitly in the PR description which stage adds resume, and confirming stage 1 won't be merged to main ahead of it (or is gated).

7. Duplicated cache-miss check

handle_open_session (lib.rs:637) and open_session_new_instance (lib.rs:858) now both do get_ta_bin(..).is_none(). Consider a single detection point. Also note the single-instance/sibling path never reaches the new check — presumably intentional (TA already resident), but worth a comment.

8. Duplicate binding (`lib.rs:635` and `lib.rs:649`)

let ta_uuid = ta_req_info.uuid.ok_or(OpteeSmcReturnCode::EBadCmd)?;

appears twice; the second is a redundant shadow. Please drop it.

9. Inconsistent error mapping (`lib.rs:597`, `lib.rs:608`)

.map_err(|_| OpteeSmcReturnCode::EBadCmd) on set_param_attr_type, which already returns OpteeSmcReturnCode — and it discards the more accurate ENotAvail. The neighboring set_param_value / set_param_memref_size calls just use ?. Use ? throughout.

10. `prepare_load_ta_rpc` visibility and shape (`lib.rs:583`)

It's pub in the runner crate with no external caller. Make it private, or move it into litebox_common_optee alongside the other RPC helpers so both runners can use it. Also, memref_size/memref are always 0/None today — per the repo's "no speculative flexibility" guidance, consider trimming until stage 2 actually needs them (and see #3).

11. New setters in `litebox_common_optee` (`lib.rs:2120-2162`)

No unit tests, despite the crate having a test module with existing param round-trip coverage (test_optee_rpc_args_roundtrip). These are public API on the VTL0 boundary — worth covering, especially bounds behavior and the rmem layout offsets.
set_param_memref_size's doc says "rmem parameter" but it writes data[8..16] unconditionally; it silently "succeeds" on a value param. Either document that it's layout-based and attr-agnostic (it's also valid for tmem), or validate the attr type.
set_param_attr_type overwrites the whole attr word, dropping META/NONCONTIG bits. Fine for RPC args, but worth a doc note given OpteeMsgAttr carries those flags.

12. Deleted rationale comment (`litebox_common_optee/src/lib.rs`, removed lines after `set_param_tmem`)

The "RPC does not use rmem params" note is now obsolete — but rather than deleting it outright, consider replacing it with the actual rule: optee_os maps a NULL memref to RMEM_* with an all-zero body and a registered-shm memref to RMEM_*, tmem otherwise. That's non-obvious and directly informs #2.

13. Minor

shim.get_ta_bin(&ta_uuid).is_none() clones an Arc<[u8]> just for a presence check; a contains_ta_bin/has_ta_bin would be cheaper and clearer.
OpteeShimBuilder::new().build() per OpenSession (lib.rs:634) constructs a fresh LiteBox + PageManager — consistent with existing sites (lib.rs:229, lib.rs:857), but this PR adds a second one per open-session call.
Dropping rpc_get_ta_bin is a clean no-op removal (it always returned None, and no caller depended on the fallback) — nice cleanup. 👍

Sangho Lee (sangho2)

Left some comments.

Sangho Lee (sangho2) · 2026-09-01T17:28:21Z

-            }
-            Some(ta_bin)
-        }
+        self.ta_uuid_map.get(ta_uuid)


Fine for now, but I think we at least need to maintain TODO for RPC or TA binary pinning. This works now because we never call remove_ta_bin. However, if we exercise it, loading TA binaries only at the handle_open_session function can suffer from TOCTOU issues. We need to either implement real RPC for TA loading, or pin Arc<ta_bin> until we load it into the memory. Of course, not for this PR series.

From what I checked, OP-TEE does not support TA binary pinning at all. So, we should not implement it either.

Clarification: we need pinning because we use Arc here. if no one refcounts a TA binary (including uuid map itself due to remove_ta_bin), the binary will be removed from the memory and ldelf might fail to read it.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>

Track trusted continuation state across the multi-call Dynamic TA loading sequence in VTL1. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>

Replace direct physical-address writes with bounded writes through registered shared-memory bookkeeping. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>

Initiate the first LOAD_TA if the TA is not present in secure-world cache. To resume this Dynamic TA load sequence safely, track some context in VTL1. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>

Read the TA Size from VTL0, send SHM_ALLOC to get VTL0 to allocate memory to store TA Binary. Receive the allocation in VTL1, register that memory into a shm object and send the final LOAD_TA request. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>

Load TA binary and initiate OpenSession. Initiate SHM_FREE RPC to clean up the memory allocated in VTL0. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>

Praveen K Paladugu (praveen-pk) · 2026-09-16T15:05:38Z

fyi, I am working on a fix for race conditions while invoking Dynamic TA in multi-instance case. I will push that as a follow up PR.

Sangho Lee (sangho2)

I think this PR is functionally fine (except for known bugs which are already noted), but need some code cleaning and refactoring.

Sangho Lee (sangho2) · 2026-09-16T19:07:34Z

+}
+
+/// Prepare a shared-memory allocation RPC request to be sent to normal world.
+pub fn prepare_shm_alloc_rpc(


nits: any particular reasons to make these functions (prepare_shm_alloc_rpc, prepare_shm_free_rpc, ...) free functions? It would be better to make them OpteeRpcArgs methods. We can use &mut self and drop _rpc.

As most of these methods to modify rpc_args, makes sense to make them OpteeRpcArgs methods. I will update this in next version.

Sangho Lee (sangho2) · 2026-09-16T21:08:43Z

+
+    rpc_msg_args
+        .set_param_attr_type(0, OpteeMsgAttrType::ValueInput)
+        .map_err(|_| OpteeSmcReturnCode::EBadCmd)?;


do we need .map_err()? Other prepare_*_rpc functions don't do this. inconsistency.

make sense to drop map_err and make this consistent with rest of the methods

Sangho Lee (sangho2) · 2026-09-16T21:13:42Z

+    shm_ref: u64,
+) -> Result<(), OpteeSmcReturnCode> {
+    rpc_msg_args.cmd = OpteeRpcCommand::ShmFree;
+    rpc_msg_args.ret = TeeResult::GenericError;


nits: this doesn't have a comment unlike the above and below functions. Perhaps better to make only the above one have a comment, and let this one and the below one refer to it.

I will find a way to make this consistent across these methods

Sangho Lee (sangho2) · 2026-09-16T21:24:46Z

+    #[test]
+    fn test_rpc_context_id_roundtrip() {
+        for context_id in [0, 1, u32::MAX] {
+            let mut args = OpteeSmcArgs::default();
+            args.set_rpc_context_id(context_id);
+            assert_eq!(args.get_rpc_context_id(), Ok(context_id));
+        }
+    }
+
+    #[cfg(target_pointer_width = "64")]
+    #[test]
+    fn test_rpc_context_id_rejects_upper_bits() {
+        let mut args = OpteeSmcArgs::default();
+        args.args[3] = (u32::MAX as usize) + 1;
+        assert_eq!(args.get_rpc_context_id(), Err(OpteeSmcReturnCode::EBadCmd));
+    }


These tests are too trivial, not worthwhile to maintain.

Sangho Lee (sangho2) · 2026-09-16T21:33:58Z

+    let is_return_from_rpc = smc_args.func_id() == Ok(OpteeSmcFunction::ReturnFromRpc);
+    let smc_result = if is_return_from_rpc {
+        let context_id = match smc_args.get_rpc_context_id() {
+            Ok(context_id) => context_id,
+            Err(error) => {
+                smc_args.set_return_code(error);
+                return *smc_args;
+            }
+        };
+        let Some(registered_shm_ref) = rpc_context_map().get_registered_shm_ref(context_id) else {
+            smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd);
+            return *smc_args;
+        };
+        let Some(regd_shm_offset) = rpc_context_map().get_regd_shm_offset(context_id) else {
+            smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd);
+            return *smc_args;
+        };
+        read_optee_msg_args_from_regd_shm(platform, registered_shm_ref, regd_shm_offset).and_then(
+            |(msg_args, rpc_args, msg_args_phys_addr)| {
+                Ok(OpteeSmcResult::ReturnFromRpc {
+                    msg_args,
+                    rpc_args: rpc_args.ok_or(OpteeSmcReturnCode::EBadAddr)?,
+                    msg_args_phys_addr,
+                })
+            },
+        )
+    } else {
+        handle_optee_smc_args(platform, &mut smc_args)
+    };
+    let Ok(smc_result) = smc_result else {
+        if is_return_from_rpc && let Ok(context_id) = smc_args.get_rpc_context_id() {
+            discard_rpc_context(context_id);
+        }


The RPC branch (is_return_from_rpc == true) already knows context_id. why not call discard_rpc_context at the end of that branch?

This was an attempt to consolidate discard_rpc_context calls. If not, this, I have to call it in every error case.

As I am already doing this in the rpc helper methods, I will pull discard_rpc_context into the if block and call it when necessary.

Sangho Lee (sangho2) · 2026-09-17T00:00:53Z

+        let mut contexts = self.inner.lock();
+        if contexts.len() >= self.max_contexts {
+            return Err(RpcContextError::Full);
+        }
+
+        // With N active entries, at least one of N + 1 consecutive IDs is free.
+        for _ in 0..=contexts.len() {
+            let context_id = self.next_id.fetch_add(1, Ordering::Relaxed);


double synchronization: self.inner.lock() and self.next_id.fetch_add(). next_id doesn't need to be atomic in this case.

Sangho Lee (sangho2) · 2026-09-17T00:07:34Z

+    }
+
+    /// Get the current stage for `context_id`.
+    pub fn get_curr_stage(&self, context_id: u32) -> Option<RpcStage> {


most of this get_* methods are just indirections. Why not return RpcContext using context_id and call its methods? they are already public.

Sangho Lee (sangho2) · 2026-09-17T00:08:07Z

+    }
+
+    /// Return whether there are no active RPC contexts.
+    pub fn is_empty(&self) -> bool {


test-only function.

Sangho Lee (sangho2) · 2026-09-17T00:11:51Z

+
+/// Trusted state for an RPC-backed Dynamic TA request.
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct RpcContext {


I think RpcContext should be enum because its field have stage-specific meanings. The current representation looks more like C not Rust.

Sangho Lee (sangho2) · 2026-09-17T00:14:45Z

+        &self,
+        context_id: u32,
+        expected: RpcStage,
+        next: RpcStage,


Why does this function need to get next? It should be able to figure out the next stage by itself because the current state machine simple enough (e.g., LoadTaSize should be followed by ShmAlloc if there is no error)?

Sangho Lee (sangho2) reviewed Aug 20, 2026

View reviewed changes

Praveen K Paladugu (praveen-pk) force-pushed the dyn_ta/stage1 branch from d909e53 to 5576449 Compare August 25, 2026 17:12

Praveen K Paladugu (praveen-pk) commented Aug 28, 2026

View reviewed changes

Comment thread litebox_common_optee/src/lib.rs

Praveen K Paladugu (praveen-pk) force-pushed the dyn_ta/stage1 branch from 5576449 to d08bb5a Compare August 31, 2026 21:10

Sangho Lee (sangho2) reviewed Sep 1, 2026

View reviewed changes

Praveen K Paladugu (praveen-pk) and others added 6 commits September 15, 2026 19:17

optee: Add Dynamic TA RPC protocol helpers

bdddc22

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>

optee: Add trusted context tracking for Dynamic TA RPCs

64564db

Track trusted continuation state across the multi-call Dynamic TA loading sequence in VTL1. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>

Praveen K Paladugu (praveen-pk) force-pushed the dyn_ta/stage1 branch from d08bb5a to d96e276 Compare September 15, 2026 19:34

Praveen K Paladugu (praveen-pk) changed the title ~~[Draft] Initiate LOAD_TA RPC if TA not found~~ [Draft] Dynamic TA Support Sep 16, 2026

Sangho Lee (sangho2) reviewed Sep 17, 2026

View reviewed changes

Conversation

Praveen K Paladugu (praveen-pk) commented Aug 20, 2026

Uh oh!

Sangho Lee (sangho2) left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Correctness

1. UUID is byte-swapped on the wire (litebox_runner_lvbs/src/lib.rs:595-603)

2. Stale rmem fields leak into the RPC (prepare_load_ta_rpc, lib.rs:583-616)

3. set_param_memref_size and set_param_rmem clobber each other (lib.rs:609-614)

4. unwrap() on a normal-world-driven path (lib.rs:554)

5. Behavior regression when rpc_args is None (lib.rs:642)

Design / follow-up

6. Nothing can resume the RPC yet

7. Duplicated cache-miss check

8. Duplicate binding (lib.rs:635 and lib.rs:649)

9. Inconsistent error mapping (lib.rs:597, lib.rs:608)

10. prepare_load_ta_rpc visibility and shape (lib.rs:583)

11. New setters in litebox_common_optee (lib.rs:2120-2162)

12. Deleted rationale comment (litebox_common_optee/src/lib.rs, removed lines after set_param_tmem)

13. Minor

Uh oh!

Sangho Lee (sangho2) left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Sangho Lee (sangho2) Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Praveen K Paladugu (praveen-pk) commented Sep 16, 2026

Uh oh!

Sangho Lee (sangho2) left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Sangho Lee (sangho2) left a comment •

edited

Loading

1. UUID is byte-swapped on the wire (`litebox_runner_lvbs/src/lib.rs:595-603`)

2. Stale rmem fields leak into the RPC (`prepare_load_ta_rpc`, `lib.rs:583-616`)

3. `set_param_memref_size` and `set_param_rmem` clobber each other (`lib.rs:609-614`)

4. `unwrap()` on a normal-world-driven path (`lib.rs:554`)

5. Behavior regression when `rpc_args` is `None` (`lib.rs:642`)

8. Duplicate binding (`lib.rs:635` and `lib.rs:649`)

9. Inconsistent error mapping (`lib.rs:597`, `lib.rs:608`)

10. `prepare_load_ta_rpc` visibility and shape (`lib.rs:583`)

11. New setters in `litebox_common_optee` (`lib.rs:2120-2162`)

12. Deleted rationale comment (`litebox_common_optee/src/lib.rs`, removed lines after `set_param_tmem`)

Sangho Lee (sangho2) Sep 3, 2026 •

edited

Loading