Skip to main content

binaryninja/similarity/session/
receiver.rs

1use super::SimilaritySessionCompletion;
2use crate::rc::{CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable};
3use crate::similarity::node::SimilaritySessionNode;
4use crate::similarity::provider::CoreSimilarityProvider;
5use crate::similarity::SimilarityEntityId;
6use binaryninjacore_sys::*;
7use std::ffi::c_void;
8
9/// Receives start and update notifications from a similarity session.
10pub trait SimilaritySessionReceiver: Send + Sync + 'static {
11    /// Called when a session run starts.
12    ///
13    /// This is mainly used to get the completion state for the run.
14    fn on_started(&self, _completion: &SimilaritySessionCompletion) {}
15
16    /// Called when a provider's results, resolution state, or applied metadata changes.
17    fn on_updated(
18        &self,
19        node: &SimilaritySessionNode,
20        provider: &CoreSimilarityProvider,
21        entities: &[SimilarityEntityId],
22    );
23}
24
25/// A core-backed similarity session receiver.
26pub struct CoreSimilaritySessionReceiver {
27    pub(crate) handle: *mut BNSimilaritySessionReceiver,
28}
29
30impl CoreSimilaritySessionReceiver {
31    pub unsafe fn from_raw(handle: *mut BNSimilaritySessionReceiver) -> Self {
32        Self { handle }
33    }
34
35    pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySessionReceiver) -> Ref<Self> {
36        Ref::new(Self { handle })
37    }
38
39    /// Wraps a custom session receiver in a core receiver.
40    pub fn create<C: SimilaritySessionReceiver>(receiver: C) -> Ref<CoreSimilaritySessionReceiver> {
41        let receiver = Box::into_raw(Box::new(receiver));
42        let mut callbacks = BNCustomSimilaritySessionReceiver {
43            context: receiver.cast(),
44            externalRefTaken: None,
45            externalRefReleased: None,
46            onStarted: Some(cb_on_started::<C>),
47            onUpdated: Some(cb_on_updated::<C>),
48            free: Some(cb_receiver_free::<C>),
49        };
50        let raw_receiver = unsafe { BNCreateCustomSimilaritySessionReceiver(&mut callbacks) };
51        unsafe { CoreSimilaritySessionReceiver::ref_from_raw(raw_receiver) }
52    }
53}
54
55impl SimilaritySessionReceiver for CoreSimilaritySessionReceiver {
56    fn on_started(&self, completion: &SimilaritySessionCompletion) {
57        unsafe { BNSimilaritySessionReceiverNotifyStart(self.handle, completion.handle) }
58    }
59
60    fn on_updated(
61        &self,
62        node: &SimilaritySessionNode,
63        provider: &CoreSimilarityProvider,
64        entities: &[SimilarityEntityId],
65    ) {
66        let raw_entities: Vec<BNSimilarityEntityId> =
67            entities.iter().copied().map(Into::into).collect();
68        unsafe {
69            BNSimilaritySessionReceiverNotifyBatch(
70                self.handle,
71                node.handle,
72                provider.handle,
73                raw_entities.as_ptr(),
74                raw_entities.len(),
75            )
76        }
77    }
78}
79
80unsafe impl Send for CoreSimilaritySessionReceiver {}
81unsafe impl Sync for CoreSimilaritySessionReceiver {}
82
83impl ToOwned for CoreSimilaritySessionReceiver {
84    type Owned = Ref<Self>;
85
86    fn to_owned(&self) -> Self::Owned {
87        unsafe { RefCountable::inc_ref(self) }
88    }
89}
90
91unsafe impl RefCountable for CoreSimilaritySessionReceiver {
92    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
93        Ref::new(Self {
94            handle: BNNewSimilaritySessionReceiverReference(handle.handle),
95        })
96    }
97
98    unsafe fn dec_ref(handle: &Self) {
99        BNFreeSimilaritySessionReceiver(handle.handle);
100    }
101}
102
103impl CoreArrayProvider for CoreSimilaritySessionReceiver {
104    type Raw = *mut BNSimilaritySessionReceiver;
105    type Context = ();
106    type Wrapped<'a> = Guard<'a, Self>;
107}
108
109unsafe impl CoreArrayProviderInner for CoreSimilaritySessionReceiver {
110    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
111        BNFreeSimilaritySessionReceiverList(raw, count)
112    }
113
114    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> {
115        Guard::new(Self::from_raw(*raw), context)
116    }
117}
118
119unsafe extern "C" fn cb_on_updated<C: SimilaritySessionReceiver>(
120    ctxt: *mut c_void,
121    node: *mut BNSimilaritySessionNode,
122    provider: *mut BNSimilarityProvider,
123    entities: *const BNSimilarityEntityId,
124    count: usize,
125) {
126    ffi_wrap!("SimilaritySessionReceiver::on_updated", unsafe {
127        let ctxt: &C = &*(ctxt as *const C);
128        let node = SimilaritySessionNode::from_raw(node);
129        let provider = CoreSimilarityProvider { handle: provider };
130        let entity_slice = crate::ffi::slice_from_raw_parts(entities, count);
131        let mapped_entities: Vec<SimilarityEntityId> =
132            entity_slice.iter().copied().map(Into::into).collect();
133        ctxt.on_updated(&node, &provider, &mapped_entities);
134    })
135}
136
137unsafe extern "C" fn cb_on_started<C: SimilaritySessionReceiver>(
138    ctxt: *mut c_void,
139    completion: *mut BNSimilaritySessionCompletion,
140) {
141    ffi_wrap!("SimilaritySessionReceiver::on_started", unsafe {
142        let ctxt: &C = &*(ctxt as *const C);
143        let completion = SimilaritySessionCompletion::from_raw(completion);
144        ctxt.on_started(&completion);
145    })
146}
147
148unsafe extern "C" fn cb_receiver_free<C: SimilaritySessionReceiver>(ctxt: *mut c_void) {
149    ffi_wrap!("SimilaritySessionReceiver::free", unsafe {
150        let _ = Box::from_raw(ctxt as *mut C);
151    })
152}