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

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

set_param_rmem would overwrite the memref size that set_pram_memref_size wrote. is this expected? If it is, better not to call set_param_memref_size if memref is Some to avoid confusion.

Also, I wonder whether RmemOutput is correct here. According to the old study (the deleted comment), RPC does use Tmem for this purpose. Is this difference because we are not doing real RPC yet? In that case, better to document it.

[Praveen K Paladugu (praveen-pk)]

https://optee.readthedocs.io/en/latest/architecture/trusted_applications.html#loading-ree-fs-ta is the protocol for LOAD_TA.
From what I have traced in optee_os and optee_client:

First LOAD_TA

send LOAD_TA with memref_size set to 0.
If memref_size it set to 0, tee-supplicant will return the size of the TA in bytes.

SHM_ALLOC

Get REE to allocate some memory to load the TA.

Second LOAD_TA

Register the TMEM into a shm and invoke a Second LOAD_TA. REE will actually load the TA into the provided shm in this call.

While tracing the first load_ta, the code branches of based on the input memref_size of 0 without checking the type. Further digging I noticed that the argument should be Tmem. I will fix in the next update.

If it is, better not to call set_param_memref_size if memref is Some to avoid confusion.

As I mentioned above code handles the first and second LOAD_TA calls as described above. I do see that memref_size does get clobbers if memref is valid (during second LOAD_TA). I will fix this.