1 use crate::debug::kprobe::args::KprobeInfo;
2 use crate::libs::rwlock::RwLock;
3 use crate::libs::spinlock::SpinLock;
4 use alloc::collections::BTreeMap;
5 use alloc::sync::Arc;
6 use alloc::vec::Vec;
7 use kprobe::{Kprobe, KprobeBuilder, KprobeOps, KprobePoint};
8 use system_error::SystemError;
9
10 pub mod args;
11 #[cfg(feature = "kprobe_test")]
12 mod test;
13
14 pub type LockKprobe = Arc<RwLock<Kprobe>>;
15 pub static KPROBE_MANAGER: SpinLock<KprobeManager> = SpinLock::new(KprobeManager::new());
16 static KPROBE_POINT_LIST: SpinLock<BTreeMap<usize, Arc<KprobePoint>>> =
17 SpinLock::new(BTreeMap::new());
18
19 /// 管理所有的kprobe探测点
20 #[derive(Debug, Default)]
21 pub struct KprobeManager {
22 break_list: BTreeMap<usize, Vec<LockKprobe>>,
23 debug_list: BTreeMap<usize, Vec<LockKprobe>>,
24 }
25
26 impl KprobeManager {
new() -> Self27 pub const fn new() -> Self {
28 KprobeManager {
29 break_list: BTreeMap::new(),
30 debug_list: BTreeMap::new(),
31 }
32 }
33 /// # 插入一个kprobe
34 ///
35 /// ## 参数
36 /// - `kprobe`: kprobe的实例
insert_kprobe(&mut self, kprobe: LockKprobe)37 pub fn insert_kprobe(&mut self, kprobe: LockKprobe) {
38 let probe_point = kprobe.read().probe_point().clone();
39 self.insert_break_point(probe_point.break_address(), kprobe.clone());
40 self.insert_debug_point(probe_point.debug_address(), kprobe);
41 }
42
43 /// # 向break_list中插入一个kprobe
44 ///
45 /// ## 参数
46 /// - `address`: kprobe的地址, 由`KprobePoint::break_address()`或者`KprobeBuilder::probe_addr()`返回
47 /// - `kprobe`: kprobe的实例
insert_break_point(&mut self, address: usize, kprobe: LockKprobe)48 fn insert_break_point(&mut self, address: usize, kprobe: LockKprobe) {
49 let list = self.break_list.entry(address).or_default();
50 list.push(kprobe);
51 }
52
53 /// # 向debug_list中插入一个kprobe
54 ///
55 /// ## 参数
56 /// - `address`: kprobe的单步执行地址，由`KprobePoint::debug_address()`返回
57 /// - `kprobe`: kprobe的实例
insert_debug_point(&mut self, address: usize, kprobe: LockKprobe)58 fn insert_debug_point(&mut self, address: usize, kprobe: LockKprobe) {
59 let list = self.debug_list.entry(address).or_default();
60 list.push(kprobe);
61 }
62
get_break_list(&self, address: usize) -> Option<&Vec<LockKprobe>>63 pub fn get_break_list(&self, address: usize) -> Option<&Vec<LockKprobe>> {
64 self.break_list.get(&address)
65 }
66
get_debug_list(&self, address: usize) -> Option<&Vec<LockKprobe>>67 pub fn get_debug_list(&self, address: usize) -> Option<&Vec<LockKprobe>> {
68 self.debug_list.get(&address)
69 }
70
71 /// # 返回一个地址上注册的kprobe数量
72 ///
73 /// ## 参数
74 /// - `address`: kprobe的地址, 由`KprobePoint::break_address()`或者`KprobeBuilder::probe_addr()`返回
kprobe_num(&self, address: usize) -> usize75 pub fn kprobe_num(&self, address: usize) -> usize {
76 self.break_list_len(address)
77 }
78
79 #[inline]
break_list_len(&self, address: usize) -> usize80 fn break_list_len(&self, address: usize) -> usize {
81 self.break_list
82 .get(&address)
83 .map(|list| list.len())
84 .unwrap_or(0)
85 }
86 #[inline]
debug_list_len(&self, address: usize) -> usize87 fn debug_list_len(&self, address: usize) -> usize {
88 self.debug_list
89 .get(&address)
90 .map(|list| list.len())
91 .unwrap_or(0)
92 }
93
94 /// # 移除一个kprobe
95 ///
96 /// ## 参数
97 /// - `kprobe`: kprobe的实例
remove_kprobe(&mut self, kprobe: &LockKprobe)98 pub fn remove_kprobe(&mut self, kprobe: &LockKprobe) {
99 let probe_point = kprobe.read().probe_point().clone();
100 self.remove_one_break(probe_point.break_address(), kprobe);
101 self.remove_one_debug(probe_point.debug_address(), kprobe);
102 }
103
104 /// # 从break_list中移除一个kprobe
105 ///
106 /// 如果没有其他kprobe注册在这个地址上，则删除列表
107 ///
108 /// ## 参数
109 /// - `address`: kprobe的地址, 由`KprobePoint::break_address()`或者`KprobeBuilder::probe_addr()`返回
110 /// - `kprobe`: kprobe的实例
remove_one_break(&mut self, address: usize, kprobe: &LockKprobe)111 fn remove_one_break(&mut self, address: usize, kprobe: &LockKprobe) {
112 if let Some(list) = self.break_list.get_mut(&address) {
113 list.retain(|x| !Arc::ptr_eq(x, kprobe));
114 }
115 if self.break_list_len(address) == 0 {
116 self.break_list.remove(&address);
117 }
118 }
119
120 /// # 从debug_list中移除一个kprobe
121 ///
122 /// 如果没有其他kprobe注册在这个地址上，则删除列表
123 ///
124 /// ## 参数
125 /// - `address`: kprobe的单步执行地址，由`KprobePoint::debug_address()`返回
126 /// - `kprobe`: kprobe的实例
remove_one_debug(&mut self, address: usize, kprobe: &LockKprobe)127 fn remove_one_debug(&mut self, address: usize, kprobe: &LockKprobe) {
128 if let Some(list) = self.debug_list.get_mut(&address) {
129 list.retain(|x| !Arc::ptr_eq(x, kprobe));
130 }
131 if self.debug_list_len(address) == 0 {
132 self.debug_list.remove(&address);
133 }
134 }
135 }
136
137 #[cfg(feature = "kprobe_test")]
138 #[allow(unused)]
139 /// This function is only used for testing kprobe
kprobe_test()140 pub fn kprobe_test() {
141 test::kprobe_test();
142 }
143
144 /// # 注册一个kprobe
145 ///
146 /// 该函数会根据`symbol`查找对应的函数地址，如果找不到则返回错误。
147 ///
148 /// ## 参数
149 /// - `kprobe_info`: kprobe的信息
register_kprobe(kprobe_info: KprobeInfo) -> Result<LockKprobe, SystemError>150 pub fn register_kprobe(kprobe_info: KprobeInfo) -> Result<LockKprobe, SystemError> {
151 let kprobe_builder = KprobeBuilder::try_from(kprobe_info)?;
152 let address = kprobe_builder.probe_addr();
153 let existed_point = KPROBE_POINT_LIST.lock().get(&address).map(Clone::clone);
154 let kprobe = match existed_point {
155 Some(existed_point) => {
156 kprobe_builder
157 .with_probe_point(existed_point.clone())
158 .install()
159 .0
160 }
161 None => {
162 let (kprobe, probe_point) = kprobe_builder.install();
163 KPROBE_POINT_LIST.lock().insert(address, probe_point);
164 kprobe
165 }
166 };
167 let kprobe = Arc::new(RwLock::new(kprobe));
168 KPROBE_MANAGER.lock().insert_kprobe(kprobe.clone());
169 Ok(kprobe)
170 }
171
172 /// # 注销一个kprobe
173 ///
174 /// ## 参数
175 /// - `kprobe`: 已安装的kprobe
unregister_kprobe(kprobe: LockKprobe)176 pub fn unregister_kprobe(kprobe: LockKprobe) {
177 let kprobe_addr = kprobe.read().probe_point().break_address();
178 KPROBE_MANAGER.lock().remove_kprobe(&kprobe);
179 // 如果没有其他kprobe注册在这个地址上，则删除探测点
180 if KPROBE_MANAGER.lock().kprobe_num(kprobe_addr) == 0 {
181 KPROBE_POINT_LIST.lock().remove(&kprobe_addr);
182 }
183 }
184