From dc3748e1b9130f3da75927114543f53e60cf7f91 Mon Sep 17 00:00:00 2001
From: ChanTsune <41658782+ChanTsune@users.noreply.github.com>
Date: Tue, 10 Dec 2024 18:47:53 +0900
Subject: [PATCH] :bug: Deallocations violate GlobalAlloc's safety contract

https://github.com/gyscos/zstd-rs/issues/248

https://github.com/Portable-Network-Archive/liblzma-rs/issues/130
---
 liblzma-sys/src/wasm_shim.rs | 60 ++++++++++++++++++++++++++++--------
 1 file changed, 47 insertions(+), 13 deletions(-)

diff --git a/liblzma-sys/src/wasm_shim.rs b/liblzma-sys/src/wasm_shim.rs
index f2d678d3..8cb890fc 100644
--- a/liblzma-sys/src/wasm_shim.rs
+++ b/liblzma-sys/src/wasm_shim.rs
@@ -1,30 +1,23 @@
 use core::ffi::{c_char, c_int, c_void};
-use std::alloc::{alloc, dealloc, Layout};
+use std::alloc::{alloc, alloc_zeroed, dealloc, Layout};
 
 #[no_mangle]
 pub extern "C" fn rust_lzma_wasm_shim_malloc(size: usize) -> *mut c_void {
-    unsafe {
-        let layout = Layout::from_size_align_unchecked(size, 1);
-        alloc(layout).cast()
-    }
+    wasm_shim_alloc::<false>(size)
 }
 
 #[no_mangle]
 pub extern "C" fn rust_lzma_wasm_shim_calloc(nmemb: usize, size: usize) -> *mut c_void {
-    unsafe {
-        let layout = Layout::from_size_align_unchecked(size * nmemb, 1);
-        alloc(layout).cast()
-    }
+    // note: calloc expects the allocation to be zeroed
+    wasm_shim_alloc::<true>(nmemb * size)
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn rust_lzma_wasm_shim_free(ptr: *mut c_void) {
-    if ptr == std::ptr::null_mut() {
+    if ptr.is_null() {
         return;
     }
-    // layout is not actually used
-    let layout = Layout::from_size_align_unchecked(1, 1);
-    dealloc(ptr.cast(), layout);
+    wasm_shim_free(ptr)
 }
 
 #[no_mangle]
@@ -95,3 +88,44 @@ pub unsafe extern "C" fn rust_lzma_wasm_shim_memchr(
             s.add(p) as *mut c_void
         })
 }
+
+const USIZE_ALIGN: usize = core::mem::align_of::<usize>();
+const USIZE_SIZE: usize = core::mem::size_of::<usize>();
+
+#[inline]
+fn wasm_shim_alloc<const ZEROED: bool>(size: usize) -> *mut c_void {
+    // in order to recover the size upon free, we store the size below the allocation
+    // special alignment is never requested via the malloc API,
+    // so it's not stored, and usize-alignment is used
+    // memory layout: [size] [allocation]
+
+    let full_alloc_size = size + USIZE_SIZE;
+
+    unsafe {
+        let layout = Layout::from_size_align_unchecked(full_alloc_size, USIZE_ALIGN);
+
+        let ptr = if ZEROED {
+            alloc_zeroed(layout)
+        } else {
+            alloc(layout)
+        };
+
+        // SAFETY: ptr is usize-aligned and we've allocated sufficient memory
+        ptr.cast::<usize>().write(full_alloc_size);
+
+        ptr.add(USIZE_SIZE).cast()
+    }
+}
+
+unsafe fn wasm_shim_free(ptr: *mut c_void) {
+    // the layout for the allocation needs to be recovered for dealloc
+    // - the size must be recovered from directly below the allocation
+    // - the alignment will always by USIZE_ALIGN
+
+    let alloc_ptr = ptr.sub(USIZE_SIZE);
+    // SAFETY: the allocation routines must uphold having a valid usize below the provided pointer
+    let full_alloc_size = alloc_ptr.cast::<usize>().read();
+
+    let layout = Layout::from_size_align_unchecked(full_alloc_size, USIZE_ALIGN);
+    dealloc(alloc_ptr.cast(), layout);
+}