rs_malloc_tracker 1.0.1

#![allow(clippy::deref_addrof)] // necessary for static muts

use std::{
    ffi::c_void,
    sync::{LazyLock, Mutex, RwLock},
    thread,
};

#[cfg(test)]
use std::sync::mpsc::Sender;

use ctor::{ctor, dtor};
use libc::{MAP_ANONYMOUS, PR_SET_VMA, PR_SET_VMA_ANON_NAME, c_long, c_ulong, prctl};
use log::*;
use no_alloc_channel::{Channel, SendHalf};
use nodit::Interval;
use nodit::{NoditMap, interval::ie};

use crate::{c_stringify, fn_impl, get_frame, shadow};

fn_impl!(
    extern "C" fn(
        addr: *mut c_void,
        len: usize,
        prot: i32,
        flags: i32,
        fildes: i32,
        off: usize,
    ) -> *mut c_void
);
fn_impl!(extern "C" fn(addr: *mut c_void, len: usize) -> i32);

enum Operation {
    Mmap((usize, usize, Mmap)),
    Munmap((usize, usize)),
    #[cfg(test)]
    TestFence(Sender<()>), // used by tests to wait for pending operations to be processed
    Shutdown,
}

#[cfg(test)]
pub fn test_wait_mmap_processed() {
    use std::sync::mpsc::channel;

    let (sender, receiver) = channel::<()>();
    SENDER
        .lock()
        .unwrap()
        .as_mut()
        .unwrap()
        .try_send(Operation::TestFence(sender))
        .unwrap();
    receiver.recv().unwrap();
}

#[allow(dead_code)]
#[derive(Clone, Debug)]
pub struct Mmap {
    pub flags: i32,
    pub ip: Option<usize>,
}

// Heapless channel so we can pass over mmap() allocations without blocking
// this allows us to instrument jemalloc() without deadlocking its mmap calls
static SENDER: LazyLock<Mutex<Option<SendHalf<'_, Operation, 64>>>> =
    LazyLock::new(|| Mutex::new(None));

#[unsafe(no_mangle)]
pub static MMAP_ZONES: LazyLock<RwLock<NoditMap<usize, Interval<usize>, Mmap>>> =
    LazyLock::new(|| RwLock::new(NoditMap::default()));
#[unsafe(no_mangle)]
pub static mut MMAP_SEND_MISSES: LazyLock<RwLock<usize>> = LazyLock::new(|| RwLock::new(0));
#[unsafe(no_mangle)]
pub static mut MMAP_SEND_FAILS: LazyLock<RwLock<usize>> = LazyLock::new(|| RwLock::new(0));

#[ctor]
fn init() {
    static mut C: Channel<Operation, 64> = Channel::new();

    let (mut producer, mut consumer) = unsafe { (*(&raw mut C)).split() };
    SENDER.lock().unwrap().replace(producer);

    thread::spawn(move || {
        futures::executor::block_on(async move {
            debug!("mmap consumer thread starting");

            loop {
                match consumer.receive().await {
                    Operation::Mmap((ptr, len, map)) => {
                        MMAP_ZONES
                            .write()
                            .unwrap()
                            .insert_overwrite(ie(ptr, ptr + len), map);
                    }
                    Operation::Munmap((ptr, len)) => {
                        MMAP_ZONES
                            .write()
                            .unwrap()
                            .cut(ie(ptr, ptr + len))
                            .collect::<Vec<_>>();
                    }
                    #[cfg(test)]
                    Operation::TestFence(s) => s.send(()).unwrap(),
                    Operation::Shutdown => break,
                }
            }

            debug!("mmap consumer thread disconnecting");
        })
    });
}

#[dtor]
fn deinit() {
    if let Some(s) = SENDER.lock().unwrap().as_mut() {
        s.try_send(Operation::Shutdown).unwrap_or(())
    };
}

#[inline(always)]
fn handle_mmap(ptr: *mut c_void, len: usize, flags: i32) {
    let ip = get_frame!(2).map(|i| i as usize);

    if flags & MAP_ANONYMOUS != 0 {
        unsafe {
            prctl(
                PR_SET_VMA,
                PR_SET_VMA_ANON_NAME as c_long,
                ptr as c_ulong,
                len as c_ulong,
                c_stringify!("instrumented-anon-mmap"),
            )
        };
    }

    SENDER
        .lock()
        .unwrap()
        .as_mut()
        .map(|s| {
            s.try_send(Operation::Mmap((ptr as usize, len, Mmap { flags, ip })))
                .inspect_err(|_| unsafe { (*(*(&raw mut MMAP_SEND_FAILS)).write().unwrap()) += 1 })
                .unwrap_or(())
        })
        .ok_or(())
        .inspect_err(|_| unsafe { (*(*(&raw mut MMAP_SEND_MISSES)).write().unwrap()) += 1 })
        .unwrap_or(());
}

shadow!(mmap(addr: *mut c_void, len: usize, prot: i32, flags: i32, fildes: i32, off: usize) -> *mut c_void {
    let ptr = f!()(addr, len, prot, flags, fildes, off);
    if ptr != std::ptr::without_provenance_mut(usize::MAX) { // MAP_FAILED
        handle_mmap(ptr, len, flags);
    }
    ptr
});

shadow!(mmap64(addr: *mut c_void, len: usize, prot: i32, flags: i32, fildes: i32, off: usize) -> *mut c_void {
    let ptr = f!()(addr, len, prot, flags, fildes, off);
    if ptr != std::ptr::without_provenance_mut(usize::MAX) { // MAP_FAILED
        handle_mmap(ptr, len, flags);
    }
    ptr
});

shadow!(munmap(addr: *mut c_void, len: usize) -> i32 {
    let ret = f!()(addr, len);

    if ret == 0 {
        SENDER.lock().unwrap().as_mut().map(|s| s.try_send(
            Operation::Munmap((addr as usize, len))
        ).unwrap_or(())).unwrap_or(());
    }

    ret
});