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 { 27 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的实例 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的实例 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的实例 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 63 pub fn get_break_list(&self, address: usize) -> Option<&Vec<LockKprobe>> { 64 self.break_list.get(&address) 65 } 66 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()`返回 75 pub fn kprobe_num(&self, address: usize) -> usize { 76 self.break_list_len(address) 77 } 78 79 #[inline] 80 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] 87 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的实例 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的实例 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的实例 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 140 pub fn kprobe_test() { 141 test::kprobe_test(); 142 } 143 144 /// # 注册一个kprobe 145 /// 146 /// 该函数会根据`symbol`查找对应的函数地址，如果找不到则返回错误。 147 /// 148 /// ## 参数 149 /// - `kprobe_info`: kprobe的信息 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 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