[Draft] Initiate LOAD_TA RPC if TA not found

[Praveen K Paladugu (praveen-pk)]

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)]

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.

[Praveen K Paladugu (praveen-pk)]

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

[Sangho Lee (sangho2)]

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.