panda/callbacks/
closure.rs

1use std::{
2    collections::HashMap,
3    ffi::c_void,
4    sync::{
5        atomic::{AtomicU64, Ordering},
6        RwLock,
7    },
8};
9
10use once_cell::sync::OnceCell;
11
12use crate::sys::{hwaddr, target_ptr_t, CPUState, MachineState, Monitor, TranslationBlock};
13use crate::{sys, PluginHandle};
14
15/// A reference to a given callback slot which can be used to install,
16/// enable, disable, or otherwise reference, a closure-based callback.
17///
18/// Since this is a reference to a callback slot and does not include storage
19/// for the callback itself, it can be trivially copied, as well as included in
20/// the callback itself (for the purposes of enabling/disabling).
21///
22/// ## Example
23///
24/// ```
25/// use panda::prelude::*;
26/// use panda::Callback;
27///
28/// let mut count = 0;
29/// let bb_callback = Callback::new();
30/// bb_callback.before_block_exec(move |_, _| {
31///     count += 1;
32///     println!("Basic block #{}", count);
33///
34///     if count > 5 {
35///         bb_callback.disable();
36///     }
37/// });
38///
39/// Panda::new()
40///    .generic("x86_64")
41///    .run();
42/// ```
43///
44/// ## Note
45///
46/// Callback closures must have a static lifetime in order to live past the end of the
47/// function. This means that the only references a callback can include are references
48/// to static variables or leaked objects on the heap (See `Box::leak` for more info).
49///
50/// If you'd like to reference shared data without leaking, this can be accomplished via
51/// reference counting. See [`Arc`] for more info. If you want to capture data owned
52/// by the current function without sharing it, you can mark your closure as `move` in
53/// order to move all the variables you capture into your closure. (Such as in the above
54/// example, where `count` is moved into the closure for modification)
55///
56/// [`Arc`]: https://doc.rust-lang.org/std/sync/struct.Arc.html
57#[repr(transparent)]
58#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
59pub struct Callback(u64);
60
61static CURRENT_CALLBACK_ID: AtomicU64 = AtomicU64::new(0);
62
63impl Callback {
64    /// Create a new callback slot which can then be used to install or modify
65    /// a given callback.
66    pub fn new() -> Self {
67        Self(CURRENT_CALLBACK_ID.fetch_add(1, Ordering::SeqCst))
68    }
69
70    /// Enable the callback assigned to the given slot, if any.
71    pub fn enable(&self) {
72        let callbacks = CALLBACKS.read().unwrap();
73        if let Some(callback) = callbacks.get(&self.0) {
74            unsafe {
75                sys::panda_enable_callback_with_context(
76                    get_plugin_ref(),
77                    callback.cb_kind,
78                    callback.trampoline,
79                    callback.closure_ref as *mut c_void,
80                );
81            }
82        }
83    }
84
85    /// Disable the callback assigned to the given slot, if any.
86    pub fn disable(&self) {
87        let callbacks = CALLBACKS.read().unwrap();
88
89        if let Some(callback) = callbacks.get(&self.0) {
90            unsafe {
91                sys::panda_disable_callback_with_context(
92                    get_plugin_ref(),
93                    callback.cb_kind,
94                    callback.trampoline,
95                    callback.closure_ref as *mut c_void,
96                );
97            }
98        }
99    }
100}
101
102struct ClosureCallback {
103    closure_ref: *mut *mut c_void,
104    cb_kind: sys::panda_cb_type,
105    trampoline: sys::panda_cb_with_context,
106    drop_fn: unsafe fn(*mut *mut c_void),
107}
108
109unsafe impl Sync for ClosureCallback {}
110unsafe impl Send for ClosureCallback {}
111
112lazy_static::lazy_static! {
113    static ref CALLBACKS: RwLock<HashMap<u64, ClosureCallback>> = RwLock::new(HashMap::new());
114}
115
116static PLUGIN_REF: OnceCell<u64> = OnceCell::new();
117
118#[doc(hidden)]
119pub fn set_plugin_ref(plugin: *mut PluginHandle) {
120    let _ = PLUGIN_REF.set(plugin as u64);
121}
122
123fn get_plugin_ref() -> *mut c_void {
124    *PLUGIN_REF.get_or_init(|| &PLUGIN_REF as *const _ as u64) as _
125}
126
127fn install_closure_callback(id: u64, callback: ClosureCallback) {
128    unsafe {
129        sys::panda_register_callback_with_context(
130            get_plugin_ref(),
131            callback.cb_kind,
132            callback.trampoline,
133            callback.closure_ref as *mut c_void,
134        );
135    }
136
137    CALLBACKS.write().unwrap().insert(id, callback);
138}
139
140impl std::ops::Drop for ClosureCallback {
141    fn drop(&mut self) {
142        unsafe { (self.drop_fn)(self.closure_ref) }
143    }
144}
145
146panda_macros::define_closure_callbacks!();