Supports tracepoint-based eBPF applications.

Signed-off-by: Godones <chenlinfeng25@outlook.com>

Author: Godones

Cargo.lock

@@ -75,7 +75,7 @@ tracepoint.workspace = true
 static-keys.workspace = true
 
 # kbpf-basic = { path = "../../ax-project/kbpf-basic/" }
-kbpf-basic = { git = "https://github.com/Starry-OS/kbpf-basic" }
+kbpf-basic = { git = "https://github.com/Starry-OS/kbpf-basic", rev = "ef8ade5" }
 rbpf = { git = "https://github.com/qmonnet/rbpf.git", default-features = false, rev = "f31e471" }
 
 [target.'cfg(not(target_arch = "loongarch64"))'.dependencies]

api/Cargo.toml

@@ -63,6 +63,7 @@ impl FileLike for BpfMap {
 
 pub fn bpf_map_create(attr: &bpf_attr) -> AxResult<isize> {
     let map_meta = BpfMapMeta::try_from(attr).map_err(bpferror_to_axerr)?;
+    axlog::info!("The map attr is {:#?}", map_meta);
     let unified_map =
         kbpf_basic::map::bpf_map_create::<PerCpuImpl>(map_meta).map_err(bpferror_to_axerr);
     if let Err(e) = &unified_map {

api/src/bpf/map.rs

@@ -96,7 +96,6 @@ impl KernelAuxiliaryOps for EbpfKernelAuxiliary {
         let ret = func(&mut *unified_map);
         drop(unified_map);
         // avoid double free
-        // log::error!("get_unified_map_from_ptr: ret: {:?}", ret);
         let _ = Arc::into_raw(map);
         ret
     }
@@ -105,15 +104,13 @@ impl KernelAuxiliaryOps for EbpfKernelAuxiliary {
     where
         F: FnOnce(&mut UnifiedMap) -> kbpf_basic::Result<R>,
     {
-        // log::error!("get_unified_map_from_fd: map_fd: {}", map_fd);
         let file = get_file_like(map_fd as _).map_err(|_| BpfError::NotFound)?;
         let bpf_map = file.into_any().downcast::<BpfMap>().unwrap();
         let unified_map = &mut bpf_map.unified_map();
         func(unified_map)
     }
 
     fn get_unified_map_ptr_from_fd(map_fd: u32) -> kbpf_basic::Result<*const u8> {
-        // log::error!("get_unified_map_ptr_from_fd: map_fd: {}", map_fd);
         let file = get_file_like(map_fd as _).map_err(|_| BpfError::NotFound)?;
         let bpf_map = file.into_any().downcast::<BpfMap>().unwrap();
         let map_ptr = Arc::into_raw(bpf_map) as usize;
@@ -149,7 +146,6 @@ impl KernelAuxiliaryOps for EbpfKernelAuxiliary {
 
     fn string_from_user_cstr(ptr: *const u8) -> kbpf_basic::Result<String> {
         let str = vm_load_string(ptr).map_err(|_| BpfError::InvalidArgument)?;
-        // axlog::info!("string_from_user_cstr: string: {:?}", str);
         Ok(str)
     }
 

api/src/bpf/tansform.rs

@@ -36,6 +36,7 @@ pub fn init() {
 
     tracepoint::tracepoint_init();
     bpf::init_bpf();
+    perf::perf_event_init();
 
     if axconfig::plat::CPU_NUM > 1 {
         panic!("SMP is not supported");

api/src/lib.rs

@@ -1,32 +1,38 @@
-use alloc::sync::Arc;
 use core::{any::Any, convert::Into, fmt::Debug};
 
 use axerrno::AxResult;
 use axhal::paging::PageSize;
 use axio::{IoEvents, PollSet, Pollable};
-use axmm::backend::SharedPages;
+use axmm::backend::{alloc_frames, dealloc_frames};
 use kbpf_basic::{
     linux_bpf::perf_event_sample_format,
     perf::{PerfProbeArgs, bpf::BpfPerfEvent},
 };
+use memory_addr::PhysAddr;
 
 use super::PerfEventOps;
 
 pub struct BpfPerfEventWrapper {
     inner: BpfPerfEvent,
     poll_ready: PollSet,
+    phys_addr: Option<(PhysAddr, usize)>,
 }
 
 impl BpfPerfEventWrapper {
     pub fn new(inner: BpfPerfEvent) -> Self {
         BpfPerfEventWrapper {
             inner,
             poll_ready: PollSet::new(),
+            phys_addr: None,
         }
     }
 
     pub fn write_event(&mut self, data: &[u8]) -> AxResult<()> {
         // TODO: remove unwrap
+        if self.phys_addr.is_none() {
+            axlog::warn!("BpfPerfEventWrapper: first write_event, mmap not done yet");
+            return Ok(());
+        }
         self.inner.write_event(data).unwrap();
         if self.inner.enabled() {
             self.poll_ready.wake();
@@ -71,24 +77,42 @@ impl PerfEventOps for BpfPerfEventWrapper {
         length: usize,
         prot: crate::syscall::MmapProt,
         flags: crate::syscall::MmapFlags,
-        _offset: usize,
+        offset: usize,
     ) -> AxResult<isize> {
-        axlog::warn!(
+        axlog::info!(
             "BpfPerfEventWrapper::mmap prot:{:?} flags:{:?}",
             prot,
             flags
         );
-        let back = axmm::backend::Backend::new_shared(
-            start,
-            Arc::new(SharedPages::new(length, PageSize::Size4K)?),
-        );
-        aspace.map(start, length, prot.into(), true, back)?;
 
-        self.inner.do_mmap(start.as_usize(), length, 0).unwrap();
+        let phys_addr = alloc_frames(
+            true,
+            PageSize::Size4K,
+            length / PageSize::Size4K as usize,
+            axalloc::UsageKind::PageCache,
+        )?;
+        let page_virt = axhal::mem::phys_to_virt(phys_addr);
+
+        aspace.map_linear(start, phys_addr, length, prot.into())?;
+
+        self.inner
+            .do_mmap(page_virt.as_usize(), length, offset)
+            .unwrap();
+
+        self.phys_addr = Some((phys_addr, length / PageSize::Size4K as usize));
+
         Ok(start.as_usize() as isize)
     }
 }
 
+impl Drop for BpfPerfEventWrapper {
+    fn drop(&mut self) {
+        if let Some((phys_addr, nums)) = self.phys_addr {
+            dealloc_frames(phys_addr, nums);
+        }
+    }
+}
+
 impl Pollable for BpfPerfEventWrapper {
     fn poll(&self) -> axio::IoEvents {
         if self.inner.readable() {

api/src/perf/bpf.rs

@@ -69,6 +69,27 @@ impl PerfEventOps for KprobePerfEvent {
     fn as_any_mut(&mut self) -> &mut dyn Any {
         self
     }
+
+    fn set_bpf_prog(&mut self, bpf_prog: Arc<dyn FileLike>) -> AxResult<()> {
+        let bpf_prog = bpf_prog.into_any().downcast::<BpfProg>().unwrap();
+        let prog = bpf_prog.insns();
+        let prog = unsafe { core::slice::from_raw_parts_mut(prog.as_ptr() as *mut u8, prog.len()) };
+        let mut vm = EbpfVmRaw::new(Some(prog)).unwrap();
+        for (key, value) in BPF_HELPER_FUN_SET.iter() {
+            vm.register_helper(*key, *value).unwrap();
+        }
+        vm.register_allowed_memory(0..u64::MAX);
+        static CALLBACK_ID: AtomicU32 = AtomicU32::new(0);
+
+        let id = CALLBACK_ID.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
+
+        // create a callback to execute the ebpf prog
+        let callback = Box::new(KprobePerfCallBack::new(bpf_prog, vm));
+        // update callback for kprobe
+        self.kprobe.register_event_callback(id, callback);
+        self.callback_list.push(id);
+        Ok(())
+    }
 }
 
 pub struct KprobePerfCallBack {
@@ -100,29 +121,6 @@ impl CallBackFunc for KprobePerfCallBack {
     }
 }
 
-impl KprobePerfEvent {
-    pub fn set_bpf_prog(&mut self, bpf_prog: Arc<dyn FileLike>) -> AxResult<()> {
-        let bpf_prog = bpf_prog.into_any().downcast::<BpfProg>().unwrap();
-        let prog = bpf_prog.insns();
-        let prog = unsafe { core::slice::from_raw_parts_mut(prog.as_ptr() as *mut u8, prog.len()) };
-        let mut vm = EbpfVmRaw::new(Some(prog)).unwrap();
-        for (key, value) in BPF_HELPER_FUN_SET.iter() {
-            vm.register_helper(*key, *value).unwrap();
-        }
-        vm.register_allowed_memory(0..u64::MAX);
-        static CALLBACK_ID: AtomicU32 = AtomicU32::new(0);
-
-        let id = CALLBACK_ID.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
-
-        // create a callback to execute the ebpf prog
-        let callback = Box::new(KprobePerfCallBack::new(bpf_prog, vm));
-        // update callback for kprobe
-        self.kprobe.register_event_callback(id, callback);
-        self.callback_list.push(id);
-        Ok(())
-    }
-}
-
 fn perf_probe_arg_to_kprobe_builder(args: &PerfProbeArgs) -> KprobeBuilder<KprobeAuxiliary> {
     let symbol = &args.name;
     let addr = ksym::addr_from_symbol(symbol).unwrap() as usize;

api/src/perf/kprobe.rs

@@ -1,21 +1,27 @@
 mod bpf;
 mod kprobe;
+mod tracepoint;
 
-use alloc::{boxed::Box, sync::Arc};
+use alloc::{
+    boxed::Box,
+    sync::{Arc, Weak},
+};
 use core::{any::Any, ffi::c_void, fmt::Debug};
 
 use axerrno::{AxError, AxResult};
 use axio::Pollable;
+use hashbrown::HashMap;
 use kbpf_basic::{
     linux_bpf::{perf_event_attr, perf_type_id},
     perf::{PerfEventIoc, PerfProbeArgs},
 };
 use kspin::{SpinNoPreempt, SpinNoPreemptGuard};
+use lazyinit::LazyInit;
 
 use crate::{
     bpf::tansform::EbpfKernelAuxiliary,
     file::{FileLike, Kstat, add_file_like, get_file_like},
-    perf::{bpf::BpfPerfEventWrapper, kprobe::KprobePerfEvent},
+    perf::bpf::BpfPerfEventWrapper,
 };
 
 pub trait PerfEventOps: Pollable + Send + Sync + Debug {
@@ -26,7 +32,9 @@ pub trait PerfEventOps: Pollable + Send + Sync + Debug {
     fn custom_mmap(&self) -> bool {
         false
     }
-
+    fn set_bpf_prog(&mut self, _bpf_prog: Arc<dyn FileLike>) -> AxResult<()> {
+        Err(AxError::OperationNotSupported)
+    }
     fn mmap(
         &mut self,
         _aspace: &mut axmm::AddrSpace,
@@ -104,11 +112,7 @@ impl FileLike for PerfEvent {
                 let file = get_file_like(bpf_prog_fd as _)?;
 
                 let mut event = self.event.lock();
-                let kprobe_event = event
-                    .as_any_mut()
-                    .downcast_mut::<KprobePerfEvent>()
-                    .ok_or(AxError::InvalidInput)?;
-                kprobe_event.set_bpf_prog(file)?;
+                event.set_bpf_prog(file)?;
             }
         }
         Ok(0)
@@ -143,7 +147,7 @@ pub fn perf_event_open(
     let args =
         PerfProbeArgs::try_from_perf_attr::<EbpfKernelAuxiliary>(attr, pid, cpu, group_fd, flags)
             .unwrap();
-    axlog::warn!("perf_event_process: {:#?}", args);
+    axlog::info!("perf_event_process: {:#?}", args);
     let event: Box<dyn PerfEventOps> = match args.type_ {
         // Kprobe
         // See /sys/bus/event_source/devices/kprobe/type
@@ -155,21 +159,49 @@ pub fn perf_event_open(
             let bpf_event = bpf::perf_event_open_bpf(args);
             Box::new(bpf_event)
         }
+        perf_type_id::PERF_TYPE_TRACEPOINT => {
+            let tracepoint_event = tracepoint::perf_event_open_tracepoint(args)?;
+            Box::new(tracepoint_event)
+        }
         _ => {
             unimplemented!("perf_event_process: unknown type: {:?}", args);
         }
     };
-    let event = Arc::new(PerfEvent::new(event));
-    let fd = add_file_like(event, false).map(|fd| fd as _);
-    fd
+    let event = Arc::new(PerfEvent::new(event)) as Arc<dyn FileLike>;
+    let fd = add_file_like(event.clone(), false).map(|fd| fd as _)?;
+
+    PERF_FILE
+        .get()
+        .unwrap()
+        .lock()
+        .insert(fd, Arc::downgrade(&event));
+
+    axlog::info!("perf_event_open: fd: {:?}", fd);
+    Ok(fd as _)
+}
+
+static PERF_FILE: LazyInit<SpinNoPreempt<HashMap<usize, Weak<dyn FileLike>>>> = LazyInit::new();
+
+pub fn perf_event_init() {
+    PERF_FILE.init_once(SpinNoPreempt::new(HashMap::new()));
 }
 
 pub fn perf_event_output(_ctx: *mut c_void, fd: usize, _flags: u32, data: &[u8]) -> AxResult<()> {
-    // axlog::error!("perf_event_output: fd: {}, data: {:?}", fd, data.len());
-    let file = get_file_like(fd as _)?;
+    let mut perf_file_map = PERF_FILE.get().unwrap().lock();
+    let perf_file_weak = perf_file_map.get(&fd).ok_or(AxError::NotFound)?;
+
+    let Some(file) = perf_file_weak.upgrade() else {
+        // The file has been dropped, remove the weak reference from the map
+        perf_file_map.remove(&fd);
+        return Err(AxError::NotFound);
+    };
+
     let bpf_event_file = file.into_any().downcast::<PerfEvent>().unwrap();
     let mut event = bpf_event_file.event();
-    let event = event.as_any_mut().downcast_mut::<BpfPerfEventWrapper>().unwrap();
+    let event = event
+        .as_any_mut()
+        .downcast_mut::<BpfPerfEventWrapper>()
+        .unwrap();
     event.write_event(data).unwrap();
     Ok(())
 }