binaryninja/similarity/session/
resolver.rs

1use super::{SimilaritySession, SimilaritySessionCompletion};
2use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable};
3use crate::settings::Settings;
4use crate::similarity::node::SimilaritySessionNode;
5use crate::similarity::{SimilarityEntityId, SimilaritySessionResolverId};
6use crate::string::IntoCStr;
7use binaryninjacore_sys::*;
8use std::ffi::c_void;
9
10/// Creates similarity result resolvers with the given settings.
11pub trait SimilaritySessionResolverType: Sync + 'static {
12    type SimilaritySessionResolver: SimilaritySessionResolver;
13
14    /// The resolver name shown to users.
15    const NAME: &'static str;
16    /// A short description of the resolver.
17    const DESCRIPTION: &'static str;
18
19    /// Creates a resolver for a session.
20    fn create_resolver(
21        &self,
22        session: &SimilaritySession,
23        settings: &Settings,
24    ) -> Self::SimilaritySessionResolver;
25
26    /// Returns the settings used to configure new resolvers, if any.
27    fn default_settings(&self) -> Option<Ref<Settings>>;
28}
29
30/// Selects provider results for session entities.
31///
32/// Calls for nodes in the same processing group may run at the same time. Do not keep the session after a callback
33/// returns.
34pub trait SimilaritySessionResolver: Send + Sync + 'static {
35    /// Replaces this resolver's settings only if they are valid.
36    ///
37    /// Return `false` without changing the current settings when the new settings are invalid or updates are not
38    /// supported. Use
39    /// [`SimilaritySession::update_resolver_settings`] so affected entities are resolved again.
40    fn update_settings(&self, _settings: &Settings) -> bool {
41        false
42    }
43
44    /// Prepares a node after its entities are created and before providers run.
45    ///
46    /// This is useful for large graphs where not all views are available, but the resolver needs to
47    /// change the scheduled entities or add new entities.
48    fn prepare_for_node(
49        &self,
50        _session: &SimilaritySession,
51        _node: &SimilaritySessionNode,
52        _completion: &SimilaritySessionCompletion,
53        _resolver_id: SimilaritySessionResolverId,
54    ) {
55    }
56
57    /// Selects results after a node's providers have run.
58    ///
59    /// To select a result, call [`SimilaritySessionNode::set_resolved_result`] for the given entity.
60    /// To request another provider and resolver round, call [`SimilaritySessionNode::add_scheduled_entity`].
61    fn resolve_for_node(
62        &self,
63        session: &SimilaritySession,
64        node: &SimilaritySessionNode,
65        entities: &[SimilarityEntityId],
66        completion: &SimilaritySessionCompletion,
67        resolver_id: SimilaritySessionResolverId,
68    );
69}
70
71/// Registers a similarity session resolver type.
72pub fn register_similarity_session_resolver<C>(
73    resolver_ty: C,
74) -> (&'static C, CoreSimilaritySessionResolverType)
75where
76    C: SimilaritySessionResolverType,
77{
78    let name = C::NAME.to_cstr();
79    let description = C::DESCRIPTION.to_cstr();
80    // Resolver types remain registered for the lifetime of the process.
81    let leaked_resolver: &'static C = Box::leak(Box::new(resolver_ty));
82    let result = unsafe {
83        BNRegisterSimilaritySessionResolverType(
84            name.as_ptr(),
85            description.as_ptr(),
86            &mut BNCustomSimilaritySessionResolverType {
87                context: leaked_resolver as *const C as *mut c_void,
88                create: Some(cb_create_resolver::<C>),
89                getDefaultSettings: Some(cb_resolver_default_settings::<C>),
90            },
91        )
92    };
93    let core_resolver_ty = unsafe { CoreSimilaritySessionResolverType::from_raw(result) };
94    (leaked_resolver, core_resolver_ty)
95}
96
97/// A registered similarity session resolver type.
98pub struct CoreSimilaritySessionResolverType {
99    pub(crate) handle: *mut BNSimilaritySessionResolverType,
100}
101
102impl CoreSimilaritySessionResolverType {
103    pub unsafe fn from_raw(handle: *mut BNSimilaritySessionResolverType) -> Self {
104        Self { handle }
105    }
106
107    /// Returns a registered resolver type by name.
108    pub fn by_name(name: &str) -> Option<Self> {
109        let name = name.to_cstr();
110        let raw_type = unsafe { BNGetSimilaritySessionResolverTypeByName(name.as_ptr()) };
111        match raw_type.is_null() {
112            true => None,
113            false => Some(unsafe { Self::from_raw(raw_type) }),
114        }
115    }
116
117    /// Returns all registered resolver types.
118    pub fn all() -> Array<CoreSimilaritySessionResolverType> {
119        let mut count = 0;
120        let result = unsafe { BNGetSimilaritySessionResolverTypeList(&mut count) };
121        unsafe { Array::new(result, count, ()) }
122    }
123
124    /// Returns the registered name.
125    pub fn name(&self) -> String {
126        unsafe {
127            crate::string::BnString::into_string(BNSimilaritySessionResolverTypeGetName(
128                self.handle,
129            ))
130        }
131    }
132
133    /// Returns the resolver description.
134    pub fn description(&self) -> String {
135        unsafe {
136            crate::string::BnString::into_string(BNSimilaritySessionResolverTypeGetDescription(
137                self.handle,
138            ))
139        }
140    }
141
142    /// Creates a resolver for a session, if supported.
143    pub fn create_resolver(
144        &self,
145        session: &SimilaritySession,
146        settings: &Settings,
147    ) -> Option<Ref<CoreSimilaritySessionResolver>> {
148        let handle = unsafe {
149            BNSimilaritySessionResolverTypeCreateResolver(
150                self.handle,
151                session.handle,
152                settings.handle,
153            )
154        };
155        (!handle.is_null()).then(|| unsafe { CoreSimilaritySessionResolver::ref_from_raw(handle) })
156    }
157
158    /// Returns the default resolver settings, if available.
159    pub fn default_settings(&self) -> Option<Ref<Settings>> {
160        let handle = unsafe { BNSimilaritySessionResolverTypeGetDefaultSettings(self.handle) };
161        (!handle.is_null()).then(|| unsafe { Settings::ref_from_raw(handle) })
162    }
163}
164
165impl CoreArrayProvider for CoreSimilaritySessionResolverType {
166    type Raw = *mut BNSimilaritySessionResolverType;
167    type Context = ();
168    type Wrapped<'a> = CoreSimilaritySessionResolverType;
169}
170
171unsafe impl CoreArrayProviderInner for CoreSimilaritySessionResolverType {
172    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
173        BNFreeSimilaritySessionResolverTypeList(raw)
174    }
175    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
176        CoreSimilaritySessionResolverType::from_raw(*raw)
177    }
178}
179
180/// A core-backed similarity session resolver.
181pub struct CoreSimilaritySessionResolver {
182    pub(crate) handle: *mut BNSimilaritySessionResolver,
183}
184
185impl CoreSimilaritySessionResolver {
186    pub unsafe fn from_raw(handle: *mut BNSimilaritySessionResolver) -> Self {
187        Self { handle }
188    }
189
190    pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySessionResolver) -> Ref<Self> {
191        Ref::new(Self { handle })
192    }
193
194    /// Wraps a custom resolver for `session`.
195    pub fn create<C: SimilaritySessionResolver>(
196        ty: &CoreSimilaritySessionResolverType,
197        session: &SimilaritySession,
198        resolver: C,
199    ) -> Ref<CoreSimilaritySessionResolver> {
200        let resolver = Box::into_raw(Box::new(resolver));
201        let mut callbacks = BNCustomSimilaritySessionResolver {
202            context: resolver.cast(),
203            externalRefTaken: None,
204            externalRefReleased: None,
205            updateSettings: Some(cb_update_resolver_settings::<C>),
206            prepareForNode: Some(cb_prepare_for_node::<C>),
207            resolveForNode: Some(cb_resolve_for_node::<C>),
208            free: Some(cb_resolver_free::<C>),
209        };
210        let raw_resolver = unsafe {
211            BNCreateCustomSimilaritySessionResolver(ty.handle, session.handle, &mut callbacks)
212        };
213        unsafe { CoreSimilaritySessionResolver::ref_from_raw(raw_resolver) }
214    }
215
216    /// Returns the resolver's ID.
217    pub fn id(&self) -> SimilaritySessionResolverId {
218        unsafe { BNSimilaritySessionResolverGetId(self.handle) }.into()
219    }
220
221    /// Returns the registered type that created this resolver.
222    pub fn resolver_type(&self) -> CoreSimilaritySessionResolverType {
223        let handle = unsafe { BNSimilaritySessionResolverGetType(self.handle) };
224        unsafe { CoreSimilaritySessionResolverType::from_raw(handle) }
225    }
226}
227
228impl SimilaritySessionResolver for CoreSimilaritySessionResolver {
229    fn prepare_for_node(
230        &self,
231        session: &SimilaritySession,
232        node: &SimilaritySessionNode,
233        completion: &SimilaritySessionCompletion,
234        _resolver_id: SimilaritySessionResolverId,
235    ) {
236        unsafe {
237            BNSimilaritySessionResolverPrepareForNode(
238                self.handle,
239                session.handle,
240                node.handle,
241                completion.handle,
242            )
243        }
244    }
245
246    fn resolve_for_node(
247        &self,
248        session: &SimilaritySession,
249        node: &SimilaritySessionNode,
250        entities: &[SimilarityEntityId],
251        completion: &SimilaritySessionCompletion,
252        _resolver_id: SimilaritySessionResolverId,
253    ) {
254        let raw_entities: Vec<BNSimilarityEntityId> =
255            entities.iter().copied().map(Into::into).collect();
256        unsafe {
257            BNSimilaritySessionResolverResolveForNode(
258                self.handle,
259                session.handle,
260                node.handle,
261                raw_entities.as_ptr(),
262                raw_entities.len(),
263                completion.handle,
264            )
265        }
266    }
267}
268
269unsafe impl Send for CoreSimilaritySessionResolver {}
270unsafe impl Sync for CoreSimilaritySessionResolver {}
271
272impl ToOwned for CoreSimilaritySessionResolver {
273    type Owned = Ref<Self>;
274
275    fn to_owned(&self) -> Self::Owned {
276        unsafe { RefCountable::inc_ref(self) }
277    }
278}
279
280unsafe impl RefCountable for CoreSimilaritySessionResolver {
281    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
282        Ref::new(Self {
283            handle: BNNewSimilaritySessionResolverReference(handle.handle),
284        })
285    }
286
287    unsafe fn dec_ref(handle: &Self) {
288        BNFreeSimilaritySessionResolver(handle.handle);
289    }
290}
291
292impl CoreArrayProvider for CoreSimilaritySessionResolver {
293    type Raw = *mut BNSimilaritySessionResolver;
294    type Context = ();
295    type Wrapped<'a> = Guard<'a, Self>;
296}
297
298unsafe impl CoreArrayProviderInner for CoreSimilaritySessionResolver {
299    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
300        BNFreeSimilaritySessionResolverList(raw, count)
301    }
302
303    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> {
304        Guard::new(Self::from_raw(*raw), context)
305    }
306}
307
308unsafe extern "C" fn cb_create_resolver<C: SimilaritySessionResolverType>(
309    ctxt: *mut c_void,
310    session: *mut BNSimilaritySession,
311    settings: *mut BNSettings,
312) -> *mut BNSimilaritySessionResolver {
313    ffi_wrap!("SimilaritySessionResolverType::create_resolver", unsafe {
314        let ctxt: &C = &*(ctxt as *const C);
315        let session = SimilaritySession::from_raw(session);
316        let settings = Settings::from_raw(settings);
317        let resolver = ctxt.create_resolver(&session, &settings);
318        let core_type = CoreSimilaritySessionResolverType::by_name(C::NAME).unwrap();
319        let core_resolver = CoreSimilaritySessionResolver::create(&core_type, &session, resolver);
320        Ref::into_raw(core_resolver).handle
321    })
322}
323
324unsafe extern "C" fn cb_resolver_default_settings<C: SimilaritySessionResolverType>(
325    ctxt: *mut c_void,
326) -> *mut BNSettings {
327    ffi_wrap!("SimilaritySessionResolverType::default_settings", unsafe {
328        let ctxt: &C = &*(ctxt as *const C);
329        ctxt.default_settings()
330            .map(|settings| Ref::into_raw(settings).handle)
331            .unwrap_or(std::ptr::null_mut())
332    })
333}
334
335unsafe extern "C" fn cb_update_resolver_settings<C: SimilaritySessionResolver>(
336    ctxt: *mut c_void,
337    settings: *mut BNSettings,
338) -> bool {
339    ffi_wrap!("SimilaritySessionResolver::update_settings", unsafe {
340        let ctxt: &C = &*(ctxt as *const C);
341        let settings = Settings::from_raw(settings);
342        ctxt.update_settings(&settings)
343    })
344}
345
346unsafe extern "C" fn cb_resolve_for_node<C: SimilaritySessionResolver>(
347    ctxt: *mut c_void,
348    session: *mut BNSimilaritySession,
349    node: *mut BNSimilaritySessionNode,
350    entities: *const BNSimilarityEntityId,
351    entity_count: usize,
352    completion: *mut BNSimilaritySessionCompletion,
353    resolver_id: BNSimilaritySessionResolverId,
354) {
355    ffi_wrap!("SimilaritySessionResolver::resolve_for_node", unsafe {
356        let ctxt: &C = &*(ctxt as *const C);
357        let session = SimilaritySession::from_raw(session);
358        let node = SimilaritySessionNode::from_raw(node);
359        let entity_slice = crate::ffi::slice_from_raw_parts(entities, entity_count);
360        let mapped_entities: Vec<SimilarityEntityId> =
361            entity_slice.iter().copied().map(Into::into).collect();
362        let completion = SimilaritySessionCompletion::from_raw(completion);
363        ctxt.resolve_for_node(
364            &session,
365            &node,
366            &mapped_entities,
367            &completion,
368            resolver_id.into(),
369        );
370    })
371}
372
373unsafe extern "C" fn cb_prepare_for_node<C: SimilaritySessionResolver>(
374    ctxt: *mut c_void,
375    session: *mut BNSimilaritySession,
376    node: *mut BNSimilaritySessionNode,
377    completion: *mut BNSimilaritySessionCompletion,
378    resolver_id: BNSimilaritySessionResolverId,
379) {
380    ffi_wrap!("SimilaritySessionResolver::prepare_for_node", unsafe {
381        let ctxt: &C = &*(ctxt as *const C);
382        let session = SimilaritySession::from_raw(session);
383        let node = SimilaritySessionNode::from_raw(node);
384        let completion = SimilaritySessionCompletion::from_raw(completion);
385        ctxt.prepare_for_node(&session, &node, &completion, resolver_id.into());
386    })
387}
388
389unsafe extern "C" fn cb_resolver_free<C: SimilaritySessionResolver>(ctxt: *mut c_void) {
390    ffi_wrap!("SimilaritySessionResolver::free", unsafe {
391        let _ = Box::from_raw(ctxt as *mut C);
392    })
393}
binaryninja/similarity/session/resolver.rs

binaryninja/similarity/session/
resolver.rs
