crashtracker_runtime_stacks.rs

use std::cmp;
use std::ffi::{c_char, c_int, c_void};
use std::ptr;
use std::slice;
use std::sync::Once;

use libdd_crashtracker::RuntimeStackFrame;

// We define these raw system calls here to be used within signal handler context.
// These direct C functions are preferred over going through Rust wrappers
extern "C" {
    fn pipe(pipefd: *mut [c_int; 2]) -> c_int;
    fn read(fd: c_int, buf: *mut c_void, count: usize) -> isize;
    fn close(fd: c_int) -> c_int;
    fn fcntl(fd: c_int, cmd: c_int, arg: c_int) -> c_int;
}

/// Check whether the current thread holds the GIL.
/// This is signal-safe: it's just a thread-local read + comparison.
/// Returns false if the GIL is not held (e.g. during a ctypes foreign function call),
/// in which case it's unsafe to call any Python C API functions.
///
/// `PyGILState_Check` is not limited-API: it lives in CPython's
/// `Include/cpython/pystate.h` and is absent from `Misc/stable_abi.toml`
/// (3.15 lists Ensure / Release / GetThisThreadState only). pyo3-ffi therefore
/// omits it under `Py_LIMITED_API` (3.15 via `PYO3_USE_ABI3_FORWARD_COMPATIBILITY`).
/// This runs in a crash signal handler, so `PyGILState_Ensure` is not a
/// substitute (not async-signal-safe). On limited-API builds we cannot probe,
/// so skip Python C API use rather than acquire or guess.
#[cfg(all(unix, not(Py_LIMITED_API)))]
unsafe fn gil_is_held() -> bool {
    pyo3_ffi::PyGILState_Check() != 0
}

#[cfg(all(unix, Py_LIMITED_API))]
unsafe fn gil_is_held() -> bool {
    false
}

/************************************************************
 Emit runtime stacktrace as string using _Py_DumpTracebackThreads /
 PyUnstable_DumpTracebackThreads
************************************************************/

// CPython 3.15 renamed _Py_DumpTracebackThreads to PyUnstable_DumpTracebackThreads
// and added a 4th `max_threads` parameter.  We use the 4-arg signature universally:
// on older CPython the extra register argument is harmless (x86_64 SysV / ARM64 AAPCS
// pass the first >=4 integer args in registers), and on 3.15+ passing 0 means
// "use the default limit".
type PyDumpTracebackThreadsFn = unsafe extern "C" fn(
    fd: c_int,
    interp: *mut pyo3_ffi::PyInterpreterState,
    current_tstate: *mut pyo3_ffi::PyThreadState,
    max_threads: isize,
) -> *const c_char;

// Function pointer type for PyFrame_GetBack
type PyFrameGetBackFn =
    unsafe extern "C" fn(*mut pyo3_ffi::PyFrameObject) -> *mut pyo3_ffi::PyFrameObject;

// Cached function pointers to avoid dlsym during crash
static mut DUMP_TRACEBACK_FN: Option<PyDumpTracebackThreadsFn> = None;
static DUMP_TRACEBACK_INIT: Once = Once::new();
static mut FRAME_GET_BACK_FN: Option<PyFrameGetBackFn> = None;
static FRAME_GET_BACK_INIT: Once = Once::new();

const MAX_TRACEBACK_SIZE: usize = 8 * 1024; // 8KB

// Attempt to resolve the CPython traceback-dump function at runtime.
// CPython <= 3.14 exports `_Py_DumpTracebackThreads` (3 args).
// CPython >= 3.15 replaced it with `PyUnstable_DumpTracebackThreads` (4 args, adds max_threads).
// We try the new name first, then fall back to the old one.
pub unsafe fn init_dump_traceback_fn() {
    DUMP_TRACEBACK_INIT.call_once(|| {
        #[cfg(unix)]
        {
            extern "C" {
                fn dlsym(
                    handle: *mut std::ffi::c_void,
                    symbol: *const std::ffi::c_char,
                ) -> *mut std::ffi::c_void;
            }

            const RTLD_DEFAULT: *mut std::ffi::c_void = ptr::null_mut();

            // Try the CPython 3.15+ public unstable name first.
            let mut symbol_ptr = dlsym(RTLD_DEFAULT, c"PyUnstable_DumpTracebackThreads".as_ptr());

            // Fall back to the private name used through CPython 3.14.
            if symbol_ptr.is_null() {
                symbol_ptr = dlsym(RTLD_DEFAULT, c"_Py_DumpTracebackThreads".as_ptr());
            }

            if !symbol_ptr.is_null() {
                DUMP_TRACEBACK_FN =
                    Some(std::mem::transmute::<*mut c_void, PyDumpTracebackThreadsFn>(symbol_ptr));
            }
        }

        #[cfg(not(unix))]
        {
            // DUMP_TRACEBACK_FN remains None on non-Unix platforms
        }
    });
}

// Get the cached function pointer; should only be called after init_dump_traceback_fn
pub unsafe fn get_cached_dump_traceback_fn() -> Option<PyDumpTracebackThreadsFn> {
    DUMP_TRACEBACK_FN
}

// PyFrame_GetBack is in CPython's public C API through 3.14+,
// but is absent from the limited API (Py_LIMITED_API). Rather than using compile-time
// cfg guards that break under PYO3_USE_ABI3_FORWARD_COMPATIBILITY, we resolve the symbol
// at runtime by using dlsym. This is safe because init runs at crashtracker registration time,
// not during the crash signal handler.
pub unsafe fn init_frame_get_back_fn() {
    FRAME_GET_BACK_INIT.call_once(|| {
        #[cfg(unix)]
        {
            extern "C" {
                fn dlsym(
                    handle: *mut std::ffi::c_void,
                    symbol: *const std::ffi::c_char,
                ) -> *mut std::ffi::c_void;
            }

            const RTLD_DEFAULT: *mut std::ffi::c_void = ptr::null_mut();

            let symbol_ptr = dlsym(RTLD_DEFAULT, c"PyFrame_GetBack".as_ptr());

            if !symbol_ptr.is_null() {
                FRAME_GET_BACK_FN = Some(std::mem::transmute::<*mut c_void, PyFrameGetBackFn>(
                    symbol_ptr,
                ));
            }
        }

        #[cfg(not(unix))]
        {
            // FRAME_GET_BACK_FN remains None on non-Unix platforms
        }
    });
}

#[cfg(Py_LIMITED_API)]
unsafe fn get_cached_frame_get_back_fn() -> Option<PyFrameGetBackFn> {
    FRAME_GET_BACK_FN
}

unsafe fn dump_python_traceback_as_string(
    emit_stacktrace_string: unsafe extern "C" fn(*const c_char),
) {
    if !gil_is_held() {
        emit_stacktrace_string(c"<python_runtime_stacktrace_unavailable>".as_ptr());
        return;
    }

    // Use function linked during registration
    let dump_fn = match get_cached_dump_traceback_fn() {
        Some(func) => func,
        None => {
            emit_stacktrace_string(c"<python_runtime_stacktrace_unavailable>".as_ptr());
            return;
        }
    };

    // Create a pipe to capture CPython internal traceback dump. _Py_DumpTracebackThreads writes to
    // a fd. Reading and writing to pipe is signal-safe. We stack allocate a buffer in the beginning,
    // and use it to read the output
    let mut pipefd: [c_int; 2] = [0, 0];
    if pipe(&mut pipefd as *mut [c_int; 2]) != 0 {
        emit_stacktrace_string(c"<pipe_creation_failed>".as_ptr());
        return;
    }

    let read_fd = pipefd[0];
    let write_fd = pipefd[1];

    if fcntl(read_fd, libc::F_SETFL as c_int, libc::O_NONBLOCK as c_int) == -1 {
        // Non-fatal: log and continue; read may block but won't crash.
        let msg = b"<fcntl_O_NONBLOCK_failed>\0";
        libc::write(
            libc::STDERR_FILENO,
            msg.as_ptr() as *const libc::c_void,
            msg.len(),
        );
    }

    // Use null thread state for signal-safety; CPython will dump all threads.
    // max_threads=0 tells CPython 3.15+ to use its default limit.
    let error_msg = dump_fn(write_fd, ptr::null_mut(), ptr::null_mut(), 0);

    // Ignore close return: EINTR cannot occur on Linux for close, and
    // any other error means the fd is already gone — nothing to recover.
    let _ = close(write_fd);

    if !error_msg.is_null() {
        let _ = close(read_fd);
        emit_stacktrace_string(error_msg as *const c_char);
        return;
    }

    let mut buffer = [0u8; MAX_TRACEBACK_SIZE];
    let bytes_read = read(
        read_fd,
        buffer.as_mut_ptr() as *mut c_void,
        MAX_TRACEBACK_SIZE,
    );

    let _ = close(read_fd);

    if bytes_read > 0 {
        let bytes_read = bytes_read as usize;
        if bytes_read < MAX_TRACEBACK_SIZE {
            buffer[bytes_read] = 0;
        } else {
            // Buffer is full; add truncation indicator
            let truncation_msg = b"\n[TRUNCATED]\0";
            let msg_len = truncation_msg.len();
            if MAX_TRACEBACK_SIZE >= msg_len {
                let start_pos = MAX_TRACEBACK_SIZE - msg_len;
                buffer[start_pos..].copy_from_slice(truncation_msg);
            }
        }
        emit_stacktrace_string(buffer.as_ptr() as *const c_char);
        return;
    }

    emit_stacktrace_string(c"<traceback_read_failed>".as_ptr());
}

pub unsafe extern "C" fn native_runtime_stack_string_callback(
    emit_stacktrace_string: unsafe extern "C" fn(*const c_char),
) {
    dump_python_traceback_as_string(emit_stacktrace_string);
}

/************************************************************
 Emit runtime stacktrace as frames using by traversing the frame object chain
************************************************************/

#[cfg(unix)]
const FRAME_FUNCTION_CAP: usize = 256;
#[cfg(unix)]
const FRAME_FILE_CAP: usize = 512;

#[cfg(unix)]
unsafe fn capture_frames_via_python(emit_frame: unsafe extern "C" fn(&RuntimeStackFrame)) {
    if !gil_is_held() {
        return;
    }
    let current = pyo3_ffi::PyThreadState_Get();
    if !current.is_null() {
        let _ = collect_and_emit_frames_for_thread(current, emit_frame);
    }
}

#[cfg(unix)]
unsafe fn collect_and_emit_frames_for_thread(
    tstate: *mut pyo3_ffi::PyThreadState,
    emit_frame: unsafe extern "C" fn(&RuntimeStackFrame),
) -> bool {
    if tstate.is_null() {
        return false;
    }

    let mut emitted = false;
    let mut frame = thread_top_frame(tstate);

    while !frame.is_null() {
        if emit_python_frame(frame, emit_frame) {
            emitted = true;
        }
        frame = advance_frame(frame);
    }

    emitted
}

unsafe fn thread_top_frame(tstate: *mut pyo3_ffi::PyThreadState) -> *mut pyo3_ffi::PyFrameObject {
    if tstate.is_null() {
        ptr::null_mut()
    } else {
        // We won't support Python 3.9
        #[cfg(not(Py_3_10))]
        let frame: *mut pyo3_ffi::PyFrameObject = ptr::null_mut();

        #[cfg(Py_3_10)]
        let frame = pyo3_ffi::PyThreadState_GetFrame(tstate);

        #[cfg(not(Py_3_11))]
        {
            if !frame.is_null() {
                // On CPython <= 3.10 it returns/holds a borrowed frame, so we need to
                // Py_XINCREF to keep it alive while we traverse.
                pyo3_ffi::Py_XINCREF(frame as *mut pyo3_ffi::PyObject);
            }
        }
        frame
    }
}

unsafe fn advance_frame(frame: *mut pyo3_ffi::PyFrameObject) -> *mut pyo3_ffi::PyFrameObject {
    if frame.is_null() {
        return ptr::null_mut();
    }
    // PyFrame_GetBack is absent from pyo3-ffi's stable/limited-API bindings (Py_LIMITED_API),
    // which is activated for Python 3.15+ via PYO3_USE_ABI3_FORWARD_COMPATIBILITY.
    // On Py_LIMITED_API builds the crashtracker captures only the top frame rather than the full stack.
    #[cfg(not(Py_LIMITED_API))]
    let back = pyo3_ffi::PyFrame_GetBack(frame);

    #[cfg(Py_LIMITED_API)]
    let back = match get_cached_frame_get_back_fn() {
        Some(get_back) => get_back(frame),
        None => ptr::null_mut(),
    };

    pyo3_ffi::Py_DecRef(frame as *mut pyo3_ffi::PyObject);
    back
}

#[cfg(unix)]
unsafe fn emit_python_frame(
    frame: *mut pyo3_ffi::PyFrameObject,
    emit_frame: unsafe extern "C" fn(&RuntimeStackFrame),
) -> bool {
    if frame.is_null() {
        return false;
    }

    let file_view = get_code_attr_utf8_view(frame, b"co_filename\0");

    // For versions < 3.11, we should use `co_name`
    #[cfg(not(Py_3_11))]
    let function_view = get_code_attr_utf8_view(frame, b"co_name\0");

    // For versions 3.11+, we should use qualified name
    #[cfg(Py_3_11)]
    let function_view = get_code_attr_utf8_view(frame, b"co_qualname\0");

    let line_number = pyo3_ffi::PyFrame_GetLineNumber(frame);

    let file_slice = file_view
        .as_ref()
        .map(|view| view.truncated(FRAME_FILE_CAP))
        .unwrap_or(&[]);
    let function_slice = function_view
        .as_ref()
        .map(|view| view.truncated(FRAME_FUNCTION_CAP))
        .unwrap_or(&[]);

    let runtime_frame = RuntimeStackFrame {
        line: if line_number < 0 {
            0
        } else {
            line_number as u32
        },
        column: 0,
        function: function_slice,
        file: file_slice,
        type_name: &[],
    };

    emit_frame(&runtime_frame);

    if let Some(view) = file_view {
        pyo3_ffi::Py_DecRef(view.obj);
    }
    if let Some(view) = function_view {
        pyo3_ffi::Py_DecRef(view.obj);
    }
    true
}

struct Utf8View {
    ptr: *const u8,
    len: usize,
    obj: *mut pyo3_ffi::PyObject,
}

impl Utf8View {
    /// Safety: caller must ensure the underlying PyObject outlives the returned slice.
    /// This should be safe since no garbage collection should be happening while suspended in
    /// crash context
    fn truncated(&self, cap: usize) -> &[u8] {
        let len = cmp::min(self.len, cap);
        unsafe { slice::from_raw_parts(self.ptr, len) }
    }
}

/// Returns a view into a PyUnicode attribute without allocating; caller must Py_DecRef `obj`.
#[cfg(unix)]
unsafe fn get_code_attr_utf8_view(
    frame: *mut pyo3_ffi::PyFrameObject,
    attr: &[u8],
) -> Option<Utf8View> {
    let code_obj = pyo3_ffi::PyFrame_GetCode(frame) as *mut pyo3_ffi::PyObject;
    if code_obj.is_null() {
        return None;
    }
    let attr_obj = pyo3_ffi::PyObject_GetAttrString(code_obj, attr.as_ptr() as *const c_char);
    pyo3_ffi::Py_DecRef(code_obj);
    if attr_obj.is_null() {
        return None;
    }

    let mut size: pyo3_ffi::Py_ssize_t = 0;
    let data = pyo3_ffi::PyUnicode_AsUTF8AndSize(attr_obj, &mut size);
    if data.is_null() || size <= 0 {
        pyo3_ffi::Py_DecRef(attr_obj);
        return None;
    }

    Some(Utf8View {
        ptr: data as *const u8,
        len: size as usize,
        obj: attr_obj,
    })
}

unsafe fn dump_python_traceback_as_frames(emit_frame: unsafe extern "C" fn(&RuntimeStackFrame)) {
    #[cfg(unix)]
    {
        if emit_frame as usize == 0 {
            return;
        }

        capture_frames_via_python(emit_frame);
    }

    #[cfg(not(unix))]
    {
        let _ = emit_frame;
    }
}

pub unsafe extern "C" fn native_runtime_stack_frame_callback(
    emit_frame: unsafe extern "C" fn(&RuntimeStackFrame),
) {
    dump_python_traceback_as_frames(emit_frame);
}