From 92ce415317d9085611c4b2d12d46a322e25b05b3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20L=C3=B6bel?= <loebel.marvin@gmail.com>
Date: Wed, 21 Oct 2020 22:11:09 +0200
Subject: [PATCH] Prevent `try_zeroed_box<T>()` from reserving `size_of<T>()`
 space on the stack. (#43)

* Add test

* Change try_zeroed_box implementation to not allocate space for T on the stack

* Add second test
---
 src/allocation.rs  | 16 +++++++++++++++-
 tests/std_tests.rs | 16 ++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/allocation.rs b/src/allocation.rs
index 06d62fa..07de49f 100644
--- a/src/allocation.rs
+++ b/src/allocation.rs
@@ -11,6 +11,7 @@ use alloc::{
   boxed::Box,
   vec::Vec,
 };
+use core::convert::TryInto;
 
 /// As [`try_cast_box`](try_cast_box), but unwraps for you.
 #[inline]
@@ -55,7 +56,20 @@ pub fn try_cast_box<A: Pod, B: Pod>(
 #[inline]
 pub fn try_zeroed_box<T: Zeroable>() -> Result<Box<T>, ()> {
   if size_of::<T>() == 0 {
-    return Ok(Box::new(T::zeroed()));
+    // This will not allocate but simple create a dangling slice pointer.
+    // NB: We go the way via a push to `Vec<T>` to ensure the compiler
+    // does not allocate space for T on the stack even if the branch
+    // would not be taken.
+    let mut vec = Vec::with_capacity(1);
+    vec.resize_with(1, || T::zeroed());
+    let ptr: Box<[T; 1]> = vec.into_boxed_slice().try_into().ok().unwrap();
+    debug_assert!(
+      align_of::<[T; 1]>() == align_of::<T>()
+        && size_of::<[T; 1]>() == size_of::<T>()
+    );
+    // NB: We basically do the same as in try_cast_box here:
+    let ptr: Box<T> = unsafe { Box::from_raw(Box::into_raw(ptr) as *mut _) };
+    return Ok(ptr);
   }
   let layout =
     Layout::from_size_align(size_of::<T>(), align_of::<T>()).unwrap();
diff --git a/tests/std_tests.rs b/tests/std_tests.rs
index 0e0fb48..c8ea9e3 100644
--- a/tests/std_tests.rs
+++ b/tests/std_tests.rs
@@ -27,3 +27,19 @@ fn test_transparent_vtabled() {
   let s = format!("{}", v_mut);
   assert_eq!(s, "100");
 }
+
+#[test]
+#[cfg(feature = "extern_crate_alloc")]
+fn test_large_box_alloc() {
+  type SuperPage = [[u8; 4096]; 4096];
+  let _: Box<SuperPage> = try_zeroed_box().unwrap();
+}
+
+#[test]
+#[cfg(feature = "extern_crate_alloc")]
+fn test_zero_sized_box_alloc() {
+  #[repr(align(4096))]
+  struct Empty;
+  unsafe impl Zeroable for Empty {}
+  let _: Box<Empty> = try_zeroed_box().unwrap();
+}