From 290b3379f60c0578174861c6ac8fb93f1ce83a3f Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Tue, 17 Aug 2021 10:50:56 -0700
Subject: [PATCH] Adjust API to match __cxa_demangle

Bug: 178565008
Test: cargo test + generate tombstone
Change-Id: Iaf76af67a7f3a6323e926075a8ecd08b1d381db0
---
 include/rustc_demangle.h |   9 +-
 src/lib.rs               | 339 +++++++++++++++++++++++++++------------
 2 files changed, 238 insertions(+), 110 deletions(-)

diff --git a/include/rustc_demangle.h b/include/rustc_demangle.h
index 61c4aa1..e7ee2ca 100644
--- a/include/rustc_demangle.h
+++ b/include/rustc_demangle.h
@@ -5,11 +5,10 @@
 extern "C" {
 #endif
 
-// Demangles symbol given in `mangled` argument into `out` buffer
-//
-// Returns 0 if `mangled` is not Rust symbol or if `out` buffer is too small
-// Returns 1 otherwise
-int rustc_demangle(const char *mangled, char *out, size_t out_size);
+// For size_t
+#include <stddef.h>
+
+char *rustc_demangle(const char *mangled, char *out, size_t *len, int *status);
 
 #ifdef __cplusplus
 }
diff --git a/src/lib.rs b/src/lib.rs
index 51e3103..7610145 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,160 +1,289 @@
 extern crate rustc_demangle;
 
+use std::alloc::{GlobalAlloc, Layout, System};
 use std::io::Write;
 use std::os::raw::{c_char, c_int};
+use std::ptr;
+use std::result;
+
+type Result<T> = result::Result<T, Status>;
+
+/// Convenience function to set return status if a location was provided.
+unsafe fn set_status(status: *mut c_int, val: c_int) {
+    if !status.is_null() {
+        *status = val;
+    }
+}
+
+/// Region from the system allocator for demangler output. We use the
+/// system allocator because the intended client is C/C++ code which
+/// may not be using the Rust allocator.
+struct SystemBuffer {
+    buf: *mut u8,
+    size: usize,
+    size_out: *mut usize,
+}
+
+impl SystemBuffer {
+    const DEFAULT_BUFFER_SIZE: usize = 1024;
+    fn new(size: usize) -> Result<Self> {
+        let buf = unsafe { System.alloc_zeroed(Layout::from_size_align_unchecked(size, 1)) };
+        if buf.is_null() {
+            Err(Status::AllocFailure)
+        } else {
+            Ok(Self {
+                buf,
+                size,
+                size_out: ptr::null_mut(),
+            })
+        }
+    }
+    /// Safety: If buf is non-null, size must be non-null and point to the
+    /// non-zero size of the buffer provided in buf.
+    /// Takes ownership of the buffer passed in (and may reallocate it).
+    /// size must outlive the resulting buffer if non-null.
+    unsafe fn from_raw(buf: *mut c_char, size: *mut usize) -> Result<Self> {
+        if buf.is_null() {
+            if !size.is_null() {
+                *size = Self::DEFAULT_BUFFER_SIZE;
+            }
+            let fresh = Self::new(Self::DEFAULT_BUFFER_SIZE)?;
+            Ok(Self {
+                size_out: size,
+                ..fresh
+            })
+        } else {
+            Ok(Self {
+                buf: buf as *mut u8,
+                size: *size,
+                size_out: size,
+            })
+        }
+    }
+    fn as_mut_slice(&mut self) -> &mut [u8] {
+        unsafe { std::slice::from_raw_parts_mut(self.buf, self.size) }
+    }
+    fn resize(&mut self) -> Result<()> {
+        let new_size = self.size * 2;
+        let new_buf = unsafe {
+            System.realloc(
+                self.buf,
+                Layout::from_size_align_unchecked(self.size, 1),
+                new_size,
+            )
+        };
+        if new_buf.is_null() {
+            Err(Status::AllocFailure)
+        } else {
+            self.buf = new_buf;
+            self.size = new_size;
+            if !self.size_out.is_null() {
+                unsafe {
+                    *self.size_out = new_size;
+                }
+            }
+            Ok(())
+        }
+    }
+}
 
 /// C-style interface for demangling.
-/// Demangles symbol given in `mangled` argument into `out` buffer
+/// Demangles symbol given in `mangled` argument into `out` buffer.
+///
+/// This interface is a drop-in replacement for `__cxa_demangle`, but for
+/// Rust demangling.
+///
+/// If `out` is null, a buffer will be allocated using the system allocator
+/// to contain the results.
+/// If `out` is non-null, `out_size` must be a pointer to the current size
+/// of the buffer, and `out` must come from the system allocator.
+/// If `out_size` is non-null, the size of the output buffer will be written
+/// to it.
+///
+/// If `status` is non-null, it will be set to one of the following values:
+/// * 0: Demangling succeeded
+/// * -1: Allocation failure
+/// * -2: Name did not demangle
+/// * -3: Invalid arguments
+///
+/// Returns null if `mangled` is not Rust symbol or demangling failed.
+/// Returns the buffer containing the demangled symbol name otherwise.
 ///
 /// Unsafe as it handles buffers by raw pointers.
 ///
-/// Returns 0 if `mangled` is not Rust symbol or if `out` buffer is too small
-/// Returns 1 otherwise
+/// For non-null `out`, `out_size` represents a slight deviation from the
+/// `__cxa_demangle` behavior. For `__cxa_demangle`, the buffer must be at
+/// *least* the provided size. For `rustc_demangle`, it must be the exact
+/// buffer size because it is used in the reconstruction of the `Layout`
+/// for use with `::realloc`.
 #[no_mangle]
 pub unsafe extern "C" fn rustc_demangle(
     mangled: *const c_char,
     out: *mut c_char,
-    out_size: usize,
-) -> c_int {
+    out_size: *mut usize,
+    status: *mut c_int,
+) -> *mut c_char {
+    match rustc_demangle_native(mangled, out, out_size) {
+        Ok(demangled) => {
+            set_status(status, 0);
+            demangled
+        }
+        Err(e) => {
+            set_status(status, e as c_int);
+            ptr::null_mut()
+        }
+    }
+}
+
+enum Status {
+    AllocFailure = -1,
+    DemangleFailure = -2,
+    InvalidArgs = -3,
+}
+
+unsafe fn rustc_demangle_native(
+    mangled: *const c_char,
+    out: *mut c_char,
+    out_size: *mut usize,
+) -> Result<*mut c_char> {
+    if mangled.is_null() {
+        return Err(Status::InvalidArgs);
+    }
     let mangled_str = match std::ffi::CStr::from_ptr(mangled).to_str() {
         Ok(s) => s,
-        Err(_) => return 0,
+        Err(_) => return Err(Status::InvalidArgs),
     };
+
+    if !out.is_null() {
+        if out_size.is_null() {
+            return Err(Status::InvalidArgs);
+        }
+        if *out_size == 0 {
+            return Err(Status::InvalidArgs);
+        }
+    }
+
+    let mut out_buf = SystemBuffer::from_raw(out, out_size)?;
+
     match rustc_demangle::try_demangle(mangled_str) {
         Ok(demangle) => {
-            let mut out_slice = std::slice::from_raw_parts_mut(out as *mut u8, out_size);
-            match write!(out_slice, "{:#}\0", demangle) {
-                Ok(_) => return 1,
-                Err(_) => return 0,
+            while write!(out_buf.as_mut_slice(), "{:#}\0", demangle).is_err() {
+                out_buf.resize()?;
             }
+            Ok(out_buf.as_mut_slice().as_mut_ptr() as *mut c_char)
         }
-        Err(_) => return 0,
+        Err(_) => Err(Status::DemangleFailure),
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use std;
-    use std::os::raw::c_char;
+    use std::alloc::{GlobalAlloc, Layout, System};
+    use std::os::raw::{c_char, c_int};
+    use std::ptr;
+
+    struct DemangleResult {
+        out_buf: *mut u8,
+        out_size: usize,
+        status: c_int,
+    }
+
+    impl Drop for DemangleResult {
+        fn drop(&mut self) {
+            if !self.out_buf.is_null() {
+                unsafe {
+                    System.dealloc(
+                        self.out_buf,
+                        Layout::from_size_align_unchecked(self.out_size, 1),
+                    );
+                }
+            }
+        }
+    }
+
+    impl DemangleResult {
+        fn as_slice(&self) -> &[u8] {
+            unsafe { std::slice::from_raw_parts(self.out_buf, self.out_size) }
+        }
+    }
+
+    fn demangle(mangled: &str, alloc_size: usize) -> DemangleResult {
+        unsafe { raw_demangle(mangled.as_ptr() as *const c_char, alloc_size) }
+    }
+
+    unsafe fn raw_demangle(mangled: *const c_char, alloc_size: usize) -> DemangleResult {
+        let mut out_size: usize = alloc_size;
+        let mut status: c_int = 0;
+        let out_buf: *mut c_char = if out_size != 0 {
+            System.alloc(Layout::from_size_align_unchecked(out_size, 1)) as *mut c_char
+        } else {
+            ptr::null_mut()
+        };
+        ptr::write_bytes(out_buf, '*' as u8, out_size);
+
+        let res = super::rustc_demangle(mangled, out_buf, &mut out_size, &mut status);
+        DemangleResult {
+            out_buf: res as *mut u8,
+            out_size,
+            status,
+        }
+    }
+
     #[test]
     fn demangle_c_str_large() {
-        let mangled = "_ZN4testE\0";
-        let mut out_buf: Vec<u8> = vec![42; 8];
-        let res = unsafe {
-            super::rustc_demangle(
-                mangled.as_ptr() as *const c_char,
-                out_buf.as_mut_ptr() as *mut c_char,
-                8,
-            )
-        };
-        assert_eq!(res, 1);
-        let out_str = std::str::from_utf8(&out_buf[..5]).unwrap();
+        let res = demangle("_ZN4testE\0", 8);
+        assert_eq!(res.status, 0);
+        let out_str = core::str::from_utf8(&res.as_slice()[..5]).unwrap();
         assert_eq!(out_str, "test\0");
     }
 
     #[test]
     fn demangle_c_str_exact() {
-        let mangled = "_ZN4testE\0";
-        let mut out_buf: Vec<u8> = vec![42; 8];
-        let res = unsafe {
-            super::rustc_demangle(
-                mangled.as_ptr() as *const c_char,
-                out_buf.as_mut_ptr() as *mut c_char,
-                5,
-            )
-        };
-        assert_eq!(res, 1);
-        let out_str = std::str::from_utf8(&out_buf).unwrap();
+        let res = demangle("_ZN4testE\0", 8);
+        assert_eq!(res.status, 0);
+        // No reallocation necessary, so our * fill should be present
+        let out_str = core::str::from_utf8(res.as_slice()).unwrap();
         assert_eq!(out_str, "test\0***");
     }
 
     #[test]
     fn demangle_c_str_small() {
-        let mangled = "_ZN4testE\0";
-        let mut out_buf: Vec<u8> = vec![42; 8];
-        let res = unsafe {
-            super::rustc_demangle(
-                mangled.as_ptr() as *const c_char,
-                out_buf.as_mut_ptr() as *mut c_char,
-                4,
-            )
-        };
-        assert_eq!(res, 0);
-        let out_str = std::str::from_utf8(&out_buf[4..]).unwrap();
-        assert_eq!(out_str, "****");
-    }
-
-    #[test]
-    fn demangle_c_str_smaller() {
-        let mangled = "_ZN4testE\0";
-        let mut out_buf: Vec<u8> = vec![42; 8];
-        let res = unsafe {
-            super::rustc_demangle(
-                mangled.as_ptr() as *const c_char,
-                out_buf.as_mut_ptr() as *mut c_char,
-                3,
-            )
-        };
-        assert_eq!(res, 0);
-        let out_str = std::str::from_utf8(&out_buf[3..]).unwrap();
-        assert_eq!(out_str, "*****");
+        let res = demangle("_ZN4testE\0", 4);
+        assert_eq!(res.status, 0);
+        // demangle should have realloced
+        assert_ne!(res.out_size, 4);
+        // Only check the start, since the reallocation means our * fill may
+        // be absent.
+        let out_str = core::str::from_utf8(&res.as_slice()[..5]).unwrap();
+        assert_eq!(out_str, "test\0");
     }
 
     #[test]
-    fn demangle_c_str_zero() {
-        let mangled = "_ZN4testE\0";
-        let mut out_buf: Vec<u8> = vec![42; 8];
-        let res = unsafe {
-            super::rustc_demangle(
-                mangled.as_ptr() as *const c_char,
-                out_buf.as_mut_ptr() as *mut c_char,
-                0,
-            )
-        };
-        assert_eq!(res, 0);
-        let out_str = std::str::from_utf8(&out_buf).unwrap();
-        assert_eq!(out_str, "********");
+    fn demangle_c_str_alloc() {
+        let res = demangle("_ZN4testE\0", 0);
+        assert_eq!(res.status, 0);
+        // demangle should have allocated
+        assert_ne!(res.out_size, 0);
+        let out_str = core::str::from_utf8(&res.as_slice()[..5]).unwrap();
+        assert_eq!(out_str, "test\0");
     }
 
     #[test]
     fn demangle_c_str_not_rust_symbol() {
-        let mangled = "la la la\0";
-        let mut out_buf: Vec<u8> = vec![42; 8];
-        let res = unsafe {
-            super::rustc_demangle(
-                mangled.as_ptr() as *const c_char,
-                out_buf.as_mut_ptr() as *mut c_char,
-                8,
-            )
-        };
-        assert_eq!(res, 0);
+        let res = demangle("la la la\0", 8);
+        assert_eq!(res.status, -2);
     }
 
     #[test]
     fn demangle_c_str_null() {
-        let mangled = "\0";
-        let mut out_buf: Vec<u8> = vec![42; 8];
-        let res = unsafe {
-            super::rustc_demangle(
-                mangled.as_ptr() as *const c_char,
-                out_buf.as_mut_ptr() as *mut c_char,
-                8,
-            )
-        };
-        assert_eq!(res, 0);
+        let res = demangle("\0", 8);
+        assert_eq!(res.status, -2);
     }
 
     #[test]
     fn demangle_c_str_invalid_utf8() {
         let mangled = [116, 101, 115, 116, 165, 0];
-        let mut out_buf: Vec<u8> = vec![42; 8];
-        let res = unsafe {
-            super::rustc_demangle(
-                mangled.as_ptr() as *const c_char,
-                out_buf.as_mut_ptr() as *mut c_char,
-                8,
-            )
-        };
-        assert_eq!(res, 0);
+        let res = unsafe { raw_demangle(mangled.as_ptr() as *const c_char, 8) };
+        assert_eq!(res.status, -2);
     }
 }
-- 
2.33.0.rc1.237.g0d66db33f3-goog