Skip to content

Commit 96e9621

Browse files
sangho2Sangho Lee
sangho2
and
Sangho Lee
authored
Add system PTA map_zi and unmap (#939)
This PR adds the system PTA's `map_zi` and `unmap` functions. --------- Co-authored-by: Sangho Lee <[email protected]>
1 parent c597b7d commit 96e9621

6 files changed

Lines changed: 219 additions & 40 deletions

File tree

‎litebox_common_optee/src/lib.rs‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ pub enum SyscallRequest<Platform: litebox::platform::RawPointerProvider> {
8080
ta_sess_id: u32,
8181
cancel_req_to: u32,
8282
cmd_id: u32,
83-
params: Platform::RawConstPointer<UteeParams>,
83+
params: Platform::RawMutPointer<UteeParams>,
8484
ret_orig: Platform::RawMutPointer<TeeOrigin>,
8585
},
8686
CheckAccessRights {
@@ -194,7 +194,7 @@ impl<Platform: litebox::platform::RawPointerProvider> SyscallRequest<Platform> {
194194
ta_sess_id: u32::try_from(ctx.syscall_arg(0)).map_err(|_| Errno::EINVAL)?,
195195
cancel_req_to: u32::try_from(ctx.syscall_arg(1)).map_err(|_| Errno::EINVAL)?,
196196
cmd_id: u32::try_from(ctx.syscall_arg(2)).map_err(|_| Errno::EINVAL)?,
197-
params: Platform::RawConstPointer::from_usize(ctx.syscall_arg(3)),
197+
params: Platform::RawMutPointer::from_usize(ctx.syscall_arg(3)),
198198
ret_orig: Platform::RawMutPointer::from_usize(ctx.syscall_arg(4)),
199199
},
200200
TeeSyscallNr::CheckAccessRights => SyscallRequest::CheckAccessRights {
@@ -510,6 +510,18 @@ impl UteeParams {
510510
(0..Self::TEE_NUM_PARAMS).all(|i| self.get_type(i).is_ok_and(|t| t == expected[i]))
511511
}
512512

513+
/// Return `true` if any parameter is an output or inout type, i.e., the
514+
/// command may write results that must be copied back to the caller.
515+
pub fn needs_copy_back(&self) -> bool {
516+
use TeeParamType::{MemrefInout, MemrefOutput, ValueInout, ValueOutput};
517+
(0..Self::TEE_NUM_PARAMS).any(|i| {
518+
matches!(
519+
self.get_type(i),
520+
Ok(ValueOutput | ValueInout | MemrefOutput | MemrefInout)
521+
)
522+
})
523+
}
524+
513525
pub fn get_type(&self, index: usize) -> Result<TeeParamType, Errno> {
514526
let type_byte = match index {
515527
0 => self.types.type_0(),

‎litebox_shim_optee/src/lib.rs‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -469,8 +469,24 @@ impl Task {
469469
params,
470470
ret_orig,
471471
} => {
472-
if let Some(params) = params.read_at_offset(0) {
473-
self.sys_invoke_ta_command(ta_sess_id, cancel_req_to, cmd_id, params, ret_orig)
472+
if let Some(mut params_copied) = params.read_at_offset(0) {
473+
self.sys_invoke_ta_command(
474+
ta_sess_id,
475+
cancel_req_to,
476+
cmd_id,
477+
&mut params_copied,
478+
ret_orig,
479+
)
480+
.and_then(|cleanup| {
481+
if !params_copied.needs_copy_back()
482+
|| params.write_at_offset(0, params_copied).is_some()
483+
{
484+
Ok(())
485+
} else {
486+
cleanup.run(self);
487+
Err(TeeResult::AccessDenied)
488+
}
489+
})
474490
} else {
475491
Err(TeeResult::BadParameters)
476492
}
@@ -664,7 +680,19 @@ impl Task {
664680
pad_begin,
665681
pad_end,
666682
flags,
667-
} => self.sys_map_zi(va, num_bytes, pad_begin, pad_end, flags),
683+
} => match va.read_at_offset(0) {
684+
Some(hint) => self
685+
.sys_map_zi(hint, num_bytes, pad_begin, pad_end, flags)
686+
.and_then(|(mapped, cleanup)| {
687+
if va.write_at_offset(0, mapped).is_some() {
688+
Ok(())
689+
} else {
690+
cleanup.run(self);
691+
Err(TeeResult::AccessDenied)
692+
}
693+
}),
694+
None => Err(TeeResult::BadParameters),
695+
},
668696
LdelfSyscallRequest::OpenBin {
669697
uuid,
670698
uuid_size,

‎litebox_shim_optee/src/syscalls/ldelf.rs‎

Lines changed: 38 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Copyright (c) Microsoft Corporation.
22
// Licensed under the MIT license.
33

4+
use crate::syscalls::Cleanup;
45
use crate::{Task, UserMutPtr};
56
use litebox::mm::linux::PAGE_SIZE;
67
use litebox::platform::{RawConstPointer, RawMutPointer, SystemInfoProvider as _};
@@ -68,28 +69,27 @@ impl Task {
6869
}
6970

7071
/// OP-TEE's syscall to map zero-initialized memory with padding.
71-
/// This function pads `pad_begin` bytes before and `pad_end` bytes after the
72-
/// zero-initialized `num_bytes` bytes. `va` can contain a hint address which
73-
/// is `pad_begin` bytes lower than the starting address of the memory region.
74-
/// (`start - pad_begin`, ..., `start`, ..., `start + num_bytes`, ..., `start + num_bytes + pad_end`)
75-
/// Memory regions between `start - pad_begin` and `start` and between
76-
/// `start + num_bytes` and `start + num_bytes + pad_end` are reserved and must not be used.
72+
///
73+
/// Maps `pad_begin + num_bytes + pad_end` bytes (rounded up to a page) and
74+
/// zero-initializes the `num_bytes` usable region. `va` is a page-aligned
75+
/// hint for the *base of the whole mapping* (`0` means no hint). The usable
76+
/// region thus starts at `start = va + pad_begin`; the `pad_begin`/`pad_end`
77+
/// regions are reserved and must not be accessed.
78+
///
79+
/// On success, returns `start` plus a `Cleanup` that unmaps the usable
80+
/// region. The caller communicates the address back to userspace and must
81+
/// run the cleanup if that write-back fails.
7782
pub fn sys_map_zi(
7883
&self,
79-
va: UserMutPtr<usize>,
84+
va: usize,
8085
num_bytes: usize,
8186
pad_begin: usize,
8287
pad_end: usize,
8388
flags: LdelfMapFlags,
84-
) -> Result<(), TeeResult> {
85-
let Some(addr) = va.read_at_offset(0) else {
86-
return Err(TeeResult::BadParameters);
87-
};
88-
89+
) -> Result<(usize, Cleanup), TeeResult> {
8990
#[cfg(debug_assertions)]
9091
litebox_util_log::debug!(
91-
va:% = format_args!("{:#x}", va.as_usize()),
92-
addr:% = format_args!("{:#x}", addr),
92+
va:% = format_args!("{:#x}", va),
9393
num_bytes:% = num_bytes,
9494
flags:% = format_args!("{:#x}", flags);
9595
"sys_map_zi"
@@ -101,20 +101,28 @@ impl Task {
101101
}
102102
// TODO: Check whether flags contains `LDELF_MAP_FLAG_SHAREABLE` once we support sharing of file-based mappings.
103103

104+
// OP-TEE requires the address hint and padding to be page-aligned.
105+
if !va.is_multiple_of(PAGE_SIZE)
106+
|| !pad_begin.is_multiple_of(PAGE_SIZE)
107+
|| !pad_end.is_multiple_of(PAGE_SIZE)
108+
{
109+
return Err(TeeResult::AccessConflict);
110+
}
111+
104112
let total_size = Self::checked_map_size(num_bytes, pad_begin, pad_end)?;
105-
if addr.checked_add(total_size).is_none() {
113+
if va.checked_add(total_size).is_none() {
106114
return Err(TeeResult::BadParameters);
107115
}
108116
// `sys_map_zi` always creates read/writeable mapping.
109117
//
110118
// We map with PROT_READ_WRITE first, then mprotect padding regions to PROT_NONE.
111119
let mut flags = MapFlags::MAP_PRIVATE | MapFlags::MAP_ANONYMOUS;
112-
if addr != 0 {
120+
if va != 0 {
113121
flags |= MapFlags::MAP_FIXED;
114122
}
115123

116124
let addr = self
117-
.sys_mmap(addr, total_size, ProtFlags::PROT_READ_WRITE, flags, -1, 0)
125+
.sys_mmap(va, total_size, ProtFlags::PROT_READ_WRITE, flags, -1, 0)
118126
.map_err(|_| TeeResult::OutOfMemory)?;
119127
let guard = MmapGuard::new(self, addr, total_size);
120128

@@ -143,9 +151,12 @@ impl Task {
143151
);
144152
}
145153

146-
let _ = va.write_at_offset(0, padded_start);
147154
guard.disarm();
148-
Ok(())
155+
let cleanup = Cleanup::Unmap {
156+
addr: padded_start,
157+
len: pad_end_start - padded_start,
158+
};
159+
Ok((padded_start, cleanup))
149160
}
150161

151162
/// OP-TEE's syscall to open a TA binary.
@@ -215,6 +226,14 @@ impl Task {
215226
return Err(TeeResult::BadParameters);
216227
}
217228

229+
// OP-TEE requires the address hint and padding to be page-aligned.
230+
if !addr.is_multiple_of(PAGE_SIZE)
231+
|| !pad_begin.is_multiple_of(PAGE_SIZE)
232+
|| !pad_end.is_multiple_of(PAGE_SIZE)
233+
{
234+
return Err(TeeResult::AccessConflict);
235+
}
236+
218237
if self.ta_handle_map.get(handle).is_none() {
219238
return Err(TeeResult::BadParameters);
220239
}

‎litebox_shim_optee/src/syscalls/mod.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33

44
//! Syscalls Handlers
55
6+
use crate::{Task, UserMutPtr};
7+
use litebox::platform::RawConstPointer as _;
8+
69
pub(crate) mod cryp;
710
pub(crate) mod ldelf;
811
pub(crate) mod mm;
@@ -11,3 +14,26 @@ pub(crate) mod tee;
1114

1215
#[cfg(test)]
1316
pub(crate) mod tests;
17+
18+
/// Undo a syscall/command's side effects if dispatch fails after the command
19+
/// succeeded (currently only when copying results back out to the guest fails).
20+
#[derive(Default)]
21+
#[must_use = "must be run when dispatch fails after the command, or the side effect leaks"]
22+
pub(crate) enum Cleanup {
23+
#[default]
24+
None,
25+
/// Unmap a region. `addr` must be page-aligned; `len` is rounded up by `sys_munmap`.
26+
Unmap { addr: usize, len: usize },
27+
}
28+
29+
impl Cleanup {
30+
/// Undo the side effect. Runs only on an error path, so failures are ignored.
31+
pub(crate) fn run(self, task: &Task) {
32+
match self {
33+
Self::None => {}
34+
Self::Unmap { addr, len } => {
35+
let _ = task.sys_munmap(UserMutPtr::<u8>::from_usize(addr), len);
36+
}
37+
}
38+
}
39+
}

‎litebox_shim_optee/src/syscalls/pta.rs‎

Lines changed: 93 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,19 @@
44
//! Implementation of pseudo TAs (PTAs) which export system services as
55
//! the functions of built-in TAs.
66
7+
use crate::syscalls::Cleanup;
78
use crate::{Task, UserConstPtr, UserMutPtr};
89
use alloc::vec;
910
use alloc::vec::Vec;
1011
use hmac::{Hmac, Mac};
12+
use litebox::mm::linux::PAGE_SIZE;
1113
use litebox::platform::{
1214
DerivedKeyError, DerivedKeyProvider, KDFParams, RawConstPointer as _, RawMutPointer as _,
1315
};
1416
use litebox::utils::TruncateExt;
1517
use litebox_common_optee::{
16-
HUK_SUBKEY_MAX_LEN, HukSubkeyUsage, TaFlags, TeeParamType, TeeResult, TeeUuid, UteeParams,
18+
HUK_SUBKEY_MAX_LEN, HukSubkeyUsage, LdelfMapFlags, TaFlags, TeeParamType, TeeResult, TeeUuid,
19+
UteeParams,
1720
};
1821
use num_enum::TryFromPrimitive;
1922
use sha2::Sha256;
@@ -48,8 +51,8 @@ impl PseudoTa {
4851
self,
4952
task: &Task,
5053
cmd_id: u32,
51-
params: &UteeParams,
52-
) -> Result<(), TeeResult> {
54+
params: &mut UteeParams,
55+
) -> Result<Cleanup, TeeResult> {
5356
let _busy = task.try_set_busy(self)?;
5457
match self {
5558
Self::System => SystemPta::invoke_command(task, cmd_id, params),
@@ -235,10 +238,19 @@ impl SystemPta {
235238
}
236239

237240
/// Handle a command of the system PTA.
238-
fn invoke_command(task: &Task, cmd_id: u32, params: &UteeParams) -> Result<(), TeeResult> {
239-
#[allow(clippy::single_match_else)]
241+
///
242+
/// See `Cleanup` for the returned rollback; most commands have no cleanup.
243+
fn invoke_command(
244+
task: &Task,
245+
cmd_id: u32,
246+
params: &mut UteeParams,
247+
) -> Result<Cleanup, TeeResult> {
240248
match PtaSystemCommandId::try_from(cmd_id).map_err(|_| TeeResult::BadParameters)? {
241-
PtaSystemCommandId::DeriveTaUniqueKey => Self::derive_ta_unique_key(task, params),
249+
PtaSystemCommandId::DeriveTaUniqueKey => {
250+
Self::derive_ta_unique_key(task, params).map(|()| Cleanup::None)
251+
}
252+
PtaSystemCommandId::MapZi => Self::map_zi(task, params),
253+
PtaSystemCommandId::Unmap => Self::unmap(task, params).map(|()| Cleanup::None),
242254
_ => {
243255
#[cfg(debug_assertions)]
244256
todo!("support other system PTA commands {cmd_id}");
@@ -346,6 +358,81 @@ impl SystemPta {
346358

347359
Ok(())
348360
}
361+
362+
fn map_zi(task: &Task, params: &mut UteeParams) -> Result<Cleanup, TeeResult> {
363+
use TeeParamType::{None, ValueInout, ValueInput};
364+
365+
if !params.has_types([ValueInput, ValueInout, ValueInput, None]) {
366+
return Err(TeeResult::BadParameters);
367+
}
368+
369+
let (num_bytes, flags) = params
370+
.get_values(0)
371+
.map_err(|_| TeeResult::BadParameters)?
372+
.ok_or(TeeResult::BadParameters)?;
373+
if num_bytes == 0 {
374+
return Err(TeeResult::BadParameters);
375+
}
376+
let (addr_high, addr_low) = params
377+
.get_values(1)
378+
.map_err(|_| TeeResult::BadParameters)?
379+
.ok_or(TeeResult::BadParameters)?;
380+
let (pad_begin, pad_end) = params
381+
.get_values(2)
382+
.map_err(|_| TeeResult::BadParameters)?
383+
.ok_or(TeeResult::BadParameters)?;
384+
385+
if addr_high & 0xffff_ffff_0000_0000 != 0 || addr_low & 0xffff_ffff_0000_0000 != 0 {
386+
return Err(TeeResult::BadParameters);
387+
}
388+
let addr: usize = ((addr_high << 32) | addr_low).trunc();
389+
let (mapped, cleanup) = task.sys_map_zi(
390+
addr,
391+
num_bytes.trunc(),
392+
pad_begin.trunc(),
393+
pad_end.trunc(),
394+
LdelfMapFlags::from_bits_retain(flags.trunc()),
395+
)?;
396+
397+
// Return the mapped address to the caller via the inout value param.
398+
// This `set_values` cannot fail because the index is fixed/known.
399+
let _ = params.set_values(1, (mapped as u64) >> 32, (mapped as u64) & 0xffff_ffff);
400+
401+
// The caller runs `cleanup` (unmap) if it encounters an error.
402+
Ok(cleanup)
403+
}
404+
405+
fn unmap(task: &Task, params: &UteeParams) -> Result<(), TeeResult> {
406+
use TeeParamType::{None, ValueInput};
407+
408+
if !params.has_types([ValueInput, ValueInput, None, None]) {
409+
return Err(TeeResult::BadParameters);
410+
}
411+
412+
let (size, must_be_zero) = params
413+
.get_values(0)
414+
.map_err(|_| TeeResult::BadParameters)?
415+
.ok_or(TeeResult::BadParameters)?;
416+
if must_be_zero != 0 {
417+
return Err(TeeResult::BadParameters);
418+
}
419+
let (addr_high, addr_low) = params
420+
.get_values(1)
421+
.map_err(|_| TeeResult::BadParameters)?
422+
.ok_or(TeeResult::BadParameters)?;
423+
424+
if addr_high & 0xffff_ffff_0000_0000 != 0 || addr_low & 0xffff_ffff_0000_0000 != 0 {
425+
return Err(TeeResult::BadParameters);
426+
}
427+
let addr: usize = ((addr_high << 32) | addr_low).trunc();
428+
let size: usize = size.trunc();
429+
let size = size
430+
.checked_next_multiple_of(PAGE_SIZE)
431+
.ok_or(TeeResult::BadParameters)?;
432+
433+
task.sys_munmap(UserMutPtr::<u8>::from_usize(addr), size)
434+
.map_err(|_| TeeResult::BadParameters)
435+
}
349436
}
350437

351438
/// A KDF callback that derives a subkey from `huk` and `params.context` to be passed to

0 commit comments

Comments
 (0)