starry_kernel/
kprobe.rs

1//! Kernel probe (kprobe) subsystem for StarryOS.
2//!
3//! This module provides dynamic tracing support by allowing breakpoint
4//! insertion at kernel function entry/return points. It integrates the
5//! [`kprobe`] crate with StarryOS kernel infrastructure.
6//!
7//! # Architecture Support
8//!
9//! All four supported architectures are enabled: x86_64, riscv64, aarch64,
10//! and loongarch64. Each architecture provides TrapFrame↔PtRegs register
11//! conversion to bridge the kernel's trap frame format with the kprobe
12//! crate's portable `PtRegs` type.
13//!
14//! # Key Components
15//!
16//! - [`KernelKprobeOps`]: Platform-specific auxiliary operations for the kprobe crate
17//! - [`handle_breakpoint`]: Entry point for breakpoint exceptions (INT3/EBREAK/BRK)
18//! - [`handle_debug`]: Entry point for debug exceptions (x86_64 single-step only)
19
20use alloc::{sync::Arc, vec::Vec};
21
22use ax_kspin::{RawSpinNoIrq, SpinNoIrq};
23use ax_memory_addr::{MemoryAddr, PAGE_SIZE_4K, VirtAddr, VirtAddrRange};
24use ax_runtime::hal::{
25    cpu::{KernelTrapFrame, UserRegisters},
26    paging::MappingFlags,
27};
28use kprobe::{
29    KprobeAuxiliaryOps, KretprobeBuilder, ProbeBuilder, ProbePointList,
30    register_kprobe as kprobe_crate_register_kprobe,
31    register_kretprobe as kprobe_crate_register_kretprobe, retprobe::RetprobeInstance,
32    unregister_kprobe as kprobe_crate_unregister_kprobe,
33    unregister_kretprobe as kprobe_crate_unregister_kretprobe,
34};
35
36use crate::task::AsThread;
37
38/// Raw mutex used as the `L` type parameter for the `kprobe` crate's
39/// `ProbeManager` / `Kprobe` / `Kretprobe` (the perf subsystem refers to the
40/// concrete probe types parameterized on it — see [`KernelKprobe`] /
41/// [`KernelKretprobe`]).
42///
43/// Backed by [`ax_kspin::RawSpinNoIrq`], which disables kernel preemption and
44/// local IRQs across the critical section (`NoPreemptIrqSave` semantics, the
45/// same as the rest of the kernel's spin locks). This matters because the lock
46/// is taken on trap / kprobe-callback paths: a plain atomic spin lock that left
47/// preemption and IRQs enabled could be re-entered on the same CPU and would
48/// then deadlock spinning on a lock it already holds.
49pub type KernelRawMutex = RawSpinNoIrq;
50
51#[derive(Debug)]
52pub struct KernelKprobeOps;
53
54impl KprobeAuxiliaryOps for KernelKprobeOps {
55    fn copy_memory(src: *const u8, dst: *mut u8, len: usize, user_pid: Option<i32>) {
56        if let Some(pid) = user_pid {
57            // Uprobe arm/disarm reads the target process' original text bytes
58            // while the per-process kprobe manager spin-lock is held (IRQs
59            // disabled), so the faultable user-access path (`vm_read_slice`,
60            // which asserts IRQs enabled) cannot be used. Read through the
61            // *kernel* direct-map alias of the target page's physical frame
62            // instead — the same aliasing `set_writeable_for_address` uses to
63            // write. The text page is already resident (the loader executes the
64            // probed function before arming).
65            let task = crate::task::get_task(pid as _).expect("Failed to get task for uprobe");
66            let aspace = task.as_thread().proc_data.aspace();
67            let mm = aspace.lock();
68            let pt = mm.page_table();
69            let mut copied = 0;
70            while copied < len {
71                let vaddr = VirtAddr::from(src as usize + copied);
72                let Ok((paddr, ..)) = pt.query(vaddr) else {
73                    warn!(
74                        "kprobe copy_memory: user addr {:#x} not mapped",
75                        vaddr.as_usize()
76                    );
77                    return;
78                };
79                let page_off = vaddr.as_usize() & (PAGE_SIZE_4K - 1);
80                let chunk = core::cmp::min(len - copied, PAGE_SIZE_4K - page_off);
81                let kvaddr = ax_runtime::hal::mem::phys_to_virt(paddr);
82                unsafe {
83                    core::ptr::copy_nonoverlapping(kvaddr.as_ptr(), dst.add(copied), chunk);
84                }
85                copied += chunk;
86            }
87        } else {
88            unsafe {
89                core::ptr::copy_nonoverlapping(src, dst, len);
90            }
91        }
92    }
93
94    fn set_writeable_for_address<F: FnOnce(*mut u8)>(
95        address: usize,
96        len: usize,
97        user_pid: Option<i32>,
98        action: F,
99    ) {
100        if let Some(pid) = user_pid {
101            // User-space probe (uprobe): patch the target process' text by
102            // writing through the *kernel* direct-map alias of the page's
103            // physical frame. The user PTE keeps its read-only/exec flags
104            // untouched (no per-fire `protect` dance needed — uprobe single-step
105            // is out-of-line, see `alloc_user_exec_memory`). This runs at
106            // arm/disarm time (syscall context), so taking the sleeping aspace
107            // lock is fine. The instruction patch (≤ a few bytes) stays within
108            // the resolved page.
109            let task = crate::task::get_task(pid as _).expect("uprobe: target task gone");
110            let aspace = task.as_thread().proc_data.aspace();
111            let mm = aspace.lock();
112            let vaddr = VirtAddr::from(address);
113            let (paddr, ..) = mm
114                .page_table()
115                .query(vaddr)
116                .expect("uprobe: target address not mapped");
117            let kvaddr = ax_runtime::hal::mem::phys_to_virt(paddr);
118            action(kvaddr.as_mut_ptr());
119            ax_runtime::hal::cache::sync_kernel_text(vaddr.align_down_4k(), PAGE_SIZE_4K);
120            return;
121        }
122        let addr = VirtAddr::from(address);
123        crate::mm::patch_kernel_text(addr, len, action)
124            .expect("kprobe: set_writeable: patch kernel text failed");
125    }
126
127    fn alloc_kernel_exec_memory() -> *mut u8 {
128        let mut guard = ax_mm::kernel_aspace().lock();
129        let range = VirtAddrRange::new(guard.base(), guard.end());
130        let vaddr = guard
131            .find_free_area(guard.base(), PAGE_SIZE_4K, range)
132            .expect("kprobe: no free virtual address for exec memory");
133        guard
134            .map_alloc(
135                vaddr,
136                PAGE_SIZE_4K,
137                MappingFlags::READ | MappingFlags::WRITE | MappingFlags::EXECUTE,
138                true,
139            )
140            .expect("kprobe: map_alloc for exec memory failed");
141        vaddr.as_mut_ptr()
142    }
143
144    fn free_kernel_exec_memory(ptr: *mut u8) {
145        let vaddr = VirtAddr::from(ptr as usize);
146        let mut guard = ax_mm::kernel_aspace().lock();
147        guard
148            .unmap(vaddr, PAGE_SIZE_4K)
149            .expect("kprobe: unmap exec memory failed");
150    }
151
152    fn alloc_user_exec_memory<F: FnOnce(*mut u8)>(pid: Option<i32>, action: F) -> *mut u8 {
153        // Allocate one anonymous, user-executable page in the target process for
154        // out-of-line single-stepping (the displaced original instruction is
155        // copied here so the planted `int3` can stay armed). `action` writes
156        // that instruction through the kernel alias of the freshly-mapped frame.
157        let pid = pid.expect("uprobe: alloc_user_exec_memory needs a pid");
158        let task = crate::task::get_task(pid as _).expect("uprobe: target task gone");
159        let aspace = task.as_thread().proc_data.aspace();
160        let mut mm = aspace.lock();
161        let range = VirtAddrRange::new(mm.base(), mm.end());
162        let vaddr = mm
163            .find_free_area(mm.base(), PAGE_SIZE_4K, range, PAGE_SIZE_4K)
164            .expect("uprobe: no free user va for exec memory");
165        let backend = crate::mm::Backend::new_alloc(vaddr, PAGE_SIZE_4K, "uprobe-ols");
166        mm.map(
167            vaddr,
168            PAGE_SIZE_4K,
169            MappingFlags::READ | MappingFlags::EXECUTE | MappingFlags::USER,
170            true,
171            backend,
172        )
173        .expect("uprobe: map user exec memory failed");
174        let (paddr, ..) = mm
175            .page_table()
176            .query(vaddr)
177            .expect("uprobe: exec page not mapped after populate");
178        let kvaddr = ax_runtime::hal::mem::phys_to_virt(paddr);
179        action(kvaddr.as_mut_ptr());
180        ax_runtime::hal::cache::sync_kernel_text(vaddr, PAGE_SIZE_4K);
181        vaddr.as_mut_ptr()
182    }
183
184    fn free_user_exec_memory(pid: Option<i32>, ptr: *mut u8) {
185        let pid = pid.expect("uprobe: free_user_exec_memory needs a pid");
186        let task = crate::task::get_task(pid as _).expect("uprobe: target task gone");
187        let aspace = task.as_thread().proc_data.aspace();
188        let mut mm = aspace.lock();
189        mm.unmap(VirtAddr::from(ptr as usize), PAGE_SIZE_4K)
190            .expect("uprobe: unmap user exec memory failed");
191    }
192
193    fn insert_kretprobe_instance_to_task(instance: RetprobeInstance) {
194        let task = ax_task::current_may_uninit();
195        if let Some(task) = task {
196            let thread = task.try_as_thread();
197            if let Some(thread) = thread {
198                let mut kretprobe_instances = thread.kretprobe_stack.lock();
199                kretprobe_instances.push(instance);
200                return;
201            }
202        }
203        // If the current task is None, we can store it in a static variable
204        let mut instances = INSTANCE.lock();
205        instances.push(instance);
206    }
207
208    fn pop_kretprobe_instance_from_task() -> RetprobeInstance {
209        let task = ax_task::current_may_uninit();
210        if let Some(task) = task {
211            let thread = task.try_as_thread();
212            if let Some(thread) = thread {
213                let mut kretprobe_instances = thread.kretprobe_stack.lock();
214                return kretprobe_instances
215                    .pop()
216                    .expect("kretprobe instance stack underflow");
217            }
218        }
219        // If the current task is None, we can pop it from the static variable
220        let mut instances = INSTANCE.lock();
221        instances.pop().unwrap()
222    }
223}
224
225pub(crate) type KprobeManager = kprobe::ProbeManager<KernelRawMutex, KernelKprobeOps>;
226pub(crate) type KprobePointList = ProbePointList<KernelKprobeOps>;
227
228/// Concrete `kprobe::Kprobe` parameterized on the kernel's `RawMutex` and
229/// auxiliary ops, named to match what the perf module expects.
230pub type KernelKprobe = kprobe::Kprobe<KernelRawMutex, KernelKprobeOps>;
231/// Concrete `kprobe::Kretprobe`.
232pub type KernelKretprobe = kprobe::Kretprobe<KernelRawMutex, KernelKprobeOps>;
233/// The `KprobeAuxiliaryOps` impl, aliased under the name the perf module uses.
234pub type KprobeAuxiliary = KernelKprobeOps;
235
236static KPROBE_MANAGER: KprobeManager = KprobeManager::new();
237static KPROBE_POINT_LIST: SpinNoIrq<KprobePointList> = SpinNoIrq::new(KprobePointList::new());
238static INSTANCE: SpinNoIrq<Vec<RetprobeInstance>> = SpinNoIrq::new(Vec::new());
239
240fn with_manager<F, R>(f: F) -> R
241where
242    F: FnOnce(&KprobeManager) -> R,
243{
244    f(&KPROBE_MANAGER)
245}
246
247fn with_manager_and_list<F, R>(f: F) -> R
248where
249    F: FnOnce(&KprobeManager, &mut KprobePointList) -> R,
250{
251    let mut list = KPROBE_POINT_LIST.try_lock().unwrap();
252    f(&KPROBE_MANAGER, &mut list)
253}
254
255/// Register a kprobe into the global manager, returning the live handle.
256#[inline(never)]
257pub fn register_kprobe(builder: ProbeBuilder<KernelKprobeOps>) -> Arc<KernelKprobe> {
258    with_manager_and_list(|mgr, list| {
259        kprobe_crate_register_kprobe(mgr, list, builder).expect("Failed to register kprobe")
260    })
261}
262
263/// Unregister a previously registered kprobe.
264#[inline(never)]
265pub fn unregister_kprobe(kprobe: Arc<KernelKprobe>) {
266    with_manager_and_list(|mgr, list| kprobe_crate_unregister_kprobe(mgr, list, kprobe));
267}
268
269/// Register a kretprobe and return its live handle.
270#[inline(never)]
271pub fn register_kretprobe(builder: KretprobeBuilder<KernelRawMutex>) -> Arc<KernelKretprobe> {
272    with_manager_and_list(|mgr, list| {
273        kprobe_crate_register_kretprobe(mgr, list, builder).expect("Failed to register kretprobe")
274    })
275}
276
277/// Unregister a previously registered kretprobe.
278#[inline(never)]
279pub fn unregister_kretprobe(kretprobe: Arc<KernelKretprobe>) {
280    with_manager_and_list(|mgr, list| kprobe_crate_unregister_kretprobe(mgr, list, kretprobe));
281}
282
283pub(crate) fn trapframe_to_ptregs(tf: &UserRegisters) -> kprobe::PtRegs {
284    #[cfg(target_arch = "x86_64")]
285    {
286        kprobe::PtRegs {
287            r15: tf.r15 as usize,
288            r14: tf.r14 as usize,
289            r13: tf.r13 as usize,
290            r12: tf.r12 as usize,
291            rbp: tf.rbp as usize,
292            rbx: tf.rbx as usize,
293            r11: tf.r11 as usize,
294            r10: tf.r10 as usize,
295            r9: tf.r9 as usize,
296            r8: tf.r8 as usize,
297            rax: tf.rax as usize,
298            rcx: tf.rcx as usize,
299            rdx: tf.rdx as usize,
300            rsi: tf.rsi as usize,
301            rdi: tf.rdi as usize,
302            orig_rax: tf.vector as usize,
303            rip: tf.rip as usize,
304            cs: tf.cs as usize,
305            rflags: tf.rflags as usize,
306            rsp: tf.rsp as usize,
307            ss: tf.ss as usize,
308        }
309    }
310    #[cfg(target_arch = "riscv64")]
311    {
312        kprobe::PtRegs {
313            epc: tf.sepc,
314            ra: tf.regs.ra,
315            sp: tf.regs.sp,
316            gp: tf.regs.gp,
317            tp: tf.regs.tp,
318            t0: tf.regs.t0,
319            t1: tf.regs.t1,
320            t2: tf.regs.t2,
321            s0: tf.regs.s0,
322            s1: tf.regs.s1,
323            a0: tf.regs.a0,
324            a1: tf.regs.a1,
325            a2: tf.regs.a2,
326            a3: tf.regs.a3,
327            a4: tf.regs.a4,
328            a5: tf.regs.a5,
329            a6: tf.regs.a6,
330            a7: tf.regs.a7,
331            s2: tf.regs.s2,
332            s3: tf.regs.s3,
333            s4: tf.regs.s4,
334            s5: tf.regs.s5,
335            s6: tf.regs.s6,
336            s7: tf.regs.s7,
337            s8: tf.regs.s8,
338            s9: tf.regs.s9,
339            s10: tf.regs.s10,
340            s11: tf.regs.s11,
341            t3: tf.regs.t3,
342            t4: tf.regs.t4,
343            t5: tf.regs.t5,
344            t6: tf.regs.t6,
345            status: tf.sstatus.bits(),
346            badaddr: 0,
347            cause: 0,
348            orig_a0: tf.regs.a0,
349        }
350    }
351    #[cfg(target_arch = "aarch64")]
352    {
353        kprobe::PtRegs {
354            regs: tf.x,
355            sp: 0, // aarch64 SP is not saved in TrapFrame
356            pc: tf.elr,
357            pstate: tf.spsr,
358            orig_x0: tf.x[0],
359            syscallno: -1,
360            unused2: 0,
361        }
362    }
363    #[cfg(target_arch = "loongarch64")]
364    {
365        kprobe::PtRegs {
366            regs: [
367                tf.regs.zero,
368                tf.regs.ra,
369                tf.regs.tp,
370                tf.regs.sp,
371                tf.regs.a0,
372                tf.regs.a1,
373                tf.regs.a2,
374                tf.regs.a3,
375                tf.regs.a4,
376                tf.regs.a5,
377                tf.regs.a6,
378                tf.regs.a7,
379                tf.regs.t0,
380                tf.regs.t1,
381                tf.regs.t2,
382                tf.regs.t3,
383                tf.regs.t4,
384                tf.regs.t5,
385                tf.regs.t6,
386                tf.regs.t7,
387                tf.regs.t8,
388                tf.regs.u0,
389                tf.regs.fp,
390                tf.regs.s0,
391                tf.regs.s1,
392                tf.regs.s2,
393                tf.regs.s3,
394                tf.regs.s4,
395                tf.regs.s5,
396                tf.regs.s6,
397                tf.regs.s7,
398                tf.regs.s8,
399            ],
400            orig_a0: 0,
401            csr_era: tf.era,
402            csr_badvaddr: 0,
403            csr_crmd: 0,
404            csr_prmd: tf.prmd,
405            csr_euen: 0,
406            csr_ecfg: 0,
407            csr_estat: 0,
408        }
409    }
410}
411
412pub(crate) fn ptregs_write_back(pt: &kprobe::PtRegs, tf: &mut UserRegisters) {
413    #[cfg(target_arch = "x86_64")]
414    {
415        tf.r15 = pt.r15 as u64;
416        tf.r14 = pt.r14 as u64;
417        tf.r13 = pt.r13 as u64;
418        tf.r12 = pt.r12 as u64;
419        tf.rbp = pt.rbp as u64;
420        tf.rbx = pt.rbx as u64;
421        tf.r11 = pt.r11 as u64;
422        tf.r10 = pt.r10 as u64;
423        tf.r9 = pt.r9 as u64;
424        tf.r8 = pt.r8 as u64;
425        tf.rax = pt.rax as u64;
426        tf.rcx = pt.rcx as u64;
427        tf.rdx = pt.rdx as u64;
428        tf.rsi = pt.rsi as u64;
429        tf.rdi = pt.rdi as u64;
430        tf.rip = pt.rip as u64;
431        tf.cs = pt.cs as u64;
432        tf.vector = pt.orig_rax as u64;
433        tf.rflags = pt.rflags as u64;
434        tf.rsp = pt.rsp as u64;
435        tf.ss = pt.ss as u64;
436    }
437    #[cfg(target_arch = "riscv64")]
438    {
439        tf.sepc = pt.epc;
440        tf.regs.ra = pt.ra;
441        tf.regs.sp = pt.sp;
442        tf.regs.gp = pt.gp;
443        tf.regs.tp = pt.tp;
444        tf.regs.t0 = pt.t0;
445        tf.regs.t1 = pt.t1;
446        tf.regs.t2 = pt.t2;
447        tf.regs.s0 = pt.s0;
448        tf.regs.s1 = pt.s1;
449        tf.regs.a0 = pt.a0;
450        tf.regs.a1 = pt.a1;
451        tf.regs.a2 = pt.a2;
452        tf.regs.a3 = pt.a3;
453        tf.regs.a4 = pt.a4;
454        tf.regs.a5 = pt.a5;
455        tf.regs.a6 = pt.a6;
456        tf.regs.a7 = pt.a7;
457        tf.regs.s2 = pt.s2;
458        tf.regs.s3 = pt.s3;
459        tf.regs.s4 = pt.s4;
460        tf.regs.s5 = pt.s5;
461        tf.regs.s6 = pt.s6;
462        tf.regs.s7 = pt.s7;
463        tf.regs.s8 = pt.s8;
464        tf.regs.s9 = pt.s9;
465        tf.regs.s10 = pt.s10;
466        tf.regs.s11 = pt.s11;
467        tf.regs.t3 = pt.t3;
468        tf.regs.t4 = pt.t4;
469        tf.regs.t5 = pt.t5;
470        tf.regs.t6 = pt.t6;
471    }
472    #[cfg(target_arch = "aarch64")]
473    {
474        tf.x = pt.regs;
475        tf.elr = pt.pc;
476        tf.spsr = pt.pstate;
477    }
478    #[cfg(target_arch = "loongarch64")]
479    {
480        tf.regs.zero = pt.regs[0];
481        tf.regs.ra = pt.regs[1];
482        tf.regs.tp = pt.regs[2];
483        tf.regs.sp = pt.regs[3];
484        tf.regs.a0 = pt.regs[4];
485        tf.regs.a1 = pt.regs[5];
486        tf.regs.a2 = pt.regs[6];
487        tf.regs.a3 = pt.regs[7];
488        tf.regs.a4 = pt.regs[8];
489        tf.regs.a5 = pt.regs[9];
490        tf.regs.a6 = pt.regs[10];
491        tf.regs.a7 = pt.regs[11];
492        tf.regs.t0 = pt.regs[12];
493        tf.regs.t1 = pt.regs[13];
494        tf.regs.t2 = pt.regs[14];
495        tf.regs.t3 = pt.regs[15];
496        tf.regs.t4 = pt.regs[16];
497        tf.regs.t5 = pt.regs[17];
498        tf.regs.t6 = pt.regs[18];
499        tf.regs.t7 = pt.regs[19];
500        tf.regs.t8 = pt.regs[20];
501        tf.regs.u0 = pt.regs[21];
502        tf.regs.fp = pt.regs[22];
503        tf.regs.s0 = pt.regs[23];
504        tf.regs.s1 = pt.regs[24];
505        tf.regs.s2 = pt.regs[25];
506        tf.regs.s3 = pt.regs[26];
507        tf.regs.s4 = pt.regs[27];
508        tf.regs.s5 = pt.regs[28];
509        tf.regs.s6 = pt.regs[29];
510        tf.regs.s7 = pt.regs[30];
511        tf.regs.s8 = pt.regs[31];
512        tf.era = pt.csr_era;
513        tf.prmd = pt.csr_prmd;
514    }
515}
516
517pub fn handle_breakpoint(tf: &mut KernelTrapFrame<'_>) -> bool {
518    let mut updated = tf.snapshot();
519    let mut pt_regs = trapframe_to_ptregs(&updated);
520    let handled = with_manager(|manager| kprobe::kprobe_handler_from_break(manager, &mut pt_regs));
521    if handled.is_some() {
522        ptregs_write_back(&pt_regs, &mut updated);
523        tf.apply_registers(&updated);
524        return true;
525    }
526    false
527}
528
529#[cfg(target_arch = "x86_64")]
530pub fn handle_debug(tf: &mut KernelTrapFrame<'_>) -> bool {
531    let mut updated = tf.snapshot();
532    let mut pt_regs = trapframe_to_ptregs(&updated);
533    let handled = with_manager(|manager| kprobe::kprobe_handler_from_debug(manager, &mut pt_regs));
534    if handled.is_some() {
535        ptregs_write_back(&pt_regs, &mut updated);
536        tf.apply_registers(&updated);
537        return true;
538    }
539    false
540}
starry_kernel/kprobe.rs

starry_kernel/
kprobe.rs
