binaryninja/similarity/
provider.rs

1use super::node::SimilaritySessionNode;
2use super::render::SimilarityRenderContext;
3use super::{
4    SimilarityApplyStatus, SimilarityEntityId, SimilarityEntityRef, SimilarityProviderId,
5    SimilarityResultId, SimilaritySessionCompletion,
6};
7use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable};
8use crate::settings::Settings;
9use crate::string::{BnString, IntoCStr};
10use binaryninjacore_sys::*;
11use std::ffi::{c_char, c_void};
12use std::marker::PhantomData;
13use std::rc::Rc;
14
15/// Registers a similarity provider type.
16pub fn register_similarity_provider<C>(provider_ty: C) -> (&'static C, CoreSimilarityProviderType)
17where
18    C: SimilarityProviderType,
19{
20    let name = C::NAME.to_cstr();
21    let description = C::DESCRIPTION.to_cstr();
22    // Provider types remain registered for the lifetime of the process.
23    let leaked_provider: &'static C = Box::leak(Box::new(provider_ty));
24    let result = unsafe {
25        BNRegisterSimilarityProviderType(
26            name.as_ptr(),
27            description.as_ptr(),
28            &mut BNCustomSimilarityProviderType {
29                context: leaked_provider as *const C as *mut c_void,
30                create: Some(cb_create_provider::<C>),
31                getDefaultSettings: Some(cb_default_settings::<C>),
32            },
33        )
34    };
35    let core_provider_ty = unsafe { CoreSimilarityProviderType::from_raw(result) };
36    (leaked_provider, core_provider_ty)
37}
38
39/// Creates similarity providers with the given settings.
40pub trait SimilarityProviderType: Sync + 'static {
41    type SimilarityProvider: SimilarityProvider;
42
43    /// The provider name shown to users.
44    const NAME: &'static str;
45
46    /// A short description of the provider.
47    const DESCRIPTION: &'static str;
48
49    /// Creates a provider with the given settings.
50    fn create_provider(&self, settings: &Settings) -> Self::SimilarityProvider;
51
52    /// Returns the settings used to configure new providers, if any.
53    fn default_settings(&self) -> Option<Ref<Settings>>;
54}
55
56/// Produces and applies similarity results for session entities.
57///
58/// Visits made by a session keep the visited node and both edge endpoints active. Direct calls must provide active
59/// views.
60pub trait SimilarityProvider: Send + Sync + 'static {
61    /// Replaces this provider's settings only if they are valid.
62    ///
63    /// Return `false` without changing the current settings when the new settings are invalid or updates are not
64    /// supported. Use
65    /// [`SimilaritySession::update_provider_settings`](super::SimilaritySession::update_provider_settings) so affected
66    /// entities are scheduled again.
67    fn update_settings(&self, _settings: &Settings) -> bool {
68        false
69    }
70
71    /// Visits the node's scheduled entities and writes results for the node.
72    ///
73    /// A successful visit replaces earlier results from this provider for scheduled entities. Results for unscheduled
74    /// entities remain unchanged. Return `false` to discard the visit.
75    fn visit_node(
76        &self,
77        _node: &SimilaritySessionNode,
78        _results: &mut SimilarityProviderResults<'_>,
79        _completion: &SimilaritySessionCompletion,
80    ) -> bool {
81        true
82    }
83
84    /// Visits an edge after both endpoint nodes have been visited and writes results for the edge.
85    ///
86    /// Return `false` to discard the visit.
87    fn visit_node_edge(
88        &self,
89        _from: &SimilaritySessionNode,
90        _to: &SimilaritySessionNode,
91        _results: &mut SimilarityProviderResults<'_>,
92        _completion: &SimilaritySessionCompletion,
93    ) -> bool {
94        true
95    }
96
97    /// Returns the display name for a result.
98    fn result_name(
99        &self,
100        node: &SimilaritySessionNode,
101        entity: SimilarityEntityId,
102        result: SimilarityResultId,
103    ) -> Option<String>;
104
105    /// Applies a result.
106    ///
107    /// The default implementation performs the standard metadata transfer from the result target.
108    fn apply_result(
109        &self,
110        node: &SimilaritySessionNode,
111        entity: SimilarityEntityId,
112        result: SimilarityResultId,
113    ) -> SimilarityApplyStatus {
114        let Some(result) = node.result(result) else {
115            return SimilarityApplyStatus::SimilarityApplyFailed;
116        };
117        node.apply_target(entity, result.target)
118    }
119
120    /// Adds views for a result to a render context.
121    fn render_result(
122        &self,
123        _node: &SimilaritySessionNode,
124        _entity: SimilarityEntityId,
125        _context: &SimilarityRenderContext,
126        _result: SimilarityResultId,
127    ) {
128    }
129}
130
131/// Writes results for one provider node or edge visit.
132///
133/// Only use this writer during the provider callback that received it.
134pub struct SimilarityProviderResults<'a> {
135    handle: *mut BNSimilarityProviderResults,
136    _lifetime: PhantomData<&'a mut BNSimilarityProviderResults>,
137    _not_send_or_sync: PhantomData<Rc<()>>,
138}
139
140impl<'a> SimilarityProviderResults<'a> {
141    unsafe fn from_raw(handle: *mut BNSimilarityProviderResults) -> Self {
142        Self {
143            handle,
144            _lifetime: PhantomData,
145            _not_send_or_sync: PhantomData,
146        }
147    }
148
149    /// Adds a result for a scheduled entity and returns its ID, or zero on failure. The ID is unique
150    /// within the node. A later visit replaces the result and gives it a new ID.
151    pub fn add_result(
152        &mut self,
153        source: SimilarityEntityRef,
154        target: SimilarityEntityRef,
155        similarity: u8,
156        confidence: u8,
157    ) -> SimilarityResultId {
158        let source = BNSimilarityEntityRef::from(source);
159        let target = BNSimilarityEntityRef::from(target);
160        unsafe {
161            BNSimilarityProviderResultsAddResult(
162                self.handle,
163                &source,
164                &target,
165                similarity,
166                confidence,
167            )
168            .into()
169        }
170    }
171}
172
173/// A registered similarity provider type.
174pub struct CoreSimilarityProviderType {
175    pub(crate) handle: *mut BNSimilarityProviderType,
176}
177
178impl CoreSimilarityProviderType {
179    pub unsafe fn from_raw(handle: *mut BNSimilarityProviderType) -> Self {
180        Self { handle }
181    }
182
183    /// Returns a registered provider type by name.
184    pub fn by_name(name: &str) -> Option<CoreSimilarityProviderType> {
185        let name = name.to_cstr();
186        let raw_type = unsafe { BNGetSimilarityProviderTypeByName(name.as_ptr()) };
187        match raw_type.is_null() {
188            true => None,
189            false => Some(unsafe { Self::from_raw(raw_type) }),
190        }
191    }
192
193    /// Returns all registered provider types.
194    pub fn all() -> Array<CoreSimilarityProviderType> {
195        let mut count = 0;
196        let result = unsafe { BNGetSimilarityProviderTypeList(&mut count) };
197        unsafe { Array::new(result, count, ()) }
198    }
199
200    /// Returns the registered name.
201    pub fn name(&self) -> String {
202        unsafe { BnString::into_string(BNSimilarityProviderTypeGetName(self.handle)) }
203    }
204
205    /// Returns the provider description.
206    pub fn description(&self) -> String {
207        unsafe { BnString::into_string(BNSimilarityProviderTypeGetDescription(self.handle)) }
208    }
209
210    /// Creates a provider with the given settings, if supported. Returns `None` outside Ultimate.
211    pub fn create_provider(&self, settings: &Settings) -> Option<Ref<CoreSimilarityProvider>> {
212        let provider_raw =
213            unsafe { BNSimilarityProviderTypeCreateProvider(self.handle, settings.handle) };
214        (!provider_raw.is_null())
215            .then(|| unsafe { CoreSimilarityProvider::ref_from_raw(provider_raw) })
216    }
217
218    /// Returns the default provider settings, if available.
219    pub fn default_settings(&self) -> Option<Ref<Settings>> {
220        let settings_raw = unsafe { BNSimilarityProviderTypeGetDefaultSettings(self.handle) };
221        (!settings_raw.is_null()).then(|| unsafe { Settings::ref_from_raw(settings_raw) })
222    }
223}
224
225impl CoreArrayProvider for CoreSimilarityProviderType {
226    type Raw = *mut BNSimilarityProviderType;
227    type Context = ();
228    type Wrapped<'a> = CoreSimilarityProviderType;
229}
230
231unsafe impl CoreArrayProviderInner for CoreSimilarityProviderType {
232    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
233        BNFreeSimilarityProviderTypeList(raw)
234    }
235
236    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
237        CoreSimilarityProviderType::from_raw(*raw)
238    }
239}
240
241/// A core-backed similarity provider.
242pub struct CoreSimilarityProvider {
243    pub(crate) handle: *mut BNSimilarityProvider,
244}
245
246impl CoreSimilarityProvider {
247    pub unsafe fn from_raw(handle: *mut BNSimilarityProvider) -> Self {
248        Self { handle }
249    }
250
251    pub unsafe fn ref_from_raw(handle: *mut BNSimilarityProvider) -> Ref<Self> {
252        Ref::new(Self { handle })
253    }
254
255    /// Returns the provider's ID.
256    pub fn id(&self) -> SimilarityProviderId {
257        unsafe { BNSimilarityProviderGetId(self.handle) }.into()
258    }
259
260    /// Wraps a custom provider in a core provider.
261    pub fn create<C: SimilarityProvider>(
262        ty: &CoreSimilarityProviderType,
263        provider: C,
264    ) -> Ref<CoreSimilarityProvider> {
265        let provider = Box::into_raw(Box::new(provider));
266        let mut callbacks = BNCustomSimilarityProvider {
267            context: provider.cast(),
268            externalRefTaken: None,
269            externalRefReleased: None,
270            updateSettings: Some(cb_update_provider_settings::<C>),
271            visitNode: Some(cb_visit_node::<C>),
272            visitNodeEdge: Some(cb_visit_node_edge::<C>),
273            getName: Some(cb_provider_get_name::<C>),
274            apply: Some(cb_provider_apply::<C>),
275            render: Some(cb_provider_render::<C>),
276            free: Some(cb_provider_free::<C>),
277        };
278
279        let raw_provider = unsafe { BNCreateCustomSimilarityProvider(ty.handle, &mut callbacks) };
280        unsafe { CoreSimilarityProvider::ref_from_raw(raw_provider) }
281    }
282
283    /// Returns the registered type that created this provider.
284    pub fn provider_type(&self) -> CoreSimilarityProviderType {
285        let handle = unsafe { BNSimilarityProviderGetType(self.handle) };
286        unsafe { CoreSimilarityProviderType::from_raw(handle) }
287    }
288
289    /// Performs a complete node visit with the core managing result updates.
290    pub fn visit_node(
291        &self,
292        node: &SimilaritySessionNode,
293        completion: &SimilaritySessionCompletion,
294    ) {
295        unsafe { BNSimilarityProviderVisitNode(self.handle, node.handle, completion.handle) }
296    }
297
298    /// Performs a complete edge visit with the core managing result updates.
299    pub fn visit_node_edge(
300        &self,
301        from: &SimilaritySessionNode,
302        to: &SimilaritySessionNode,
303        completion: &SimilaritySessionCompletion,
304    ) {
305        unsafe {
306            BNSimilarityProviderVisitNodeEdge(
307                self.handle,
308                from.handle,
309                to.handle,
310                completion.handle,
311            )
312        }
313    }
314
315    /// Calls this provider's node visit with an existing result writer.
316    pub fn perform_visit_node(
317        &self,
318        node: &SimilaritySessionNode,
319        results: &mut SimilarityProviderResults<'_>,
320        completion: &SimilaritySessionCompletion,
321    ) -> bool {
322        unsafe {
323            BNSimilarityProviderPerformVisitNode(
324                self.handle,
325                node.handle,
326                results.handle,
327                completion.handle,
328            )
329        }
330    }
331
332    /// Calls this provider's edge visit with an existing result writer.
333    pub fn perform_visit_node_edge(
334        &self,
335        from: &SimilaritySessionNode,
336        to: &SimilaritySessionNode,
337        results: &mut SimilarityProviderResults<'_>,
338        completion: &SimilaritySessionCompletion,
339    ) -> bool {
340        unsafe {
341            BNSimilarityProviderPerformVisitNodeEdge(
342                self.handle,
343                from.handle,
344                to.handle,
345                results.handle,
346                completion.handle,
347            )
348        }
349    }
350
351    /// Returns the display name for a result.
352    pub fn result_name(
353        &self,
354        node: &SimilaritySessionNode,
355        entity: SimilarityEntityId,
356        result: SimilarityResultId,
357    ) -> Option<String> {
358        let name_raw = unsafe {
359            BNSimilarityProviderGetName(self.handle, node.handle, entity.into(), result.into())
360        };
361        if name_raw.is_null() {
362            return None;
363        }
364        Some(unsafe { BnString::into_string(name_raw) })
365    }
366
367    /// Applies a result.
368    pub fn apply_result(
369        &self,
370        node: &SimilaritySessionNode,
371        entity: SimilarityEntityId,
372        result: SimilarityResultId,
373    ) -> SimilarityApplyStatus {
374        unsafe { BNSimilarityProviderApply(self.handle, node.handle, entity.into(), result.into()) }
375    }
376
377    /// Adds views for a result to a render context.
378    pub fn render_result(
379        &self,
380        node: &SimilaritySessionNode,
381        entity: SimilarityEntityId,
382        context: &SimilarityRenderContext,
383        result: SimilarityResultId,
384    ) {
385        unsafe {
386            BNSimilarityProviderRender(
387                self.handle,
388                node.handle,
389                entity.into(),
390                context.handle,
391                result.into(),
392            )
393        }
394    }
395}
396
397unsafe impl Send for CoreSimilarityProvider {}
398unsafe impl Sync for CoreSimilarityProvider {}
399
400impl ToOwned for CoreSimilarityProvider {
401    type Owned = Ref<Self>;
402
403    fn to_owned(&self) -> Self::Owned {
404        unsafe { RefCountable::inc_ref(self) }
405    }
406}
407
408unsafe impl RefCountable for CoreSimilarityProvider {
409    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
410        Ref::new(Self {
411            handle: BNNewSimilarityProviderReference(handle.handle),
412        })
413    }
414
415    unsafe fn dec_ref(handle: &Self) {
416        BNFreeSimilarityProvider(handle.handle);
417    }
418}
419
420impl CoreArrayProvider for CoreSimilarityProvider {
421    type Raw = *mut BNSimilarityProvider;
422    type Context = ();
423    type Wrapped<'a> = Guard<'a, Self>;
424}
425
426unsafe impl CoreArrayProviderInner for CoreSimilarityProvider {
427    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
428        BNFreeSimilarityProviderList(raw, count)
429    }
430
431    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> {
432        Guard::new(Self::from_raw(*raw), context)
433    }
434}
435
436/// A match produced by a similarity provider.
437#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
438pub struct SimilarityResult {
439    /// The provider which produced the match.
440    pub provider_id: SimilarityProviderId,
441    /// The similarity of the two entities.
442    pub similarity: u8,
443    /// The provider's confidence in the match.
444    pub confidence: u8,
445    /// The entity matched by this result.
446    ///
447    /// This can be used to transfer metadata from the target to the source entity.
448    pub target: SimilarityEntityRef,
449}
450
451impl From<BNSimilarityResult> for SimilarityResult {
452    fn from(value: BNSimilarityResult) -> Self {
453        Self {
454            provider_id: value.providerId.into(),
455            similarity: value.similarity,
456            confidence: value.confidence,
457            target: value.target.into(),
458        }
459    }
460}
461
462impl From<SimilarityResult> for BNSimilarityResult {
463    fn from(value: SimilarityResult) -> Self {
464        Self {
465            providerId: value.provider_id.into(),
466            similarity: value.similarity,
467            confidence: value.confidence,
468            target: value.target.into(),
469        }
470    }
471}
472
473unsafe extern "C" fn cb_create_provider<C: SimilarityProviderType>(
474    ctxt: *mut c_void,
475    settings: *mut BNSettings,
476) -> *mut BNSimilarityProvider {
477    ffi_wrap!("SimilarityProviderType::create_provider", unsafe {
478        let ctxt: &C = &*(ctxt as *const C);
479        let settings = Settings::from_raw(settings);
480        let provider = ctxt.create_provider(&settings);
481        let core_type = CoreSimilarityProviderType::by_name(C::NAME).unwrap();
482        let core_provider = CoreSimilarityProvider::create(&core_type, provider);
483        Ref::into_raw(core_provider).handle
484    })
485}
486
487unsafe extern "C" fn cb_default_settings<C: SimilarityProviderType>(
488    ctxt: *mut c_void,
489) -> *mut BNSettings {
490    ffi_wrap!("SimilarityProviderType::default_settings", unsafe {
491        let ctxt: &C = &*(ctxt as *const C);
492        ctxt.default_settings()
493            .map(|settings| Ref::into_raw(settings).handle)
494            .unwrap_or(std::ptr::null_mut())
495    })
496}
497
498unsafe extern "C" fn cb_update_provider_settings<C: SimilarityProvider>(
499    ctxt: *mut c_void,
500    settings: *mut BNSettings,
501) -> bool {
502    ffi_wrap!("SimilarityProvider::update_settings", unsafe {
503        let ctxt: &C = &*(ctxt as *const C);
504        let settings = Settings::from_raw(settings);
505        ctxt.update_settings(&settings)
506    })
507}
508
509unsafe extern "C" fn cb_visit_node<C: SimilarityProvider>(
510    ctxt: *mut c_void,
511    node: *mut BNSimilaritySessionNode,
512    results: *mut BNSimilarityProviderResults,
513    completion: *mut BNSimilaritySessionCompletion,
514) -> bool {
515    ffi_wrap!("SimilarityProvider::visit_node", unsafe {
516        let ctxt: &C = &*(ctxt as *const C);
517        let node = SimilaritySessionNode::from_raw(node);
518        let mut results = SimilarityProviderResults::from_raw(results);
519        let completion = SimilaritySessionCompletion::from_raw(completion);
520        ctxt.visit_node(&node, &mut results, &completion)
521    })
522}
523
524unsafe extern "C" fn cb_visit_node_edge<C: SimilarityProvider>(
525    ctxt: *mut c_void,
526    from: *mut BNSimilaritySessionNode,
527    to: *mut BNSimilaritySessionNode,
528    results: *mut BNSimilarityProviderResults,
529    completion: *mut BNSimilaritySessionCompletion,
530) -> bool {
531    ffi_wrap!("SimilarityProvider::visit_node_edge", unsafe {
532        let ctxt: &C = &*(ctxt as *const C);
533        let from = SimilaritySessionNode::from_raw(from);
534        let to = SimilaritySessionNode::from_raw(to);
535        let mut results = SimilarityProviderResults::from_raw(results);
536        let completion = SimilaritySessionCompletion::from_raw(completion);
537        ctxt.visit_node_edge(&from, &to, &mut results, &completion)
538    })
539}
540
541unsafe extern "C" fn cb_provider_get_name<C: SimilarityProvider>(
542    ctxt: *mut c_void,
543    node: *mut BNSimilaritySessionNode,
544    entity: BNSimilarityEntityId,
545    result: BNSimilarityResultId,
546) -> *mut c_char {
547    ffi_wrap!("SimilarityProvider::result_name", unsafe {
548        let ctxt: &C = &*(ctxt as *const C);
549        let node = SimilaritySessionNode::from_raw(node);
550        let Some(name) = ctxt.result_name(&node, entity.into(), result.into()) else {
551            return std::ptr::null_mut();
552        };
553        BnString::into_raw(BnString::new(name))
554    })
555}
556
557unsafe extern "C" fn cb_provider_apply<C: SimilarityProvider>(
558    ctxt: *mut c_void,
559    node: *mut BNSimilaritySessionNode,
560    entity: BNSimilarityEntityId,
561    result: BNSimilarityResultId,
562) -> BNSimilarityApplyStatus {
563    ffi_wrap!("SimilarityProvider::apply_result", unsafe {
564        let ctxt: &C = &*(ctxt as *const C);
565        let node = SimilaritySessionNode::from_raw(node);
566        ctxt.apply_result(&node, entity.into(), result.into())
567    })
568}
569
570unsafe extern "C" fn cb_provider_render<C: SimilarityProvider>(
571    ctxt: *mut c_void,
572    node: *mut BNSimilaritySessionNode,
573    entity: BNSimilarityEntityId,
574    context: *mut BNSimilarityRenderContext,
575    result: BNSimilarityResultId,
576) {
577    ffi_wrap!("SimilarityProvider::render_result", unsafe {
578        let ctxt: &C = &*(ctxt as *const C);
579        let node = SimilaritySessionNode::from_raw(node);
580        let context = SimilarityRenderContext::from_raw(context);
581        ctxt.render_result(&node, entity.into(), &context, result.into());
582    })
583}
584
585unsafe extern "C" fn cb_provider_free<C: SimilarityProvider>(ctxt: *mut c_void) {
586    ffi_wrap!("SimilarityProvider::free", unsafe {
587        let _ = Box::from_raw(ctxt as *mut C);
588    })
589}
binaryninja/similarity/provider.rs

binaryninja/similarity/
provider.rs
