Make WARP function guid generation infallible

9 files changed, 84 insertions, 79 deletions

diff --git a/plugins/warp/src/bin/sigem.rs b/plugins/warp/src/bin/sigem.rs
index f0551ccc..3eaef0b1 100644
--- a/plugins/warp/src/bin/sigem.rs
+++ b/plugins/warp/src/bin/sigem.rs
@@ -78,7 +78,7 @@ fn data_from_view(view: &BinaryView) -> Data {
         .filter(|f| !f.symbol().short_name().as_str().contains("sub_") || f.has_user_annotations())
         .filter_map(|f| {
             let llil = f.low_level_il().ok()?;
-            warp_ninja::cache::cached_function(&f, &llil)
+            Some(warp_ninja::cache::cached_function(&f, &llil))
         })
         .collect::<Vec<_>>();
 
diff --git a/plugins/warp/src/cache.rs b/plugins/warp/src/cache.rs
index 79e974b4..bafc4ede 100644
--- a/plugins/warp/src/cache.rs
+++ b/plugins/warp/src/cache.rs
@@ -22,7 +22,7 @@ pub static GUID_CACHE: OnceLock<DashMap<ViewID, GUIDCache>> = OnceLock::new();
 pub fn cached_function<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
     function: &BNFunction,
     llil: &llil::Function<A, M, NonSSA<V>>,
-) -> Option<Function> {
+) -> Function {
     let view = function.view();
     let view_id = ViewID::from(view.as_ref());
     let function_cache = FUNCTION_CACHE.get_or_init(Default::default);
@@ -70,7 +70,7 @@ pub fn cached_adjacency_constraints(function: &BNFunction) -> HashSet<FunctionCo
 pub fn cached_function_guid<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
     function: &BNFunction,
     llil: &llil::Function<A, M, NonSSA<V>>,
-) -> Option<FunctionGUID> {
+) -> FunctionGUID {
     let view = function.view();
     let view_id = ViewID::from(view);
     let guid_cache = GUID_CACHE.get_or_init(Default::default);
@@ -87,7 +87,7 @@ pub fn cached_function_guid<A: Architecture, M: FunctionMutability, V: NonSSAVar
 
 #[derive(Clone, Debug, Default)]
 pub struct FunctionCache {
-    pub cache: DashMap<FunctionID, Option<Function>>,
+    pub cache: DashMap<FunctionID, Function>,
 }
 
 impl FunctionCache {
@@ -95,7 +95,7 @@ impl FunctionCache {
         &self,
         function: &BNFunction,
         llil: &llil::Function<A, M, NonSSA<V>>,
-    ) -> Option<Function> {
+    ) -> Function {
         let function_id = FunctionID::from(function);
         match self.cache.try_get_mut(&function_id) {
             TryResult::Present(function) => function.value().to_owned(),
@@ -111,7 +111,7 @@ impl FunctionCache {
 
 #[derive(Clone, Debug, Default)]
 pub struct GUIDCache {
-    pub cache: DashMap<FunctionID, Option<FunctionGUID>>,
+    pub cache: DashMap<FunctionID, FunctionGUID>,
 }
 
 impl GUIDCache {
@@ -125,16 +125,16 @@ impl GUIDCache {
                 let cs_ref_func = cs_ref.function();
                 let cs_ref_func_id = FunctionID::from(cs_ref_func);
                 if cs_ref_func_id != func_id {
-                    if let Some(cs_ref_func_llil) = cs_ref_func.low_level_il_if_available() {
-                        // Function references another function, constrain on the pattern.
-                        // TODO: If function is trivial thunk we should _also_ insert the tailcall target as a constraint.
-                        let call_site_offset: i64 = func_start as i64 - call_site.address as i64;
-                        constraints.insert(self.function_constraint(
+                    let call_site_offset: i64 = func_start as i64 - call_site.address as i64;
+                    let function_constraint = match cs_ref_func.low_level_il_if_available() {
+                        Some(cs_ref_func_llil) => self.function_constraint_with_guid(
                             cs_ref_func,
                             &cs_ref_func_llil,
                             call_site_offset,
-                        ));
-                    }
+                        ),
+                        None => self.function_constraint(cs_ref_func, call_site_offset),
+                    };
+                    constraints.insert(function_constraint);
                 }
             }
         }
@@ -154,15 +154,17 @@ impl GUIDCache {
                 if curr_func_id != func_id {
                     // NOTE: We have to get the llil here for the function which is problematic for running
                     // NOTE: within a workflow (before analysis has finished)
-                    if let Some(curr_func_llil) = curr_func.low_level_il_if_available() {
-                        // Function adjacent to another function, constrain on the pattern.
-                        let curr_addr_offset = (func_start_addr as i64) - func_start as i64;
-                        constraints.insert(self.function_constraint(
+                    // Function adjacent to another function, constrain on the pattern.
+                    let curr_addr_offset = (func_start_addr as i64) - func_start as i64;
+                    let function_constraint = match curr_func.low_level_il_if_available() {
+                        Some(curr_func_llil) => self.function_constraint_with_guid(
                             &curr_func,
                             &curr_func_llil,
                             curr_addr_offset,
-                        ));
-                    }
+                        ),
+                        None => self.function_constraint(&curr_func, curr_addr_offset),
+                    };
+                    constraints.insert(function_constraint);
                 }
             }
         };
@@ -183,7 +185,21 @@ impl GUIDCache {
     }
 
     /// Construct a function constraint, must pass the offset at which it is located.
-    pub fn function_constraint<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
+    pub fn function_constraint(&self, function: &BNFunction, offset: i64) -> FunctionConstraint {
+        let symbol = from_bn_symbol(&function.symbol());
+        FunctionConstraint {
+            guid: None,
+            symbol: Some(symbol),
+            offset,
+        }
+    }
+
+    /// Construct a function constraint, must pass the offset at which it is located.
+    pub fn function_constraint_with_guid<
+        A: Architecture,
+        M: FunctionMutability,
+        V: NonSSAVariant,
+    >(
         &self,
         function: &BNFunction,
         llil: &llil::Function<A, M, NonSSA<V>>,
@@ -192,7 +208,7 @@ impl GUIDCache {
         let guid = self.function_guid(function, llil);
         let symbol = from_bn_symbol(&function.symbol());
         FunctionConstraint {
-            guid,
+            guid: Some(guid),
             symbol: Some(symbol),
             offset,
         }
@@ -202,7 +218,7 @@ impl GUIDCache {
         &self,
         function: &BNFunction,
         llil: &llil::Function<A, M, NonSSA<V>>,
-    ) -> Option<FunctionGUID> {
+    ) -> FunctionGUID {
         let function_id = FunctionID::from(function);
         match self.cache.try_get_mut(&function_id) {
             TryResult::Present(function_guid) => function_guid.value().to_owned(),
diff --git a/plugins/warp/src/lib.rs b/plugins/warp/src/lib.rs
index ae14ae80..5f8d4a78 100644
--- a/plugins/warp/src/lib.rs
+++ b/plugins/warp/src/lib.rs
@@ -23,10 +23,10 @@ mod plugin;
 pub fn build_function<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
     func: &BNFunction,
     llil: &llil::Function<A, M, NonSSA<V>>,
-) -> Option<Function> {
+) -> Function {
     let bn_fn_ty = func.function_type();
-    Some(Function {
-        guid: cached_function_guid(func, llil)?,
+    Function {
+        guid: cached_function_guid(func, llil),
         symbol: from_bn_symbol(&func.symbol()),
         // TODO: Confidence should be derived from function type.
         ty: from_bn_type(&func.view(), bn_fn_ty, 255),
@@ -40,7 +40,7 @@ pub fn build_function<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
         },
         // TODO: We need more than one entry block.
         entry: entry_basic_block_guid(func, llil).map(BasicBlock::new),
-    })
+    }
 }
 
 pub fn entry_basic_block_guid<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
@@ -49,7 +49,7 @@ pub fn entry_basic_block_guid<A: Architecture, M: FunctionMutability, V: NonSSAV
 ) -> Option<BasicBlockGUID> {
     // NOTE: This is not actually the entry point. This is the highest basic block.
     let first_basic_block = sorted_basic_blocks(func).into_iter().next()?;
-    basic_block_guid(&first_basic_block, llil)
+    Some(basic_block_guid(&first_basic_block, llil))
 }
 
 /// Basic blocks sorted from high to low.
@@ -66,20 +66,20 @@ pub fn sorted_basic_blocks(func: &BNFunction) -> Vec<BNRef<BNBasicBlock<NativeBl
 pub fn function_guid<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
     func: &BNFunction,
     llil: &llil::Function<A, M, NonSSA<V>>,
-) -> Option<FunctionGUID> {
+) -> FunctionGUID {
     // TODO: Sort the basic blocks.
     let basic_blocks = sorted_basic_blocks(func);
     let basic_block_guids = basic_blocks
         .iter()
-        .filter_map(|bb| basic_block_guid(bb, llil))
+        .map(|bb| basic_block_guid(bb, llil))
         .collect::<Vec<_>>();
-    Some(FunctionGUID::from_basic_blocks(&basic_block_guids))
+    FunctionGUID::from_basic_blocks(&basic_block_guids)
 }
 
 pub fn basic_block_guid<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
     basic_block: &BNBasicBlock<NativeBlock>,
     llil: &llil::Function<A, M, NonSSA<V>>,
-) -> Option<BasicBlockGUID> {
+) -> BasicBlockGUID {
     let func = basic_block.function();
     let view = func.view();
     let arch = func.arch();
@@ -112,7 +112,7 @@ pub fn basic_block_guid<A: Architecture, M: FunctionMutability, V: NonSSAVariant
         }
     }
 
-    Some(BasicBlockGUID::from(basic_block_bytes.as_slice()))
+    BasicBlockGUID::from(basic_block_bytes.as_slice())
 }
 
 #[cfg(test)]
@@ -146,9 +146,7 @@ mod tests {
                             let mut functions = inital_bv
                                 .functions()
                                 .iter()
-                                .map(|f| {
-                                    cached_function_guid(&f, &f.low_level_il().unwrap()).unwrap()
-                                })
+                                .map(|f| cached_function_guid(&f, &f.low_level_il().unwrap()))
                                 .collect::<Vec<_>>();
                             functions.sort_by_key(|guid| guid.guid);
                             insta::assert_debug_snapshot!(functions);
diff --git a/plugins/warp/src/matcher.rs b/plugins/warp/src/matcher.rs
index ae49f00a..01222a8b 100644
--- a/plugins/warp/src/matcher.rs
+++ b/plugins/warp/src/matcher.rs
@@ -286,26 +286,25 @@ impl Matcher {
         let entry_block = entry_basic_block_guid(function, llil).map(BasicBlock::new);
         if self.basic_block_filter.contains(&entry_block) {
             // Build the full function guid now
-            if let Some(warp_func_guid) = cached_function_guid(function, llil) {
-                if let Some(matched) = self.functions.get(&warp_func_guid) {
-                    if matched.len() == 1 && !is_function_trivial {
-                        on_new_match(&matched[0]);
-                    } else if let Some(matched_function) =
-                        self.match_function_from_constraints(function, &matched)
-                    {
-                        log::info!(
-                            "Found best matching function `{}`... 0x{:x}",
-                            matched_function.symbol.name,
-                            function.start()
-                        );
-                        on_new_match(matched_function);
-                    } else {
-                        log::error!(
-                            "Failed to find matching function `{}`... 0x{:x}",
-                            matched.len(),
-                            function.start()
-                        );
-                    }
+            let warp_func_guid = cached_function_guid(function, llil);
+            if let Some(matched) = self.functions.get(&warp_func_guid) {
+                if matched.len() == 1 && !is_function_trivial {
+                    on_new_match(&matched[0]);
+                } else if let Some(matched_function) =
+                    self.match_function_from_constraints(function, &matched)
+                {
+                    log::info!(
+                        "Found best matching function `{}`... 0x{:x}",
+                        matched_function.symbol.name,
+                        function.start()
+                    );
+                    on_new_match(matched_function);
+                } else {
+                    log::error!(
+                        "Failed to find matching function `{}`... 0x{:x}",
+                        matched.len(),
+                        function.start()
+                    );
                 }
             }
         }
diff --git a/plugins/warp/src/plugin.rs b/plugins/warp/src/plugin.rs
index 5bca373b..3373cafd 100644
--- a/plugins/warp/src/plugin.rs
+++ b/plugins/warp/src/plugin.rs
@@ -54,9 +54,7 @@ struct DebugFunction;
 impl FunctionCommand for DebugFunction {
     fn action(&self, _view: &BinaryView, func: &Function) {
         if let Ok(llil) = func.low_level_il() {
-            if let Some(function) = build_function(func, &llil) {
-                log::info!("{:#?}", function);
-            }
+            log::info!("{:#?}", build_function(func, &llil));
         }
     }
 
diff --git a/plugins/warp/src/plugin/apply.rs b/plugins/warp/src/plugin/apply.rs
index 56c09d77..78cb3be8 100644
--- a/plugins/warp/src/plugin/apply.rs
+++ b/plugins/warp/src/plugin/apply.rs
@@ -60,7 +60,7 @@ impl Command for ApplySignatureFile {
             .par_iter()
             .filter_map(|func| {
                 let llil = func.low_level_il_if_available()?;
-                let pattern = cached_function_guid(&func, &llil)?;
+                let pattern = cached_function_guid(&func, &llil);
                 Some((func, data_functions.get(&pattern)?))
             })
             .filter(|(_, sig)| sig.len() == 1)
diff --git a/plugins/warp/src/plugin/copy.rs b/plugins/warp/src/plugin/copy.rs
index ee91fb4a..d539a430 100644
--- a/plugins/warp/src/plugin/copy.rs
+++ b/plugins/warp/src/plugin/copy.rs
@@ -12,20 +12,14 @@ impl FunctionCommand for CopyFunctionGUID {
             log::error!("Could not get low level il for copied function");
             return;
         };
-        if let Some(guid) = cached_function_guid(func, &llil) {
-            log::info!(
-                "Function GUID for {}... {}",
-                func.symbol().short_name().to_string(),
-                guid
-            );
-            if let Ok(mut clipboard) = arboard::Clipboard::new() {
-                let _ = clipboard.set_text(guid.to_string());
-            }
-        } else {
-            log::error!(
-                "Failed to create GUID for function... 0x{:0x}",
-                func.start()
-            );
+        let guid = cached_function_guid(func, &llil);
+        log::info!(
+            "Function GUID for {}... {}",
+            func.symbol().short_name().to_string(),
+            guid
+        );
+        if let Ok(mut clipboard) = arboard::Clipboard::new() {
+            let _ = clipboard.set_text(guid.to_string());
         }
     }
 
diff --git a/plugins/warp/src/plugin/create.rs b/plugins/warp/src/plugin/create.rs
index 2b1ad60e..b404ae1c 100644
--- a/plugins/warp/src/plugin/create.rs
+++ b/plugins/warp/src/plugin/create.rs
@@ -31,11 +31,11 @@ impl Command for CreateSignatureFile {
             let start = Instant::now();
 
             let mut data = warp::signature::Data::default();
-            data.functions.par_extend(
-                view.functions()
-                    .par_iter()
-                    .filter_map(|func| cached_function(&func, func.low_level_il().ok()?.as_ref())),
-            );
+            data.functions
+                .par_extend(view.functions().par_iter().filter_map(|func| {
+                    let llil = func.low_level_il().ok()?;
+                    Some(cached_function(&func, &llil))
+                }));
             data.types.extend(view.types().iter().map(|ty| {
                 let ref_ty = ty.type_object().to_owned();
                 ComputedType::new(from_bn_type(&view, ref_ty, u8::MAX))
diff --git a/plugins/warp/src/plugin/find.rs b/plugins/warp/src/plugin/find.rs
index 7f089c2f..6d6c5bc1 100644
--- a/plugins/warp/src/plugin/find.rs
+++ b/plugins/warp/src/plugin/find.rs
@@ -37,7 +37,7 @@ impl Command for FindFunctionFromGUID {
                 .filter_map(|func| {
                     Some((
                         func.clone(),
-                        cached_function_guid(&func, func.low_level_il_if_available()?.as_ref())?,
+                        cached_function_guid(&func, func.low_level_il_if_available()?.as_ref()),
                     ))
                 })
                 .filter(|(_func, guid)| guid.eq(&searched_guid))