wry/src/clientmem.rs

use {
    crate::{
        client::Client,
        cpu_worker::{AsyncCpuWork, CpuJob, CpuWork, CpuWorker},
        gfx_api::{ShmMemory, ShmMemoryBacking},
        utils::{page_size::page_size, vec_ext::VecExt},
    },
    std::{
        cell::Cell,
        error::Error,
        mem::{ManuallyDrop, MaybeUninit},
        ops::Deref,
        ptr,
        rc::Rc,
        sync::atomic::{Ordering, compiler_fence},
    },
    thiserror::Error,
    uapi::{
        OwnedFd, Pod,
        c::{self, raise},
        ftruncate,
    },
};

#[derive(Debug, Error)]
pub enum ClientMemError {
    #[error("Could not install the sigbus handler")]
    SigactionFailed(#[source] crate::utils::oserror::OsError),
    #[error("A SIGBUS occurred while accessing mapped memory")]
    Sigbus,
    #[error("mmap failed")]
    MmapFailed(#[source] crate::utils::oserror::OsError),
    #[error("Length was not a multiple of the data element size")]
    InvalidLength,
}

pub struct ClientMem {
    fd: ManuallyDrop<Rc<OwnedFd>>,
    failed: Cell<bool>,
    sigbus_impossible: bool,
    data: *const [Cell<u8>],
    cpu: Option<Rc<CpuWorker>>,
}

#[derive(Clone)]
pub struct ClientMemOffset {
    mem: Rc<ClientMem>,
    offset: usize,
    data: *const [Cell<u8>],
}

impl ClientMem {
    pub fn new(
        fd: &Rc<OwnedFd>,
        len: usize,
        read_only: bool,
        client: Option<&Client>,
        cpu: Option<&Rc<CpuWorker>>,
        is_udmabuf: bool,
    ) -> Result<Self, ClientMemError> {
        Self::new2(fd, len, read_only, client, cpu, c::MAP_SHARED, is_udmabuf)
    }

    pub fn new_private(
        fd: &Rc<OwnedFd>,
        len: usize,
        read_only: bool,
        client: Option<&Client>,
        cpu: Option<&Rc<CpuWorker>>,
    ) -> Result<Self, ClientMemError> {
        Self::new2(fd, len, read_only, client, cpu, c::MAP_PRIVATE, false)
    }

    fn new2(
        fd: &Rc<OwnedFd>,
        len: usize,
        read_only: bool,
        client: Option<&Client>,
        cpu: Option<&Rc<CpuWorker>>,
        flags: c::c_int,
        is_udmabuf: bool,
    ) -> Result<Self, ClientMemError> {
        let mut sigbus_impossible = is_udmabuf;
        let mut real_size = None;
        if !sigbus_impossible
            && let Ok(seals) = uapi::fcntl_get_seals(fd.raw())
            && seals & c::F_SEAL_SHRINK != 0
            && let Ok(stat) = uapi::fstat(fd.raw())
        {
            real_size = Some(stat.st_size as usize);
            sigbus_impossible = stat.st_size as u64 >= len as u64;
        }
        if !sigbus_impossible && let Some(client) = client {
            log::debug!(
                "Client {} ({}) has created a shm buffer that might cause SIGBUS",
                client.pid_info.comm,
                client.id,
            );
        }
        let len = len.next_multiple_of(page_size());
        if let Some(real_size) = real_size
            && real_size < len
        {
            let _ = ftruncate(fd.raw(), len as _);
        }
        let data = if len == 0 {
            &mut [][..]
        } else {
            let prot = match read_only {
                true => c::PROT_READ,
                false => c::PROT_READ | c::PROT_WRITE,
            };
            unsafe {
                let data = c::mmap64(ptr::null_mut(), len, prot, flags, fd.raw(), 0);
                if data == c::MAP_FAILED {
                    return Err(ClientMemError::MmapFailed(uapi::Errno::default().into()));
                }
                std::slice::from_raw_parts_mut(data as *mut Cell<u8>, len)
            }
        };
        Ok(Self {
            fd: ManuallyDrop::new(fd.clone()),
            failed: Cell::new(false),
            sigbus_impossible,
            data,
            cpu: cpu.cloned(),
        })
    }

    pub fn len(&self) -> usize {
        self.data.len()
    }

    pub fn offset(self: &Rc<Self>, offset: usize, len: usize) -> ClientMemOffset {
        let mem = unsafe { &*self.data };
        ClientMemOffset {
            mem: self.clone(),
            offset,
            data: &mem[offset..][..len],
        }
    }

    pub fn fd(&self) -> &Rc<OwnedFd> {
        &self.fd
    }

    pub fn is_sealed_memfd(&self) -> bool {
        self.sigbus_impossible
    }
}

impl ClientMemOffset {
    pub fn pool(&self) -> &ClientMem {
        &self.mem
    }

    pub fn offset(&self) -> usize {
        self.offset
    }

    #[expect(dead_code)]
    pub fn ptr(&self) -> *const [Cell<u8>] {
        self.data
    }

    pub fn access<T, F: FnOnce(&[Cell<u8>]) -> T>(&self, f: F) -> Result<T, ClientMemError> {
        unsafe {
            if self.mem.sigbus_impossible {
                return Ok(f(&*self.data));
            }
            let mref = MemRef {
                mem: &*self.mem,
                outer: MEM.get(),
            };
            MEM.set(&mref);
            compiler_fence(Ordering::SeqCst);
            let res = f(&*self.data);
            MEM.set(mref.outer);
            compiler_fence(Ordering::SeqCst);
            match self.mem.failed.get() {
                true => Err(ClientMemError::Sigbus),
                _ => Ok(res),
            }
        }
    }

    pub fn read<T: Pod>(&self, dst: &mut Vec<T>) -> Result<(), ClientMemError> {
        if self.data.len().checked_rem(std::mem::size_of::<T>()) != Some(0) {
            return Err(ClientMemError::InvalidLength);
        }
        self.access(|v| {
            let len_elements = v.len() / std::mem::size_of::<T>();
            dst.reserve(len_elements);
            let (_, unused) = dst.split_at_spare_mut_bytes_ext();
            unused[..v.len()].copy_from_slice(uapi::as_maybe_uninit_bytes(v));
            unsafe {
                dst.set_len(dst.len() + len_elements);
            }
        })
    }
}

impl Drop for ClientMem {
    fn drop(&mut self) {
        let fd = unsafe { ManuallyDrop::take(&mut self.fd) };
        if let Some(cpu) = &self.cpu {
            let pending = cpu.submit(Box::new(CloseMemWork {
                fd: Rc::try_unwrap(fd).ok(),
                data: self.data,
            }));
            pending.detach();
        } else {
            unsafe {
                c::munmap(self.data as _, self.len());
            }
        }
    }
}

struct MemRef {
    mem: *const ClientMem,
    outer: *const MemRef,
}

thread_local! {
    static MEM: Cell<*const MemRef> = const { Cell::new(ptr::null()) };
}

unsafe fn kill() -> ! {
    unsafe {
        c::signal(c::SIGBUS, c::SIG_DFL);
        raise(c::SIGBUS);
    }
    unreachable!();
}

unsafe extern "C" fn sigbus(sig: i32, info: &c::siginfo_t, _ucontext: *mut c::c_void) {
    unsafe {
        assert_eq!(sig, c::SIGBUS);
        let mut memr_ptr = MEM.get();
        while !memr_ptr.is_null() {
            let memr = &*memr_ptr;
            let mem = &*memr.mem;
            let lo = mem.data as *mut u8 as usize;
            let hi = lo + mem.len();
            let fault_addr = info.si_addr() as usize;
            if fault_addr < lo || fault_addr >= hi {
                memr_ptr = memr.outer;
                continue;
            }
            let res = c::mmap64(
                lo as _,
                hi - lo,
                c::PROT_WRITE | c::PROT_READ,
                c::MAP_ANONYMOUS | c::MAP_PRIVATE | c::MAP_FIXED,
                -1,
                0,
            );
            if res == c::MAP_FAILED {
                kill();
            }
            mem.failed.set(true);
            return;
        }
        kill();
    }
}

pub fn init() -> Result<(), ClientMemError> {
    unsafe {
        let mut action: c::sigaction = MaybeUninit::zeroed().assume_init();
        action.sa_sigaction =
            sigbus as unsafe extern "C" fn(i32, &c::siginfo_t, *mut c::c_void) as _;
        action.sa_flags = c::SA_NODEFER | c::SA_SIGINFO;
        let res = c::sigaction(c::SIGBUS, &action, ptr::null_mut());
        match uapi::map_err!(res) {
            Ok(_) => Ok(()),
            Err(e) => Err(ClientMemError::SigactionFailed(e.into())),
        }
    }
}

struct CloseMemWork {
    fd: Option<OwnedFd>,
    data: *const [Cell<u8>],
}

unsafe impl Send for CloseMemWork {}

impl CpuJob for CloseMemWork {
    fn work(&mut self) -> &mut dyn CpuWork {
        self
    }

    fn completed(self: Box<Self>) {
        // nothing
    }
}

impl CpuWork for CloseMemWork {
    fn run(&mut self) -> Option<Box<dyn AsyncCpuWork>> {
        zone!("CloseMemWork");
        self.fd.take();
        unsafe {
            c::munmap(self.data as _, self.data.len());
        }
        None
    }
}

impl ShmMemory for ClientMemOffset {
    fn len(&self) -> usize {
        self.data.len()
    }

    fn safe_access(&self) -> ShmMemoryBacking {
        match self.mem.is_sealed_memfd() {
            true => ShmMemoryBacking::Ptr(self.data),
            false => ShmMemoryBacking::Fd(self.mem.fd.deref().clone(), self.offset),
        }
    }

    fn access(&self, f: &mut dyn FnMut(&[Cell<u8>])) -> Result<(), Box<dyn Error + Sync + Send>> {
        self.access(f).map_err(|e| e.into())
    }
}