wraith/manipulation/inline_hook/
registry.rs1use std::collections::HashMap;
6use std::sync::Mutex;
7
8static REGISTRY: Mutex<Option<HookRegistry>> = Mutex::new(None);
10
11#[derive(Debug, Clone)]
13pub struct RegisteredHook {
14 pub target: usize,
16 pub detour: usize,
18 pub trampoline: Option<usize>,
20 pub hook_type: HookType,
22 pub installed_at: std::time::Instant,
24 pub active: bool,
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum HookType {
31 Inline,
33 HotPatch,
35 MidFunction,
37 Chained,
39}
40
41#[derive(Default)]
43pub struct HookRegistry {
44 by_target: HashMap<usize, RegisteredHook>,
46 by_detour: HashMap<usize, usize>,
48}
49
50impl HookRegistry {
51 pub fn new() -> Self {
53 Self::default()
54 }
55
56 pub fn register(&mut self, hook: RegisteredHook) {
58 let target = hook.target;
59 let detour = hook.detour;
60
61 self.by_detour.insert(detour, target);
62 self.by_target.insert(target, hook);
63 }
64
65 pub fn unregister(&mut self, target: usize) -> Option<RegisteredHook> {
67 if let Some(hook) = self.by_target.remove(&target) {
68 self.by_detour.remove(&hook.detour);
69 Some(hook)
70 } else {
71 None
72 }
73 }
74
75 pub fn get(&self, target: usize) -> Option<&RegisteredHook> {
77 self.by_target.get(&target)
78 }
79
80 pub fn get_mut(&mut self, target: usize) -> Option<&mut RegisteredHook> {
82 self.by_target.get_mut(&target)
83 }
84
85 pub fn is_hooked(&self, target: usize) -> bool {
87 self.by_target.contains_key(&target)
88 }
89
90 pub fn find_by_detour(&self, detour: usize) -> Option<usize> {
92 self.by_detour.get(&detour).copied()
93 }
94
95 pub fn all(&self) -> impl Iterator<Item = &RegisteredHook> {
97 self.by_target.values()
98 }
99
100 pub fn count(&self) -> usize {
102 self.by_target.len()
103 }
104
105 pub fn clear(&mut self) {
107 self.by_target.clear();
108 self.by_detour.clear();
109 }
110
111 pub fn set_active(&mut self, target: usize, active: bool) -> bool {
113 if let Some(hook) = self.by_target.get_mut(&target) {
114 hook.active = active;
115 true
116 } else {
117 false
118 }
119 }
120}
121
122pub fn get_registry() -> std::sync::MutexGuard<'static, Option<HookRegistry>> {
124 let mut guard = REGISTRY.lock().unwrap();
125 if guard.is_none() {
126 *guard = Some(HookRegistry::new());
127 }
128 guard
129}
130
131pub fn with_registry<F, R>(f: F) -> R
133where
134 F: FnOnce(&mut HookRegistry) -> R,
135{
136 let mut guard = get_registry();
137 let registry = guard.as_mut().unwrap();
138 f(registry)
139}
140
141pub fn register_hook(
143 target: usize,
144 detour: usize,
145 trampoline: Option<usize>,
146 hook_type: HookType,
147) {
148 with_registry(|registry| {
149 registry.register(RegisteredHook {
150 target,
151 detour,
152 trampoline,
153 hook_type,
154 installed_at: std::time::Instant::now(),
155 active: true,
156 });
157 });
158}
159
160pub fn unregister_hook(target: usize) -> Option<RegisteredHook> {
162 with_registry(|registry| registry.unregister(target))
163}
164
165pub fn is_hooked(target: usize) -> bool {
167 with_registry(|registry| registry.is_hooked(target))
168}
169
170pub fn get_hook_info(target: usize) -> Option<RegisteredHook> {
172 with_registry(|registry| registry.get(target).cloned())
173}
174
175pub fn get_hooks_in_range(start: usize, end: usize) -> Vec<RegisteredHook> {
177 with_registry(|registry| {
178 registry
179 .all()
180 .filter(|h| h.target >= start && h.target < end)
181 .cloned()
182 .collect()
183 })
184}
185
186pub fn active_hook_count() -> usize {
188 with_registry(|registry| registry.all().filter(|h| h.active).count())
189}
190
191pub fn total_hook_count() -> usize {
193 with_registry(|registry| registry.count())
194}