starry-kernel 0.9.0

//! Software perf event with BPF program attachment + ringbuf output.
//!
//! The user-visible ringbuf is created on the first `mmap(perf_fd, ...)`
//! call: `BpfPerfEventWrapper::device_mmap` allocates `1 + 2^N` physically
//! contiguous 4 K pages (header page + power-of-two-page data ring) and
//! hands the kernel virtual address to `BpfPerfEvent::do_mmap`, which
//! initializes `perf_event_mmap_page` in page 0. `sys_mmap` then maps the
//! same physical range into the caller's address space, so user reads of
//! `data_head` / `data_tail` and kernel writes via `bpf_perf_event_output`
//! share one buffer.

use alloc::sync::{Arc, Weak};
use core::{
    any::Any,
    fmt::Debug,
    sync::atomic::{AtomicBool, Ordering},
};

use ax_alloc::GlobalPage;
use ax_hal::mem::virt_to_phys;
use ax_memory_addr::{PAGE_SIZE_4K, PhysAddr};
use axpoll::{IoEvents, Pollable};
use axpoll_set::PollSet;
use kbpf_basic::{
    linux_bpf::{perf_event_mmap_page, perf_event_sample_format},
    perf::{PerfProbeArgs, bpf::BpfPerfEvent},
};
use kprobe::PtRegs;
use rbpf::EbpfVmRaw;

use super::PerfEventOps;
#[cfg(target_arch = "x86_64")]
use crate::perf::BPFJitMemory;
use crate::{
    StarryError, StarryResult,
    ebpf::{BPF_HELPER_FUN_SET, error::BpfResultExt, prog::BpfProg},
    file::FileLike,
    sync::IrqMutex,
    task::future::IrqNotify,
};

/// Number of 4K pages reserved for x86_64 BPF JIT executable memory.
/// Each JIT-compiled eBPF program fits within this space; the allocation
/// is sized generously (16 KiB) since programs are typically < 1 page.
#[cfg(target_arch = "x86_64")]
const BPF_JIT_MEM_PAGES: usize = 4;

struct BpfPerfEventState {
    inner: BpfPerfEvent,
    /// Weak handle to the contiguous pages backing the ringbuf. The strong
    /// ref(s) live in the user VMA(s); `strong_count() > 0` means a live
    /// mapping still exists.
    pages: Option<Weak<GlobalPage>>,
}

impl BpfPerfEventState {
    fn is_mapped(&self) -> bool {
        self.pages
            .as_ref()
            .is_some_and(|pages| pages.strong_count() > 0)
    }
}

/// Non-sleeping output capability used by `bpf_perf_event_output`.
///
/// The task control plane owns allocation and mapping. This endpoint only
/// enters the bounded ring write, observes the already-published page anchor,
/// and emits an IRQ-safe worker notification.
#[derive(Clone)]
pub(super) struct BpfPerfOutput {
    state: Arc<IrqMutex<BpfPerfEventState>>,
    poll_notify: Arc<IrqNotify>,
}

/// Task-context readiness capability separated from mutable perf control.
#[derive(Clone)]
pub(super) struct BpfPerfPoll {
    state: Arc<IrqMutex<BpfPerfEventState>>,
    poll_ready: Arc<PollSet>,
}

impl Pollable for BpfPerfPoll {
    fn poll(&self) -> IoEvents {
        if self.state.lock().inner.readable() {
            IoEvents::IN
        } else {
            IoEvents::empty()
        }
    }

    unsafe fn register_shared(
        &self,
        sink: &mut dyn axpoll::SharedRegistrationSink,
        events: IoEvents,
    ) {
        if events.contains(IoEvents::IN) {
            unsafe { sink.register_shared(&self.poll_ready, IoEvents::IN) };
        }
    }

    unsafe fn register_exclusive(
        &self,
        sink: &mut dyn axpoll::ExclusiveRegistrationSink,
        events: IoEvents,
    ) {
        if events.contains(IoEvents::IN) {
            unsafe { sink.register_exclusive(&self.poll_ready, IoEvents::IN) };
        }
    }
}

impl BpfPerfOutput {
    pub(super) fn write_event(&self, data: &[u8]) -> StarryResult<()> {
        let notify = {
            let mut state = self.state.lock();
            if !state.is_mapped() {
                return Ok(());
            }
            state.inner.write_event(data).into_starry_result()?;
            state.inner.enabled()
        };
        if notify {
            self.poll_notify.notify_irq();
        }
        Ok(())
    }
}

/// Wraps `kbpf_basic::perf::bpf::BpfPerfEvent` with separate task-control and
/// non-sleeping output state plus a poll set so readers can wait for records.
///
/// Ownership model: the user VMA owns the ringbuf pages via the strong
/// `Arc<GlobalPage>` threaded into `DeviceMmap::Physical`'s retainer slot;
/// the shared output state keeps only a `Weak`. Consequences:
///
/// * UAF safety — the pages outlive `close(perf_fd)` (which drops this
///   wrapper) for as long as a mapping is live, because the VMA holds the
///   strong ref. A userspace read after closing the fd never observes
///   freed memory.
/// * Self-cleaning allocation — if a `device_mmap` result is never adopted
///   by a surviving VMA (a non-direct mmap path, a permission/address
///   error, or an `aspace.map` failure), the lone strong ref drops, the
///   frames free, and `is_mapped` flips back to false so the perf fd can
///   be mmap'd again instead of being wedged in `ResourceBusy`. After a
///   normal `munmap` the same thing happens, matching Linux's allowance to
///   re-`mmap` a perf fd.
///
/// The inner perf event holds a raw pointer into the page buffer; `RingPage` has no
/// destructor and is never dereferenced once the pages are gone (every
/// access is gated on [`BpfPerfEventState::is_mapped`]), so a dangling pointer
/// left after the pages free is harmless.
pub struct BpfPerfEventWrapper {
    state: Arc<IrqMutex<BpfPerfEventState>>,
    poll: BpfPerfPoll,
    poll_notify: Arc<IrqNotify>,
    poll_alive: Arc<AtomicBool>,
}

impl BpfPerfEventWrapper {
    /// Construct the wrapper around a freshly-built `BpfPerfEvent`.
    pub fn new(inner: BpfPerfEvent) -> Self {
        let poll_ready = Arc::new(PollSet::new());
        let poll_notify = Arc::new(IrqNotify::new());
        let poll_alive = Arc::new(AtomicBool::new(true));
        start_bpf_perf_notify_worker(poll_ready.clone(), poll_notify.clone(), poll_alive.clone());
        let state = Arc::new(IrqMutex::new(BpfPerfEventState { inner, pages: None }));
        Self {
            poll: BpfPerfPoll {
                state: Arc::clone(&state),
                poll_ready,
            },
            state,
            poll_notify,
            poll_alive,
        }
    }

    pub(super) fn output_handle(&self) -> BpfPerfOutput {
        BpfPerfOutput {
            state: Arc::clone(&self.state),
            poll_notify: Arc::clone(&self.poll_notify),
        }
    }

    pub(super) fn poll_handle(&self) -> BpfPerfPoll {
        self.poll.clone()
    }
}

impl Drop for BpfPerfEventWrapper {
    fn drop(&mut self) {
        self.poll_alive.store(false, Ordering::Release);
        self.poll_notify.notify();
    }
}

fn start_bpf_perf_notify_worker(
    poll_ready: Arc<PollSet>,
    poll_notify: Arc<IrqNotify>,
    poll_alive: Arc<AtomicBool>,
) {
    crate::task::spawn_kernel_thread(
        move || loop {
            poll_notify.wait();
            if !poll_alive.load(Ordering::Acquire) {
                break;
            }
            // Ring data is written before the deferred poll wake.
            unsafe { poll_ready.wake(IoEvents::IN) };
        },
        "bpf-perf-notify".into(),
    );
}

impl Debug for BpfPerfEventWrapper {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "BpfPerfEventWrapper")
    }
}

impl PerfEventOps for BpfPerfEventWrapper {
    fn enable(&mut self) -> StarryResult<()> {
        self.state.lock().inner.enable().into_starry_result()?;
        Ok(())
    }

    fn disable(&mut self) -> StarryResult<()> {
        self.state.lock().inner.disable().into_starry_result()?;
        Ok(())
    }

    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    fn device_mmap(&mut self, len: usize) -> StarryResult<(PhysAddr, Arc<dyn Any + Send + Sync>)> {
        if self.state.lock().is_mapped() {
            // Linux allows only one live mmap per perf event fd; a second
            // mapping while the first is alive would orphan it. A stale
            // `Weak` from an abandoned or munmap'd previous attempt does not
            // count (its pages are already freed), so the fd stays mmap-able.
            return Err(StarryError::ResourceBusy);
        }
        // libbpf requires `(1 + 2^N) * PAGE_SIZE` so the data region is a
        // power of two pages; `RingPage::init` enforces ≥ 2 pages total and
        // 4 K alignment. Reject anything that would trip those asserts.
        if len == 0 || !len.is_multiple_of(PAGE_SIZE_4K) {
            return Err(StarryError::InvalidInput);
        }
        let num_pages = len / PAGE_SIZE_4K;
        if num_pages < 2 || !(num_pages - 1).is_power_of_two() {
            return Err(StarryError::InvalidInput);
        }
        let mut pages = GlobalPage::alloc_contiguous(num_pages, PAGE_SIZE_4K)
            .map_err(|_| StarryError::NoMemory)?;
        pages.zero();
        let kvirt = pages.start_vaddr();
        let paddr = virt_to_phys(kvirt);
        let pages = Arc::new(pages);

        let mut state = self.state.lock();
        if state.is_mapped() {
            return Err(StarryError::ResourceBusy);
        }
        state
            .inner
            .do_mmap(kvirt.as_usize(), len, 0)
            .map_err(|_| StarryError::InvalidInput)?;
        // kbpf_basic::RingPage::init sets the data-region geometry but leaves
        // version at 0. perf checks `perf_event_mmap_page.version == 1` and
        // rejects 0 (`perf_mmap__is_mmap_ok`), so we must set it here.
        let header = kvirt.as_usize() as *mut perf_event_mmap_page;
        // SAFETY: header points at the freshly allocated header page, no reader yet.
        unsafe {
            core::ptr::addr_of_mut!((*header).version).write(1);
            core::ptr::addr_of_mut!((*header).compat_version).write(0);
        }
        // Keep only a `Weak`; hand the sole strong ref to the caller, which
        // threads it into `DeviceMmap::Physical`'s retainer so the user VMA
        // pins these frames until `munmap`/exit even if the perf fd (and this
        // wrapper) is closed first. Because the wrapper does not retain a
        // strong ref, an mmap that is abandoned or fails before a VMA adopts
        // the anchor simply frees the pages and leaves the fd mmap-able again
        // (see the type-level docs). Without the anchor the pages would free
        // under a live mapping.
        state.pages = Some(Arc::downgrade(&pages));
        drop(state);
        let anchor: Arc<dyn Any + Send + Sync> = pages;
        Ok((paddr, anchor))
    }
}

impl Pollable for BpfPerfEventWrapper {
    fn poll(&self) -> axpoll::IoEvents {
        self.poll.poll()
    }

    unsafe fn register_shared(
        &self,
        sink: &mut dyn axpoll::SharedRegistrationSink,
        events: axpoll::IoEvents,
    ) {
        unsafe { self.poll.register_shared(sink, events) };
    }

    unsafe fn register_exclusive(
        &self,
        sink: &mut dyn axpoll::ExclusiveRegistrationSink,
        events: axpoll::IoEvents,
    ) {
        unsafe { self.poll.register_exclusive(sink, events) };
    }
}

/// Build a `BpfPerfEventWrapper` from `perf_event_open` args. The upstream
/// code asserts `sample_type == PERF_SAMPLE_RAW`; we keep that assertion
/// to match the verifier contract and surface bad input early.
pub fn perf_event_open_bpf(args: PerfProbeArgs) -> BpfPerfEventWrapper {
    debug_assert_eq!(
        args.sample_type,
        Some(perf_event_sample_format::PERF_SAMPLE_RAW)
    );
    BpfPerfEventWrapper::new(BpfPerfEvent::new(args))
}

/// A loaded BPF program bundled with an `rbpf` interpreter that borrows
/// into the program's instruction buffer.
///
/// Soundness: the interpreter holds `'static`-typed borrows into resources
/// owned by this struct. Field order is therefore load-bearing: `vm` must be
/// declared first so its destructor runs before the JIT memory and program
/// instruction buffer are released. Do not reorder the fields.
pub struct OwnedEbpfVm {
    vm: EbpfVmRaw<'static>,
    #[cfg(target_arch = "x86_64")]
    /// MUST be declared after `vm` (drop order). Keeps the JIT executable
    /// memory alive for the entire lifetime of `vm`.
    _jit_exec_memory: BPFJitMemory,
    /// MUST be declared after `vm` (drop order). Keeps the instruction
    /// buffer alive for the entire lifetime of `vm`.
    _prog: Arc<BpfProg>,
}

impl OwnedEbpfVm {
    /// Build an `rbpf::EbpfVmRaw` around the program's instruction stream
    /// and register the kernel helper table on it. The returned value owns
    /// both the VM and the [`Arc<BpfProg>`] backing its instruction buffer.
    pub fn new(bpf_prog: Arc<dyn FileLike>) -> StarryResult<Self> {
        let prog = bpf_prog
            .into_any_arc()
            .downcast::<BpfProg>()
            .map_err(|_| StarryError::InvalidInput)?;
        // Extend the borrow of `prog.insns()` to `'static`. SAFETY: the
        // Arc<BpfProg> is moved into the returned `OwnedEbpfVm` together
        // with the VM, and the struct's field drop order (vm before _prog)
        // guarantees the borrower is destroyed before the buffer is freed.
        let prog_slice = prog.insns();
        let prog_slice =
            unsafe { core::slice::from_raw_parts(prog_slice.as_ptr(), prog_slice.len()) };
        let mut vm = EbpfVmRaw::new(Some(prog_slice)).map_err(|e| {
            error!("rbpf::EbpfVmRaw::new failed: {e:?}");
            StarryError::InvalidInput
        })?;

        if let Some(table) = BPF_HELPER_FUN_SET.get() {
            for (key, value) in table.iter() {
                let _ = vm.register_helper(*key, *value);
            }
        }
        // TODO: not all of the address space is accessible to a BPF program;
        // allowing the full `0..u64::MAX` range disables rbpf's bounds check
        // and lets a buggy/hostile program read arbitrary kernel memory via
        // direct loads.
        //
        // FIXME: narrow this to the legitimately-reachable context /
        // map / stack ranges once kbpf-basic exposes per-program
        // bounds.
        vm.register_allowed_memory(0..u64::MAX);

        #[cfg(target_arch = "x86_64")]
        {
            // TODO: calculate a more precise size.
            let mut jit_exec_memory = BPFJitMemory::new(BPF_JIT_MEM_PAGES)?;
            // SAFETY: `jit_exec_memory` is moved into the returned
            // `OwnedEbpfVm` after `vm`; field drop order guarantees `vm` is
            // destroyed before the executable mapping is unmapped.
            let jit_slice = unsafe { jit_exec_memory.as_static_mut_slice() };
            vm.set_jit_exec_memory(jit_slice).map_err(|e| {
                error!("rbpf::EbpfVmRaw::set_jit_exec_memory failed: {e:?}");
                StarryError::InvalidInput
            })?;

            vm.jit_compile().map_err(|e| {
                error!("rbpf::EbpfVmRaw::jit_compile failed: {e:?}");
                StarryError::InvalidInput
            })?;

            Ok(Self {
                vm,
                _jit_exec_memory: jit_exec_memory,
                _prog: prog,
            })
        }

        #[cfg(not(target_arch = "x86_64"))]
        {
            Ok(Self { vm, _prog: prog })
        }
    }

    /// Execute the wrapped BPF program with the supplied context bytes.
    ///
    /// Takes `&self`: `rbpf::EbpfVmRaw::execute_program` is itself `&self`
    /// (the interpreter keeps its scratch state on the local stack), so no
    /// exterior mutability — and therefore no lock — is required around an
    /// `OwnedEbpfVm`.
    pub fn execute_program(&self, ctx: &mut [u8]) -> Result<u64, rbpf::lib::Error> {
        #[cfg(not(target_arch = "x86_64"))]
        {
            self.vm.execute_program(ctx)
        }
        #[cfg(target_arch = "x86_64")]
        {
            unsafe { self.vm.execute_program_jit(ctx) }
        }
    }

    /// Execute the wrapped BPF program with a `PtRegs` as the single-pointer
    /// context argument the kprobe/kretprobe ABI expects.
    pub fn execute_with_ptregs(&self, pt_regs: &mut PtRegs) -> Result<u64, rbpf::lib::Error> {
        // SAFETY: kbpf-basic's kprobe-context contract passes a raw
        // pointer to `PtRegs` as the program context; we hand the same
        // bytes here.
        let probe_context = unsafe {
            core::slice::from_raw_parts_mut(
                pt_regs as *mut PtRegs as *mut u8,
                core::mem::size_of::<PtRegs>(),
            )
        };
        self.execute_program(probe_context)
    }
}

impl Debug for OwnedEbpfVm {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "OwnedEbpfVm")
    }
}

// SAFETY: the bundled `EbpfVmRaw<'static>` is an interpreter over an immutable
// instruction slice; the `Arc<BpfProg>` backing that slice is `Send + Sync`.
// `execute_program` runs entirely off `&self` and a private stack, so it is
// re-entrant and may be driven concurrently from probe-fire paths on several
// CPUs without data races. The raw-pointer fields rbpf carries internally are
// never mutated after construction, so promoting the bundle to `Send + Sync`
// is sound.
unsafe impl Send for OwnedEbpfVm {}
unsafe impl Sync for OwnedEbpfVm {}