Use checked arithmetic against TA/ldelf-controllable integers

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<usize, Self::Error> {
         // 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::<u8>::from_usize(mapped_addr);
             user_ptr

litebox_shim_optee/src/msg_handler.rs

@@ -388,10 +388,6 @@ pub struct TaRequestInfo<const ALIGN: usize> {
 /// 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<TaRequestInfo<PAGE_SIZE>, 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()

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<T: Clone, const ALIGN: usize> PhysMutPtr<T, ALIGN> {
             ));
         }
         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::<ALIGN>::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<T: Clone, const ALIGN: usize> PhysMutPtr<T, ALIGN> {
             )
             .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)?;
         }

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<u8>,
+    len: usize,
+}
+
+impl<'a> MmapGuard<'a> {
+    fn new(task: &'a Task, addr: UserMutPtr<u8>, 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<usize, TeeResult> {
+        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<usize, TeeResult> {
+        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(())
     }

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<usize> {
     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);
         }

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<usize> {
     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::<PAGE_SIZE>::new(align_up(len, PAGE_SIZE))
-                .ok_or(TeeResult::AccessConflict)?
+            NonZeroPageSize::<PAGE_SIZE>::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))