Letta MCP Server

use core::mem::MaybeUninit;

/// Unpoisons `buf` if MSAN support is enabled.
///
/// Most backends do not need to unpoison their output. Rust language- and
/// library- provided functionality unpoisons automatically. Similarly, libc
/// either natively supports MSAN and/or MSAN hooks libc-provided functions
/// to unpoison outputs on success. Only when all of these things are
/// bypassed do we need to do it ourselves.
///
/// The call to unpoison should be done as close to the write as possible.
/// For example, if the backend partially fills the output buffer in chunks,
/// each chunk should be unpoisoned individually. This way, the correctness of
/// the chunking logic can be validated (in part) using MSAN.
pub unsafe fn unpoison(buf: &mut [MaybeUninit<u8>]) {
    cfg_if! {
        if #[cfg(getrandom_msan)] {
            extern "C" {
                fn __msan_unpoison(a: *mut core::ffi::c_void, size: usize);
            }
            let a = buf.as_mut_ptr().cast();
            let size = buf.len();
            #[allow(unused_unsafe)] // TODO(MSRV 1.65): Remove this.
            unsafe {
                __msan_unpoison(a, size);
            }
        } else {
            let _ = buf;
        }
    }
}

/// Interprets the result of the `getrandom` syscall of Linux, unpoisoning any
/// written part of `buf`.
///
/// `buf` must be the output buffer that was originally passed to the `getrandom`
/// syscall.
///
/// `ret` must be the result returned by `getrandom`. If `ret` is negative or
/// larger than the length of `buf` then nothing is done.
///
/// Memory Sanitizer only intercepts `getrandom` on this condition (from its
/// source code):
/// ```c
/// #define SANITIZER_INTERCEPT_GETRANDOM \
///   ((SI_LINUX && __GLIBC_PREREQ(2, 25)) || SI_FREEBSD || SI_SOLARIS)
/// ```
/// So, effectively, we have to assume that it is never intercepted on Linux.
#[cfg(any(target_os = "android", target_os = "linux"))]
pub unsafe fn unpoison_linux_getrandom_result(buf: &mut [MaybeUninit<u8>], ret: isize) {
    if let Ok(bytes_written) = usize::try_from(ret) {
        if let Some(written) = buf.get_mut(..bytes_written) {
            #[allow(unused_unsafe)] // TODO(MSRV 1.65): Remove this.
            unsafe {
                unpoison(written)
            }
        }
    }
}