1- use crate :: cache:: container:: for_cached_containers;
1+ use crate :: cache:: container:: { for_cached_containers, for_cached_containers_mut } ;
22use crate :: cache:: {
33 cached_function_guid, insert_cached_function_match, try_cached_function_guid,
44 try_cached_function_match,
55} ;
66use crate :: convert:: { platform_to_target, to_bn_type} ;
77use crate :: matcher:: { Matcher , MatcherSettings } ;
8+ use crate :: plugin:: settings:: PluginSettings ;
89use crate :: { get_warp_tag_type, relocatable_regions} ;
910use binaryninja:: architecture:: RegisterId ;
1011use binaryninja:: background_task:: BackgroundTask ;
1112use binaryninja:: binary_view:: { BinaryView , BinaryViewExt } ;
1213use binaryninja:: command:: Command ;
14+ use binaryninja:: function:: Function as BNFunction ;
15+ use binaryninja:: rc:: Ref as BNRef ;
1316use binaryninja:: settings:: { QueryOptions , Settings } ;
1417use binaryninja:: workflow:: { activity, Activity , AnalysisContext , Workflow , WorkflowBuilder } ;
1518use itertools:: Itertools ;
@@ -47,6 +50,62 @@ impl Command for RunMatcher {
4750 }
4851}
4952
53+ /// This is a helper struct that contains all the functions that match on a given target.
54+ ///
55+ /// We build this when matching and fetching so that we can relate different functions to another.
56+ #[ derive( Debug , Clone , PartialEq , Eq , Default ) ]
57+ struct FunctionSet {
58+ functions_by_target_and_guid : HashMap < ( FunctionGUID , Target ) , Vec < BNRef < BNFunction > > > ,
59+ guids_by_target : HashMap < Target , Vec < FunctionGUID > > ,
60+ }
61+
62+ impl FunctionSet {
63+ fn from_view ( view : & BinaryView ) -> Option < Self > {
64+ let mut set = FunctionSet :: default ( ) ;
65+
66+ // TODO: Par iter this? Using dashmap
67+ set. functions_by_target_and_guid = view
68+ . functions ( )
69+ . iter ( )
70+ . filter_map ( |f| {
71+ let guid = try_cached_function_guid ( & f) ?;
72+ let target = platform_to_target ( & f. platform ( ) ) ;
73+ Some ( ( ( guid, target) , f. to_owned ( ) ) )
74+ } )
75+ . into_group_map ( ) ;
76+
77+ if set. functions_by_target_and_guid . is_empty ( ) && !view. functions ( ) . is_empty ( ) {
78+ // The user is likely trying to run the matcher on a database before guids were automatically
79+ // generated, we should alert them and ask if they would like to reanalyze.
80+ // NOTE: We only alert if we actually have the GUID activity enabled.
81+ if let Some ( sample_function) = view. functions ( ) . iter ( ) . next ( ) {
82+ let function_workflow = sample_function
83+ . workflow ( )
84+ . expect ( "Function has no workflow" ) ;
85+ if function_workflow. contains ( GUID_ACTIVITY_NAME ) {
86+ log:: error!( "No function guids in database, please reanalyze the database." ) ;
87+ } else {
88+ log:: error!(
89+ "Activity '{}' is not in workflow '{}', create function guids manually to run matcher..." ,
90+ GUID_ACTIVITY_NAME ,
91+ function_workflow. name( )
92+ )
93+ }
94+ }
95+ return None ;
96+ }
97+
98+ // TODO: Par iter this? Using dashmap
99+ set. guids_by_target = set
100+ . functions_by_target_and_guid
101+ . keys ( )
102+ . map ( |( guid, target) | ( target. clone ( ) , * guid) )
103+ . into_group_map ( ) ;
104+
105+ Some ( set)
106+ }
107+ }
108+
50109pub fn run_matcher ( view : & BinaryView ) {
51110 // TODO: Create the tag type so we dont have UB in the apply function workflow.
52111 let undo_id = view. file ( ) . begin_undo_actions ( false ) ;
@@ -63,44 +122,10 @@ pub fn run_matcher(view: &BinaryView) {
63122 let matcher_settings = MatcherSettings :: from_settings ( & view_settings, & mut query_opts) ;
64123 let matcher = Matcher :: new ( matcher_settings) ;
65124
66- // TODO: Par iter this? Using dashmap
67- let functions_by_target_and_guid: HashMap < ( FunctionGUID , Target ) , Vec < _ > > = view
68- . functions ( )
69- . iter ( )
70- . filter_map ( |f| {
71- let guid = try_cached_function_guid ( & f) ?;
72- let target = platform_to_target ( & f. platform ( ) ) ;
73- Some ( ( ( guid, target) , f. to_owned ( ) ) )
74- } )
75- . into_group_map ( ) ;
76-
77- if functions_by_target_and_guid. is_empty ( ) && !view. functions ( ) . is_empty ( ) {
78- // The user is likely trying to run the matcher on a database before guids were automatically
79- // generated, we should alert them and ask if they would like to reanalyze.
80- // NOTE: We only alert if we actually have the GUID activity enabled.
81- if let Some ( sample_function) = view. functions ( ) . iter ( ) . next ( ) {
82- let function_workflow = sample_function
83- . workflow ( )
84- . expect ( "Function has no workflow" ) ;
85- if function_workflow. contains ( GUID_ACTIVITY_NAME ) {
86- log:: error!( "No function guids in database, please reanalyze the database." ) ;
87- } else {
88- log:: error!(
89- "Activity '{}' is not in workflow '{}', create function guids manually to run matcher..." ,
90- GUID_ACTIVITY_NAME ,
91- function_workflow. name( )
92- )
93- }
94- }
125+ let Some ( function_set) = FunctionSet :: from_view ( view) else {
95126 background_task. finish ( ) ;
96127 return ;
97- }
98-
99- // TODO: Par iter this? Using dashmap
100- let guids_by_target: HashMap < Target , Vec < FunctionGUID > > = functions_by_target_and_guid
101- . keys ( )
102- . map ( |( guid, target) | ( target. clone ( ) , * guid) )
103- . into_group_map ( ) ;
128+ } ;
104129
105130 // TODO: Target gets cloned a lot.
106131 // TODO: Containers might both match on the same function. What should we do?
@@ -109,7 +134,7 @@ pub fn run_matcher(view: &BinaryView) {
109134 return ;
110135 }
111136
112- for ( target, guids) in & guids_by_target {
137+ for ( target, guids) in & function_set . guids_by_target {
113138 let function_guid_with_sources = container
114139 . sources_with_function_guids ( target, guids)
115140 . unwrap_or_default ( ) ;
@@ -140,7 +165,8 @@ pub fn run_matcher(view: &BinaryView) {
140165 return ;
141166 }
142167
143- let functions = functions_by_target_and_guid
168+ let functions = function_set
169+ . functions_by_target_and_guid
144170 . get ( & ( guid, target. clone ( ) ) )
145171 . expect ( "Function guid not found" ) ;
146172
@@ -169,6 +195,46 @@ pub fn run_matcher(view: &BinaryView) {
169195 view. update_analysis ( ) ;
170196}
171197
198+ pub fn run_fetcher ( view : & BinaryView ) {
199+ let background_task = BackgroundTask :: new ( "Fetching WARP functions..." , true ) ;
200+ let start = Instant :: now ( ) ;
201+
202+ // Build matcher
203+ let view_settings = Settings :: new ( ) ;
204+ let mut query_opts = QueryOptions :: new_with_view ( view) ;
205+ let plugin_settings = PluginSettings :: from_settings ( & view_settings, & mut query_opts) ;
206+
207+ let Some ( function_set) = FunctionSet :: from_view ( view) else {
208+ background_task. finish ( ) ;
209+ return ;
210+ } ;
211+
212+ for_cached_containers_mut ( |container| {
213+ if background_task. is_cancelled ( ) {
214+ return ;
215+ }
216+
217+ for ( target, guids) in & function_set. guids_by_target {
218+ for batch in guids. chunks ( plugin_settings. fetch_batch_size ) {
219+ if background_task. is_cancelled ( ) {
220+ break ;
221+ }
222+ let _ =
223+ container. fetch_functions ( target, & plugin_settings. allowed_source_tags , batch) ;
224+ }
225+ }
226+ } ) ;
227+
228+ if background_task. is_cancelled ( ) {
229+ log:: info!(
230+ "Fetcher was cancelled by user, you may run it again by running the 'Fetch' command."
231+ ) ;
232+ }
233+
234+ log:: info!( "Fetching took {:?}" , start. elapsed( ) ) ;
235+ background_task. finish ( ) ;
236+ }
237+
172238pub fn insert_workflow ( ) -> Result < ( ) , ( ) > {
173239 // TODO: Note: because of symbol persistence function symbol is applied in `insert_cached_function_match`.
174240 // TODO: Comments are also applied there, they are "user" like, persisted and make undo actions.
@@ -232,6 +298,11 @@ pub fn insert_workflow() -> Result<(), ()> {
232298 run_matcher ( & view) ;
233299 } ;
234300
301+ let fetcher_activity = |ctx : & AnalysisContext | {
302+ let view = ctx. view ( ) ;
303+ run_fetcher ( & view) ;
304+ } ;
305+
235306 let guid_activity = |ctx : & AnalysisContext | {
236307 let function = ctx. function ( ) ;
237308 cached_function_guid ( & function, || unsafe { ctx. lifted_il_function ( ) } ) ;
@@ -269,6 +340,14 @@ pub fn insert_workflow() -> Result<(), ()> {
269340 // TODO: Remove this once the objectivec workflow is registered on the meta workflow.
270341 add_function_activities ( Workflow :: cloned ( "core.function.objectiveC" ) ) ?;
271342
343+ let fetcher_config = activity:: Config :: action (
344+ "analysis.warp.fetcher" ,
345+ "WARP Fetcher" ,
346+ "This analysis step attempts to fetch WARP functions from network containers after the initial analysis is complete..." ,
347+ )
348+ . eligibility ( activity:: Eligibility :: auto_with_default ( false ) . run_once ( true ) ) ;
349+ let fetcher_activity = Activity :: new_with_action ( & fetcher_config, fetcher_activity) ;
350+
272351 let matcher_config = activity:: Config :: action (
273352 "analysis.warp.matcher" ,
274353 "WARP Matcher" ,
@@ -281,6 +360,7 @@ pub fn insert_workflow() -> Result<(), ()> {
281360 Workflow :: cloned ( "core.module.metaAnalysis" )
282361 . ok_or ( ( ) ) ?
283362 . activity_before ( & matcher_activity, "core.module.finishUpdate" ) ?
363+ . activity_before ( & fetcher_activity, "analysis.warp.matcher" ) ?
284364 . register ( ) ?;
285365
286366 Ok ( ( ) )
0 commit comments