wraith/manipulation/inline_hook/
registry.rs

1//! Global hook registry
2//!
3//! Tracks all installed hooks for management and conflict detection.
4
5use std::collections::HashMap;
6use std::sync::Mutex;
7
8/// global hook registry singleton
9static REGISTRY: Mutex<Option<HookRegistry>> = Mutex::new(None);
10
11/// information about a registered hook
12#[derive(Debug, Clone)]
13pub struct RegisteredHook {
14    /// target function address
15    pub target: usize,
16    /// detour function address
17    pub detour: usize,
18    /// trampoline address (if available)
19    pub trampoline: Option<usize>,
20    /// type of hook
21    pub hook_type: HookType,
22    /// when the hook was installed
23    pub installed_at: std::time::Instant,
24    /// whether the hook is currently active
25    pub active: bool,
26}
27
28/// type of hook
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum HookType {
31    /// standard inline hook (prologue replacement)
32    Inline,
33    /// hot-patch style hook
34    HotPatch,
35    /// mid-function hook
36    MidFunction,
37    /// hook in a chain
38    Chained,
39}
40
41/// hook registry for tracking installed hooks
42#[derive(Default)]
43pub struct HookRegistry {
44    /// hooks by target address
45    by_target: HashMap<usize, RegisteredHook>,
46    /// target address by detour address
47    by_detour: HashMap<usize, usize>,
48}
49
50impl HookRegistry {
51    /// create a new registry
52    pub fn new() -> Self {
53        Self::default()
54    }
55
56    /// register a new hook
57    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    /// unregister a hook by target address
66    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    /// get hook info by target address
76    pub fn get(&self, target: usize) -> Option<&RegisteredHook> {
77        self.by_target.get(&target)
78    }
79
80    /// get hook info by target address (mutable)
81    pub fn get_mut(&mut self, target: usize) -> Option<&mut RegisteredHook> {
82        self.by_target.get_mut(&target)
83    }
84
85    /// check if an address is hooked
86    pub fn is_hooked(&self, target: usize) -> bool {
87        self.by_target.contains_key(&target)
88    }
89
90    /// find target by detour address
91    pub fn find_by_detour(&self, detour: usize) -> Option<usize> {
92        self.by_detour.get(&detour).copied()
93    }
94
95    /// get all registered hooks
96    pub fn all(&self) -> impl Iterator<Item = &RegisteredHook> {
97        self.by_target.values()
98    }
99
100    /// get count of registered hooks
101    pub fn count(&self) -> usize {
102        self.by_target.len()
103    }
104
105    /// clear all hooks (does not actually remove them, just clears registry)
106    pub fn clear(&mut self) {
107        self.by_target.clear();
108        self.by_detour.clear();
109    }
110
111    /// set hook active state
112    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
122/// get the global registry (initializes if needed)
123pub 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
131/// execute a function with the registry
132pub 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
141/// register a hook in the global registry
142pub 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
160/// unregister a hook from the global registry
161pub fn unregister_hook(target: usize) -> Option<RegisteredHook> {
162    with_registry(|registry| registry.unregister(target))
163}
164
165/// check if a target is hooked
166pub fn is_hooked(target: usize) -> bool {
167    with_registry(|registry| registry.is_hooked(target))
168}
169
170/// get hook info for a target
171pub fn get_hook_info(target: usize) -> Option<RegisteredHook> {
172    with_registry(|registry| registry.get(target).cloned())
173}
174
175/// get all hooks targeting a module's address range
176pub 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
186/// get count of active hooks
187pub fn active_hook_count() -> usize {
188    with_registry(|registry| registry.all().filter(|h| h.active).count())
189}
190
191/// get total hook count
192pub fn total_hook_count() -> usize {
193    with_registry(|registry| registry.count())
194}