kprobe 0.4.0

A no_std Rust probe infrastructure crate for kprobe, kretprobe, and uprobe on multiple architectures.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! A Rust library for dynamic kernel probing (kprobes and kretprobes).
//!! This library provides a safe and ergonomic interface for registering and managing kprobes
//! and kretprobes in a kernel environment.
//!! It supports multiple architectures, including x86_64, riscv64, and loongarch64.
//!! # Features
//! - Register and unregister kprobes and kretprobes.
//! - Support for pre-handler and post-handler functions.
//! - Safe management of probe points and handlers.
//! - Architecture-specific implementations for handling breakpoints and single-stepping.
//!
#![no_std]
#![deny(missing_docs)]
extern crate alloc;

mod arch;
mod kprobe;
mod kretprobe;
mod manager;
mod uprobe;
mod uretprobe;

use alloc::sync::Arc;
use core::ops::Deref;

pub use arch::*;
pub use kprobe::*;
pub use kretprobe::*;
use lock_api::RawMutex;
pub use manager::*;
pub use uprobe::*;
pub use uretprobe::*;

use crate::arch::retprobe::RetprobeData;

/// An enum representing either a kprobe or a kretprobe.
#[derive(Debug)]
#[allow(missing_docs)]
pub enum UniProbe<L: RawMutex + 'static, F: KprobeAuxiliaryOps> {
    Kprobe(Arc<Kprobe<L, F>>),
    Kretprobe(Arc<Kretprobe<L, F>>),
    Uprobe(Arc<Uprobe<L, F>>),
    Uretprobe(Arc<Uretprobe<L, F>>),
}

impl<L: RawMutex + 'static, F: KprobeAuxiliaryOps> Deref for UniProbe<L, F> {
    type Target = Kprobe<L, F>;
    fn deref(&self) -> &Self::Target {
        match self {
            UniProbe::Kprobe(kprobe) => kprobe,
            UniProbe::Kretprobe(kretprobe) => kretprobe.kprobe(),
            UniProbe::Uprobe(uprobe) => uprobe,
            UniProbe::Uretprobe(uretprobe) => uretprobe.kprobe(),
        }
    }
}

impl<L: RawMutex + 'static, F: KprobeAuxiliaryOps> UniProbe<L, F> {
    /// Get the probe point of the probe.
    pub fn probe_point(&self) -> &Arc<ProbePoint<F>> {
        match self {
            UniProbe::Kprobe(kprobe) => kprobe.probe_point(),
            UniProbe::Kretprobe(kretprobe) => kretprobe.kprobe().probe_point(),
            UniProbe::Uprobe(uprobe) => uprobe.probe_point(),
            UniProbe::Uretprobe(uretprobe) => uretprobe.kprobe().probe_point(),
        }
    }
    /// Check if the probe is a kretprobe.
    pub fn is_kretprobe(&self) -> bool {
        matches!(self, UniProbe::Kretprobe(_))
    }
}

impl<L: RawMutex + 'static, F: KprobeAuxiliaryOps> Clone for UniProbe<L, F> {
    fn clone(&self) -> Self {
        match self {
            UniProbe::Kprobe(kprobe) => UniProbe::Kprobe(kprobe.clone()),
            UniProbe::Kretprobe(kretprobe) => UniProbe::Kretprobe(kretprobe.clone()),
            UniProbe::Uprobe(uprobe) => UniProbe::Uprobe(uprobe.clone()),
            UniProbe::Uretprobe(uretprobe) => UniProbe::Uretprobe(uretprobe.clone()),
        }
    }
}

/// Register a kretprobe.
///
/// See [`register_kprobe`] for more details.
pub fn register_kretprobe<L: RawMutex + 'static, F: KprobeAuxiliaryOps + 'static>(
    manager: &mut ProbeManager<L, F>,
    kprobe_point_list: &mut ProbePointList<F>,
    kretprobe_builder: KretprobeBuilder<L>,
) -> Arc<Kretprobe<L, F>> {
    let (entry_handler, ret_handler) = kretprobe_builder.handler();

    let kprobe_builder = ProbeBuilder::from(kretprobe_builder);
    let kprobe = kprobe::__register_kprobe(kprobe_point_list, kprobe_builder);

    let kretprobe = Kretprobe::new(kprobe, entry_handler, ret_handler);
    let kretprobe = Arc::new(kretprobe);

    let data = kretprobe.kprobe().get_data();
    let data = data.as_any().downcast_ref::<RetprobeData<L, F>>().unwrap();
    *data.retprobe.lock() = Arc::downgrade(&kretprobe);

    manager.insert_probe(UniProbe::Kretprobe(kretprobe.clone()));
    kretprobe
}

/// Unregister a kretprobe.
///
/// See [`unregister_kprobe`] for more details.
pub fn unregister_kretprobe<L: RawMutex + 'static, F: KprobeAuxiliaryOps>(
    manager: &mut ProbeManager<L, F>,
    kprobe_point_list: &mut ProbePointList<F>,
    kretprobe: Arc<Kretprobe<L, F>>,
) {
    let kprobe = kretprobe.kprobe();
    let kprobe_addr = kprobe.probe_point().break_address();
    manager.remove_probe(UniProbe::Kretprobe(kretprobe));

    if manager.kprobe_num(kprobe_addr) == 0 {
        kprobe_point_list.remove(&kprobe_addr);
    }
}