Skip to main content

binaryninja/similarity/
session.rs

1use super::graph::SimilaritySessionGraph;
2use super::provider::CoreSimilarityProvider;
3use super::{
4    SimilarityProviderId, SimilaritySessionCompletionQuery, SimilaritySessionId,
5    SimilaritySessionResolverId,
6};
7use crate::rc::{Array, Ref, RefCountable};
8use crate::settings::Settings;
9use binaryninjacore_sys::*;
10use std::time::Duration;
11
12pub mod receiver;
13pub mod resolver;
14
15pub use receiver::*;
16pub use resolver::*;
17
18/// Coordinates providers and resolvers across a graph of binaries.
19pub struct SimilaritySession {
20    pub(crate) handle: *mut BNSimilaritySession,
21}
22
23impl SimilaritySession {
24    /// Creates an empty session.
25    pub fn new() -> Ref<Self> {
26        let handle = unsafe { BNCreateSimilaritySession() };
27        unsafe { Ref::new(Self { handle }) }
28    }
29
30    pub unsafe fn from_raw(handle: *mut BNSimilaritySession) -> Self {
31        Self { handle }
32    }
33
34    pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySession) -> Ref<Self> {
35        Ref::new(Self { handle })
36    }
37
38    /// Returns the session's unique ID.
39    pub fn id(&self) -> SimilaritySessionId {
40        unsafe { BNSimilaritySessionGetId(self.handle) }.into()
41    }
42
43    /// Adds a provider and schedules entities processed by earlier runs for the next run.
44    ///
45    /// NOTE: Ignored while a run is active.
46    pub fn add_provider(&self, provider: &CoreSimilarityProvider) {
47        unsafe { BNSimilaritySessionAddProvider(self.handle, provider.handle) }
48    }
49
50    /// Removes a provider, clears its results, and marks affected entities for resolution.
51    ///
52    /// NOTE: Ignored while a run is active.
53    pub fn remove_provider(&self, provider: &CoreSimilarityProvider) {
54        unsafe { BNSimilaritySessionRemoveProvider(self.handle, provider.handle) }
55    }
56
57    /// Updates a provider already in the session and schedules previously processed entities again.
58    ///
59    /// Returns `false` during a run, when the provider is absent, or when it rejects the settings.
60    pub fn update_provider_settings(
61        &self,
62        provider: &CoreSimilarityProvider,
63        settings: &Settings,
64    ) -> bool {
65        unsafe {
66            BNSimilaritySessionUpdateProviderSettings(self.handle, provider.handle, settings.handle)
67        }
68    }
69
70    /// Returns a provider by ID.
71    pub fn provider(&self, id: SimilarityProviderId) -> Option<Ref<CoreSimilarityProvider>> {
72        let handle = unsafe { BNSimilaritySessionGetProvider(self.handle, id.into()) };
73        match handle.is_null() {
74            true => None,
75            false => unsafe { Some(CoreSimilarityProvider::ref_from_raw(handle)) },
76        }
77    }
78
79    /// Returns the session's providers.
80    pub fn providers(&self) -> Array<CoreSimilarityProvider> {
81        let mut count = 0;
82        let result = unsafe { BNSimilaritySessionGetProviders(self.handle, &mut count) };
83        unsafe { Array::new(result, count, ()) }
84    }
85
86    /// Adds a resolver created for this session and marks entities processed by earlier runs for resolution.
87    ///
88    /// NOTE: Returns `false` during a run, for a duplicate, or for a resolver created for another session.
89    pub fn add_resolver(&self, resolver: &CoreSimilaritySessionResolver) -> bool {
90        unsafe { BNSimilaritySessionAddResolver(self.handle, resolver.handle) }
91    }
92
93    /// Removes a resolver from the session.
94    ///
95    /// NOTE: Returns `false` during a run, or if it is absent or belongs to another session.
96    pub fn remove_resolver(&self, resolver: &CoreSimilaritySessionResolver) -> bool {
97        unsafe { BNSimilaritySessionRemoveResolver(self.handle, resolver.handle) }
98    }
99
100    /// Updates a resolver already in the session and marks previously processed entities for resolution.
101    ///
102    /// Returns `false` during a run, when the resolver is absent or belongs to another session, or when it rejects the
103    /// settings.
104    pub fn update_resolver_settings(
105        &self,
106        resolver: &CoreSimilaritySessionResolver,
107        settings: &Settings,
108    ) -> bool {
109        unsafe {
110            BNSimilaritySessionUpdateResolverSettings(self.handle, resolver.handle, settings.handle)
111        }
112    }
113
114    /// Returns a resolver by ID.
115    pub fn resolver(
116        &self,
117        id: SimilaritySessionResolverId,
118    ) -> Option<Ref<CoreSimilaritySessionResolver>> {
119        let handle = unsafe { BNSimilaritySessionGetResolver(self.handle, id.into()) };
120        match handle.is_null() {
121            true => None,
122            false => unsafe { Some(CoreSimilaritySessionResolver::ref_from_raw(handle)) },
123        }
124    }
125
126    /// Returns the session's resolvers.
127    pub fn resolvers(&self) -> Array<CoreSimilaritySessionResolver> {
128        let mut count = 0;
129        let result = unsafe { BNSimilaritySessionGetResolvers(self.handle, &mut count) };
130        unsafe { Array::new(result, count, ()) }
131    }
132
133    /// Adds an update receiver.
134    ///
135    /// A running session keeps using the receiver list it started with.
136    pub fn add_receiver(&self, receiver: &CoreSimilaritySessionReceiver) {
137        unsafe { BNSimilaritySessionAddReceiver(self.handle, receiver.handle) }
138    }
139
140    /// Removes an update receiver.
141    ///
142    /// A running session keeps using the receiver list it started with.
143    pub fn remove_receiver(&self, receiver: &CoreSimilaritySessionReceiver) {
144        unsafe { BNSimilaritySessionRemoveReceiver(self.handle, receiver.handle) }
145    }
146
147    /// Returns the session's update receivers.
148    pub fn receivers(&self) -> Array<CoreSimilaritySessionReceiver> {
149        let mut count = 0;
150        let result = unsafe { BNSimilaritySessionGetReceivers(self.handle, &mut count) };
151        unsafe { Array::new(result, count, ()) }
152    }
153
154    /// Returns the session graph.
155    pub fn graph(&self) -> Ref<SimilaritySessionGraph> {
156        let handle = unsafe { BNSimilaritySessionGetGraph(self.handle) };
157        unsafe { SimilaritySessionGraph::ref_from_raw(handle) }
158    }
159
160    /// Starts a background run with the current graph, providers, and resolvers.
161    ///
162    /// Changes to them are ignored until the run finishes.
163    ///
164    /// NOTE: Returns the completion state of the active run when already running.
165    pub fn run(&self) -> Ref<SimilaritySessionCompletion> {
166        unsafe { SimilaritySessionCompletion::ref_from_raw(BNSimilaritySessionRun(self.handle)) }
167    }
168}
169
170impl ToOwned for SimilaritySession {
171    type Owned = Ref<Self>;
172
173    fn to_owned(&self) -> Self::Owned {
174        unsafe { RefCountable::inc_ref(self) }
175    }
176}
177
178unsafe impl RefCountable for SimilaritySession {
179    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
180        Ref::new(Self {
181            handle: BNNewSimilaritySessionReference(handle.handle),
182        })
183    }
184    unsafe fn dec_ref(handle: &Self) {
185        BNFreeSimilaritySession(handle.handle);
186    }
187}
188
189/// Stop requests, progress, and timing for a session run.
190pub struct SimilaritySessionCompletion {
191    pub(crate) handle: *mut BNSimilaritySessionCompletion,
192}
193
194impl SimilaritySessionCompletion {
195    /// Creates a new completion state, typically only done when calling into providers and resolvers
196    /// directly instead of from a session.
197    pub fn new() -> Ref<Self> {
198        unsafe { Self::ref_from_raw(BNCreateSimilaritySessionCompletion()) }
199    }
200
201    pub unsafe fn from_raw(handle: *mut BNSimilaritySessionCompletion) -> Self {
202        Self { handle }
203    }
204
205    pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySessionCompletion) -> Ref<Self> {
206        Ref::new(Self { handle })
207    }
208
209    /// Returns whether the session run has finished.
210    pub fn is_finished(&self) -> bool {
211        unsafe { BNSimilaritySessionCompletionIsFinished(self.handle) }
212    }
213
214    /// Returns progress for the selected part of the run from `0.0` through `1.0`.
215    pub fn progress(&self, query: SimilaritySessionCompletionQuery) -> f64 {
216        let raw_query = query.into();
217        unsafe { BNSimilaritySessionCompletionGetProgress(self.handle, &raw_query) }
218    }
219
220    /// Asks the run to stop.
221    pub fn request_stop(&self) {
222        unsafe { BNSimilaritySessionCompletionRequestStop(self.handle) }
223    }
224
225    /// Returns whether a stop has been requested.
226    pub fn is_stop_requested(&self) -> bool {
227        unsafe { BNSimilaritySessionCompletionIsStopRequested(self.handle) }
228    }
229
230    /// Increases progress for a node and one provider or resolver. Progress cannot decrease.
231    ///
232    /// NOTE: Only call this from the provider or resolver selected by the query.
233    pub fn set_progress(&self, query: SimilaritySessionCompletionQuery, progress: f64) {
234        let raw_query = query.into();
235        unsafe { BNSimilaritySessionCompletionSetProgress(self.handle, &raw_query, progress) }
236    }
237
238    /// Returns timing for the selected part of the run.
239    pub fn timing(&self, query: SimilaritySessionCompletionQuery) -> Duration {
240        let raw_query = query.into();
241        Duration::from_millis(unsafe {
242            BNSimilaritySessionCompletionGetTiming(self.handle, &raw_query)
243        })
244    }
245}
246
247unsafe impl Send for SimilaritySessionCompletion {}
248unsafe impl Sync for SimilaritySessionCompletion {}
249
250impl ToOwned for SimilaritySessionCompletion {
251    type Owned = Ref<Self>;
252
253    fn to_owned(&self) -> Self::Owned {
254        unsafe { RefCountable::inc_ref(self) }
255    }
256}
257
258unsafe impl RefCountable for SimilaritySessionCompletion {
259    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
260        Ref::new(Self {
261            handle: BNNewSimilaritySessionCompletionReference(handle.handle),
262        })
263    }
264    unsafe fn dec_ref(handle: &Self) {
265        BNFreeSimilaritySessionCompletion(handle.handle);
266    }
267}