diff --git a/litebox_shim_optee/src/loader/elf.rs b/litebox_shim_optee/src/loader/elf.rs index 2cfe83167..fcbd7145c 100644 --- a/litebox_shim_optee/src/loader/elf.rs +++ b/litebox_shim_optee/src/loader/elf.rs @@ -45,7 +45,8 @@ fn read_at(elf: &ElfFileInMemory, offset: u64, buf: &mut [u8]) -> Result<(), Err if offset >= elf.buffer.len() { return Err(Errno::ENODATA); } - let end = core::cmp::min(offset + buf.len(), elf.buffer.len()); + let available = elf.buffer.len() - offset; + let end = offset + core::cmp::min(buf.len(), available); let len = end - offset; buf[..len].copy_from_slice(&elf.buffer[offset..end]); Ok(()) @@ -78,7 +79,9 @@ impl litebox_common_linux::loader::MapMemory for ElfFileInMemory<'_> { fn reserve(&mut self, len: usize, align: usize) -> Result { // Allocate a mapping large enough that even if it's maximally misaligned we can // still fit `len` bytes. - let mapping_len = len + (align.max(PAGE_SIZE) - PAGE_SIZE); + let mapping_len = len + .checked_add(align.max(PAGE_SIZE) - PAGE_SIZE) + .ok_or(Errno::ENOMEM)?; let mapping_ptr = self .task .sys_mmap( @@ -91,9 +94,14 @@ impl litebox_common_linux::loader::MapMemory for ElfFileInMemory<'_> { )? .as_usize(); - let ptr = mapping_ptr.next_multiple_of(align); - let end = ptr + len; - let mapping_end = mapping_ptr + mapping_len; + let ptr = mapping_ptr + .checked_next_multiple_of(align) + .ok_or(Errno::ENOMEM)?; + let end = ptr.checked_add(len).ok_or(Errno::ENOMEM)?; + let mapping_end = mapping_ptr.checked_add(mapping_len).ok_or(Errno::ENOMEM)?; + if end > mapping_end { + return Err(Errno::ENOMEM); + } if ptr != mapping_ptr { self.task .sys_munmap(MutPtr::from_usize(mapping_ptr), ptr - mapping_ptr)?; @@ -136,7 +144,8 @@ impl litebox_common_linux::loader::MapMemory for ElfFileInMemory<'_> { // MAP_ANONYMOUS ensures remaining bytes are zero if src is shorter than len. let offset: usize = offset.trunc(); if len > 0 && offset < self.buffer.len() { - let end = core::cmp::min(offset + len, self.buffer.len()); + let available = self.buffer.len() - offset; + let end = offset + core::cmp::min(len, available); let src = &self.buffer[offset..end]; let user_ptr = UserMutPtr::::from_usize(mapped_addr); user_ptr diff --git a/litebox_shim_optee/src/msg_handler.rs b/litebox_shim_optee/src/msg_handler.rs index b1c125101..878ca6749 100644 --- a/litebox_shim_optee/src/msg_handler.rs +++ b/litebox_shim_optee/src/msg_handler.rs @@ -388,10 +388,6 @@ pub struct TaRequestInfo { /// It copies the entire parameter data from the normal world shared memory into the secure world's /// memory to create `UteeParamOwned` structures to avoid potential data corruption during TA /// execution. -/// -/// # Panics -/// -/// Panics if any conversion from `u64` to `usize` fails. OP-TEE shim doesn't support a 32-bit environment. pub fn decode_ta_request( msg_args: &OpteeMsgArgs, ) -> Result, OpteeSmcReturnCode> { @@ -445,6 +441,12 @@ pub fn decode_ta_request( }; let num_params = msg_args.num_params as usize; + if num_params + .checked_sub(skip) + .is_none_or(|n| n > UteeParamOwned::TEE_NUM_PARAMS) + { + return Err(OpteeSmcReturnCode::EBadCmd); + } for (i, param) in msg_args .params .iter() diff --git a/litebox_shim_optee/src/ptr.rs b/litebox_shim_optee/src/ptr.rs index 2a54bac25..06a492506 100644 --- a/litebox_shim_optee/src/ptr.rs +++ b/litebox_shim_optee/src/ptr.rs @@ -94,11 +94,6 @@ fn align_down(address: usize, align: usize) -> usize { address & !(align - 1) } -#[inline] -fn align_up(len: usize, align: usize) -> usize { - len.next_multiple_of(align) -} - /// Represent a physical pointer to an object with on-demand mapping. /// - `pages`: An array of page-aligned physical addresses. We expect physical addresses in this array are /// virtually contiguous. @@ -167,18 +162,23 @@ impl PhysMutPtr { )); } let start_page = align_down(pa, ALIGN); - let end_page = align_up( - pa.checked_add(bytes).ok_or(PhysPointerError::Overflow)?, - ALIGN, - ); - let mut pages = alloc::vec::Vec::with_capacity((end_page - start_page) / ALIGN); + let end_page = pa + .checked_add(bytes) + .and_then(|end| end.checked_next_multiple_of(ALIGN)) + .ok_or(PhysPointerError::Overflow)?; + let span = end_page + .checked_sub(start_page) + .ok_or(PhysPointerError::Overflow)?; + let mut pages = alloc::vec::Vec::with_capacity(span / ALIGN); let mut current_page = start_page; while current_page < end_page { pages.push( PhysPageAddr::::new(current_page) .ok_or(PhysPointerError::InvalidPhysicalAddress(current_page))?, ); - current_page += ALIGN; + current_page = current_page + .checked_add(ALIGN) + .ok_or(PhysPointerError::Overflow)?; } Self::new(&pages, pa - start_page) } @@ -377,7 +377,10 @@ impl PhysMutPtr { ) .ok_or(PhysPointerError::Overflow)?; let start = skip / ALIGN; - let end = (skip + size).div_ceil(ALIGN); + let end = skip + .checked_add(size) + .ok_or(PhysPointerError::Overflow)? + .div_ceil(ALIGN); unsafe { self.map_range(start, end, perms)?; } diff --git a/litebox_shim_optee/src/syscalls/ldelf.rs b/litebox_shim_optee/src/syscalls/ldelf.rs index 398a1464e..d9f26f6d0 100644 --- a/litebox_shim_optee/src/syscalls/ldelf.rs +++ b/litebox_shim_optee/src/syscalls/ldelf.rs @@ -13,7 +13,57 @@ fn align_down(addr: usize, align: usize) -> usize { addr & !(align - 1) } +/// Calls `sys_munmap(addr, len)` when dropped, unless `disarm()` has been called first. +/// +/// Used to ensure a mapping created by `sys_mmap` is released on every error +/// path of `sys_map_zi` / `sys_map_bin`. After the syscall has fully succeeded +/// and ownership of the mapping has been transferred to the caller, call +/// `disarm()` to suppress the unmap. +#[must_use = "MmapGuard unmaps on drop unless disarm() is called; bind it"] +struct MmapGuard<'a> { + task: &'a Task, + addr: UserMutPtr, + len: usize, +} + +impl<'a> MmapGuard<'a> { + fn new(task: &'a Task, addr: UserMutPtr, len: usize) -> Self { + Self { task, addr, len } + } + + fn disarm(self) { + core::mem::forget(self); + } +} + +impl Drop for MmapGuard<'_> { + fn drop(&mut self) { + let _ = self.task.sys_munmap(self.addr, self.len); + } +} + impl Task { + #[inline] + fn checked_map_size( + num_bytes: usize, + pad_begin: usize, + pad_end: usize, + ) -> Result { + num_bytes + .checked_add(pad_begin) + .and_then(|t| t.checked_add(pad_end)) + .and_then(|t| t.checked_next_multiple_of(PAGE_SIZE)) + .ok_or(TeeResult::BadParameters) + } + + #[inline] + fn checked_pad_end_start(padded_start: usize, num_bytes: usize) -> Result { + padded_start + .checked_add(num_bytes) + .and_then(|end| end.checked_next_multiple_of(PAGE_SIZE)) + .ok_or(TeeResult::BadParameters) + } + /// OP-TEE's syscall to map zero-initialized memory with padding. /// This function pads `pad_begin` bytes before and `pad_end` bytes after the /// zero-initialized `num_bytes` bytes. `va` can contain a hint address which @@ -48,11 +98,7 @@ impl Task { } // TODO: Check whether flags contains `LDELF_MAP_FLAG_SHAREABLE` once we support sharing of file-based mappings. - let total_size = num_bytes - .checked_add(pad_begin) - .and_then(|t| t.checked_add(pad_end)) - .ok_or(TeeResult::BadParameters)? - .next_multiple_of(PAGE_SIZE); + let total_size = Self::checked_map_size(num_bytes, pad_begin, pad_end)?; if addr.checked_add(total_size).is_none() { return Err(TeeResult::BadParameters); } @@ -67,7 +113,12 @@ impl Task { let addr = self .sys_mmap(addr, total_size, ProtFlags::PROT_READ_WRITE, flags, -1, 0) .map_err(|_| TeeResult::OutOfMemory)?; - let padded_start = addr.as_usize() + pad_begin; + let guard = MmapGuard::new(self, addr, total_size); + + let padded_start = addr + .as_usize() + .checked_add(pad_begin) + .ok_or(TeeResult::BadParameters)?; // Unmap the padding regions to free physical memory. // Using munmap instead of mprotect(PROT_NONE) actually deallocates the frames. @@ -77,8 +128,11 @@ impl Task { let _ = self.sys_munmap(addr, pad_begin_end - addr.as_usize()); } // pad_end region: [align_up(padded_start + num_bytes, PAGE_SIZE), addr + total_size) - let pad_end_start = (padded_start + num_bytes).next_multiple_of(PAGE_SIZE); - let region_end = addr.as_usize() + total_size; + let pad_end_start = Self::checked_pad_end_start(padded_start, num_bytes)?; + let region_end = addr + .as_usize() + .checked_add(total_size) + .ok_or(TeeResult::BadParameters)?; if pad_end_start < region_end { let _ = self.sys_munmap( UserMutPtr::from_usize(pad_end_start), @@ -87,6 +141,7 @@ impl Task { } let _ = va.write_at_offset(0, padded_start); + guard.disarm(); Ok(()) } @@ -172,11 +227,7 @@ impl Task { return Err(TeeResult::BadParameters); } - let total_size = num_bytes - .checked_add(pad_begin) - .and_then(|t| t.checked_add(pad_end)) - .ok_or(TeeResult::BadParameters)? - .next_multiple_of(PAGE_SIZE); + let total_size = Self::checked_map_size(num_bytes, pad_begin, pad_end)?; if addr.checked_add(total_size).is_none() { return Err(TeeResult::BadParameters); } @@ -225,9 +276,13 @@ impl Task { 0, ) .map_err(|_| TeeResult::OutOfMemory)?; - let padded_start = addr.as_usize() + pad_begin; + let guard = MmapGuard::new(self, addr, mmap_size); + + let padded_start = addr + .as_usize() + .checked_add(pad_begin) + .ok_or(TeeResult::BadParameters)?; if padded_start == 0 { - let _ = self.sys_munmap(addr, total_size).ok(); return Err(TeeResult::BadFormat); } @@ -250,16 +305,16 @@ impl Task { } else if flags.contains(LdelfMapFlags::LDELF_MAP_FLAG_EXECUTABLE) { prot |= ProtFlags::PROT_EXEC; } + let prot_start = align_down(padded_start, PAGE_SIZE); + let prot_len = padded_start + .checked_sub(prot_start) + .and_then(|offset| offset.checked_add(num_bytes)) + .and_then(|len| len.checked_next_multiple_of(PAGE_SIZE)) + .ok_or(TeeResult::BadParameters)?; if self - .sys_mprotect( - UserMutPtr::from_usize(align_down(padded_start, PAGE_SIZE)), - (num_bytes + padded_start - align_down(padded_start, PAGE_SIZE)) - .next_multiple_of(PAGE_SIZE), - prot, - ) + .sys_mprotect(UserMutPtr::from_usize(prot_start), prot_len, prot) .is_err() { - let _ = self.sys_munmap(addr, total_size).ok(); return Err(TeeResult::AccessDenied); } @@ -271,8 +326,11 @@ impl Task { let _ = self.sys_munmap(addr, pad_begin_end - addr.as_usize()); } // pad_end region: [align_up(padded_start + num_bytes, PAGE_SIZE), addr + total_size) - let pad_end_start = (padded_start + num_bytes).next_multiple_of(PAGE_SIZE); - let region_end = addr.as_usize() + total_size; + let pad_end_start = Self::checked_pad_end_start(padded_start, num_bytes)?; + let region_end = addr + .as_usize() + .checked_add(total_size) + .ok_or(TeeResult::BadParameters)?; if pad_end_start < region_end { let _ = self.sys_munmap( UserMutPtr::from_usize(pad_end_start), @@ -281,6 +339,7 @@ impl Task { } let _ = va.write_at_offset(0, padded_start); + guard.disarm(); Ok(()) } diff --git a/litebox_shim_optee/src/syscalls/mm.rs b/litebox_shim_optee/src/syscalls/mm.rs index 5bd3fa90a..116421b8f 100644 --- a/litebox_shim_optee/src/syscalls/mm.rs +++ b/litebox_shim_optee/src/syscalls/mm.rs @@ -9,9 +9,9 @@ use litebox_common_linux::{MapFlags, ProtFlags, errno::Errno}; use crate::{Task, UserMutPtr}; #[inline] -fn align_up(addr: usize, align: usize) -> usize { +fn align_up(addr: usize, align: usize) -> Option { debug_assert!(align.is_power_of_two()); - (addr + align - 1) & !(align - 1) + addr.checked_next_multiple_of(align) } impl Task { @@ -66,7 +66,7 @@ impl Task { return Err(Errno::EINVAL); } - let aligned_len = align_up(len, PAGE_SIZE); + let aligned_len = align_up(len, PAGE_SIZE).ok_or(Errno::ENOMEM)?; if aligned_len == 0 { return Err(Errno::ENOMEM); } diff --git a/litebox_shim_optee/src/syscalls/tee.rs b/litebox_shim_optee/src/syscalls/tee.rs index 4f605b50c..6ee19b6f7 100644 --- a/litebox_shim_optee/src/syscalls/tee.rs +++ b/litebox_shim_optee/src/syscalls/tee.rs @@ -21,9 +21,9 @@ use crate::{ }; #[inline] -fn align_up(addr: usize, align: usize) -> usize { +fn align_up(addr: usize, align: usize) -> Option { debug_assert!(align.is_power_of_two()); - (addr + align - 1) & !(align - 1) + addr.checked_next_multiple_of(align) } #[inline] @@ -255,9 +255,17 @@ impl Task { let len = len .checked_add(buf.as_usize() - align_down(buf.as_usize(), PAGE_SIZE)) .ok_or(TeeResult::AccessConflict)?; - NonZeroPageSize::::new(align_up(len, PAGE_SIZE)) - .ok_or(TeeResult::AccessConflict)? + NonZeroPageSize::::new( + align_up(len, PAGE_SIZE).ok_or(TeeResult::AccessConflict)?, + ) + .ok_or(TeeResult::AccessConflict)? }; + // Reject ranges where `start + aligned_len` would wrap, so downstream + // permission lookups don't operate on a truncated address range. + let _ = start + .as_usize() + .checked_add(aligned_len.as_usize()) + .ok_or(TeeResult::AccessConflict)?; if let Some(perms) = self.global.pm.get_memory_permissions(start, aligned_len) { if (flags.contains(TeeMemoryAccessRights::TEE_MEMORY_ACCESS_READ) && !perms.contains(MemoryRegionPermissions::READ))