From fd41ef659dcd9eaa973771985edce227cdcfe084 Mon Sep 17 00:00:00 2001
From: David Tolnay <dtolnay@gmail.com>
Date: Thu, 14 Dec 2023 14:09:29 -0800
Subject: [PATCH] Check arithmetic in malloc computations

---
 src/lib.rs | 15 ++++++++++-----
 src/ops.rs | 23 +++++++++++++++++++++++
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 6865d9c..e083433 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -54,6 +54,7 @@ mod libc {
 #[macro_use]
 mod externs {
     use crate::libc;
+    use crate::ops::{die, ForceAdd as _, ForceInto as _};
     use alloc::alloc::{self as rust, Layout};
     use core::mem::{self, MaybeUninit};
     use core::ptr;
@@ -66,8 +67,10 @@ mod externs {
     const MALLOC_ALIGN: usize = mem::align_of::<usize>();
 
     pub unsafe fn malloc(size: libc::c_ulong) -> *mut libc::c_void {
-        let size = HEADER + size as usize;
-        let layout = Layout::from_size_align_unchecked(size, MALLOC_ALIGN);
+        let size = HEADER.force_add(size.force_into());
+        let layout = Layout::from_size_align(size, MALLOC_ALIGN)
+            .ok()
+            .unwrap_or_else(die);
         let memory = rust::alloc(layout);
         if memory.is_null() {
             rust::handle_alloc_error(layout);
@@ -80,11 +83,13 @@ mod externs {
         let mut memory = ptr.cast::<u8>().sub(HEADER);
         let size = memory.cast::<usize>().read();
         let layout = Layout::from_size_align_unchecked(size, MALLOC_ALIGN);
-        let new_size = HEADER + new_size as usize;
+        let new_size = HEADER.force_add(new_size.force_into());
+        let new_layout = Layout::from_size_align(new_size, MALLOC_ALIGN)
+            .ok()
+            .unwrap_or_else(die);
         memory = rust::realloc(memory, layout, new_size);
         if memory.is_null() {
-            let layout = Layout::from_size_align_unchecked(new_size, MALLOC_ALIGN);
-            rust::handle_alloc_error(layout);
+            rust::handle_alloc_error(new_layout);
         }
         memory.cast::<usize>().write(new_size);
         memory.add(HEADER).cast()
diff --git a/src/ops.rs b/src/ops.rs
index a3b906c..4597d3d 100644
--- a/src/ops.rs
+++ b/src/ops.rs
@@ -26,6 +26,12 @@ impl ForceAdd for u64 {
     }
 }
 
+impl ForceAdd for usize {
+    fn force_add(self, rhs: Self) -> Self {
+        self.checked_add(rhs).unwrap_or_else(die)
+    }
+}
+
 pub(crate) trait ForceMul: Sized {
     fn force_mul(self, rhs: Self) -> Self;
 }
@@ -48,6 +54,23 @@ impl ForceMul for u64 {
     }
 }
 
+pub(crate) trait ForceInto {
+    fn force_into<U>(self) -> U
+    where
+        Self: TryInto<U>;
+}
+
+impl<T> ForceInto for T {
+    fn force_into<U>(self) -> U
+    where
+        Self: TryInto<U>,
+    {
+        <Self as TryInto<U>>::try_into(self)
+            .ok()
+            .unwrap_or_else(die)
+    }
+}
+
 // Deterministically abort on arithmetic overflow, instead of wrapping and
 // continuing with invalid behavior.
 //