[WARP] Recover register variables and re-add tags to matched functions

3 files changed, 103 insertions, 22 deletions

diff --git a/plugins/warp/src/convert.rs b/plugins/warp/src/convert.rs
index 8272768b..bacf0024 100644
--- a/plugins/warp/src/convert.rs
+++ b/plugins/warp/src/convert.rs
@@ -5,11 +5,31 @@ use binaryninja::function::Comment as BNComment;
 use binaryninja::function::Function as BNFunction;
 use binaryninja::platform::Platform;
 use binaryninja::rc::Ref;
+use binaryninja::variable::{Variable as BNVariable, VariableSourceType};
 pub use symbol::*;
 pub use types::*;
+use warp::r#type::class::function::{Location, RegisterLocation, StackLocation};
 use warp::signature::comment::FunctionComment;
 use warp::target::Target;
 
+pub fn bn_var_to_location(bn_variable: BNVariable) -> Option<Location> {
+    match bn_variable.ty {
+        VariableSourceType::StackVariableSourceType => {
+            let stack_loc = StackLocation {
+                offset: bn_variable.storage,
+            };
+            Some(Location::Stack(stack_loc))
+        }
+        VariableSourceType::RegisterVariableSourceType => {
+            let reg_loc = RegisterLocation {
+                id: bn_variable.storage as u64,
+            };
+            Some(Location::Register(reg_loc))
+        }
+        VariableSourceType::FlagVariableSourceType => None,
+    }
+}
+
 pub fn bn_comment_to_comment(func: &BNFunction, bn_comment: BNComment) -> FunctionComment {
     let offset = (bn_comment.addr as i64) - (func.start() as i64);
     FunctionComment {
diff --git a/plugins/warp/src/lib.rs b/plugins/warp/src/lib.rs
index 852d64d6..2172ac5d 100644
--- a/plugins/warp/src/lib.rs
+++ b/plugins/warp/src/lib.rs
@@ -1,5 +1,5 @@
 use crate::cache::{cached_constraints, cached_function_guid};
-use crate::convert::{bn_comment_to_comment, from_bn_symbol, from_bn_type};
+use crate::convert::{bn_comment_to_comment, bn_var_to_location, from_bn_symbol, from_bn_type};
 use binaryninja::architecture::{
     Architecture, ImplicitRegisterExtend, Register as BNRegister, RegisterInfo,
 };
@@ -16,13 +16,15 @@ use binaryninja::low_level_il::instruction::{
 };
 use binaryninja::low_level_il::{LowLevelILRegisterKind, VisitorAction};
 use binaryninja::rc::{Ref as BNRef, Ref};
+use binaryninja::tags::TagType;
+use binaryninja::variable::RegisterValueType;
+use itertools::Itertools;
 use std::ops::Range;
 use std::path::PathBuf;
 use warp::signature::basic_block::BasicBlockGUID;
 use warp::signature::function::{Function, FunctionGUID};
+use warp::signature::variable::FunctionVariable;
 
-use binaryninja::tags::TagType;
-use binaryninja::variable::RegisterValueType;
 /// Re-export the warp crate that is used, this is useful for consumers of this crate.
 pub use warp;
 
@@ -71,6 +73,38 @@ pub fn user_signature_dir() -> PathBuf {
     binaryninja::user_directory().join("signatures/")
 }
 
+pub fn build_variables(func: &BNFunction) -> Vec<FunctionVariable> {
+    let func_start = func.start();
+    let mut variables = vec![];
+    if let Ok(mlil) = func.medium_level_il() {
+        let bn_vars = func.variables();
+        for bn_var in &bn_vars {
+            if mlil.is_var_user_defined(&bn_var.variable) {
+                // TODO: live_instruction_for_variable only works for register types.
+                if let Some(first_instr) = mlil
+                    .live_instruction_for_variable(&bn_var.variable, true)
+                    .iter()
+                    .sorted_by_key(|i| i.instr_index)
+                    .next()
+                {
+                    if let Some(var_loc) = bn_var_to_location(bn_var.variable) {
+                        let var_name = bn_var.name;
+                        let var_type =
+                            from_bn_type(&func.view(), &bn_var.ty.contents, bn_var.ty.confidence);
+                        variables.push(FunctionVariable {
+                            offset: (first_instr.address as i64) - (func_start as i64),
+                            location: var_loc,
+                            name: Some(var_name),
+                            ty: Some(var_type),
+                        })
+                    }
+                }
+            }
+        }
+    }
+    variables
+}
+
 pub fn build_function<M: FunctionMutability>(
     func: &BNFunction,
     lifted_il: &LowLevelILFunction<M, NonSSA>,
@@ -97,10 +131,7 @@ pub fn build_function<M: FunctionMutability>(
         // NOTE: We do not filter out adjacent functions here.
         constraints: cached_constraints(func, |_| true),
         comments,
-        // TODO: Gather relevant variables (only user?).
-        // TODO: Will need MLIL SSA for this to locate def sites.
-        // TODO: Add this info in a second pass?
-        variables: vec![],
+        variables: build_variables(func),
     }
 }
 
diff --git a/plugins/warp/src/plugin/workflow.rs b/plugins/warp/src/plugin/workflow.rs
index c83ae809..db95b567 100644
--- a/plugins/warp/src/plugin/workflow.rs
+++ b/plugins/warp/src/plugin/workflow.rs
@@ -8,6 +8,7 @@ use crate::convert::{
 };
 use crate::matcher::{Matcher, MatcherSettings};
 use crate::{get_warp_tag_type, relocatable_regions};
+use binaryninja::architecture::RegisterId;
 use binaryninja::background_task::BackgroundTask;
 use binaryninja::binary_view::{BinaryView, BinaryViewExt};
 use binaryninja::command::Command;
@@ -16,6 +17,7 @@ use binaryninja::workflow::{Activity, AnalysisContext, Workflow};
 use itertools::Itertools;
 use std::collections::HashMap;
 use std::time::Instant;
+use warp::r#type::class::function::{Location, RegisterLocation, StackLocation};
 use warp::signature::function::{Function, FunctionGUID};
 use warp::target::Target;
 
@@ -68,6 +70,9 @@ impl Command for RunMatcher {
 }
 
 pub fn run_matcher(view: &BinaryView) {
+    // TODO: Create the tag type so we dont have UB in the apply function workflow.
+    let _ = get_warp_tag_type(view);
+
     // Alert the user if we have no actual regions (one comes from the synthetic section).
     let regions = relocatable_regions(view);
     if regions.len() <= 1 && view.memory_map().is_activated() {
@@ -174,17 +179,41 @@ pub fn insert_workflow() {
                 let bn_comment = comment_to_bn_comment(&function, comment);
                 function.set_comment_at(bn_comment.addr, &bn_comment.comment);
             }
-            // TODO: Fix this before release.
-            // TODO: Any attempt to add a tag type will create a undo action
-            // TODO: Those are currently not thread safe when running in headless python.
-            // TODO: See Mason for more lore.
-            // function.add_tag(
-            //     &get_warp_tag_type(&view),
-            //     &matched_function.guid.to_string(),
-            //     None,
-            //     false,
-            //     None,
-            // );
+            if let Some(mlil) = ctx.mlil_function() {
+                for variable in matched_function.variables {
+                    let decl_addr = ((function.start() as i64) + variable.offset) as u64;
+                    if let Some(decl_instr) = mlil.instruction_at(decl_addr) {
+                        let decl_var = match variable.location {
+                            Location::Register(RegisterLocation { id, .. }) => {
+                                decl_instr.variable_for_register_after(RegisterId(id as u32))
+                            }
+                            Location::Stack(StackLocation { offset, .. }) => {
+                                decl_instr.variable_for_stack_location_after(offset)
+                            }
+                        };
+                        let decl_ty = match variable.ty {
+                            Some(decl_ty) => to_bn_type(&function.arch(), &decl_ty),
+                            None => {
+                                let Some(existing_var) = function.variable_type(&decl_var) else {
+                                    continue;
+                                };
+                                existing_var.contents
+                            }
+                        };
+                        let decl_name = variable
+                            .name
+                            .unwrap_or_else(|| function.variable_name(&decl_var));
+                        mlil.create_auto_var(&decl_var, &decl_ty, &decl_name, false)
+                    }
+                }
+            }
+            function.add_tag(
+                &get_warp_tag_type(&view),
+                &matched_function.guid.to_string(),
+                None,
+                false,
+                None,
+            );
         }
     };
 
@@ -208,12 +237,13 @@ pub fn insert_workflow() {
         .register_activity(&guid_activity)
         .unwrap();
     // Because we are going to impact analysis with application we must make sure the function update is triggered to continue to update analysis.
-    // TODO: need to ask why i cant do core.function.update like in the rtti plugin.
     function_meta_workflow
-        .register_activity_with_subactivities::<Vec<String>>(&apply_activity, vec![])
+        .register_activity(&apply_activity)
         .unwrap();
-    function_meta_workflow.insert("core.function.runFunctionRecognizers", [GUID_ACTIVITY_NAME]);
-    function_meta_workflow.insert("core.function.generateMediumLevelIL", [APPLY_ACTIVITY_NAME]);
+    function_meta_workflow
+        .insert_after("core.function.runFunctionRecognizers", [GUID_ACTIVITY_NAME]);
+    function_meta_workflow
+        .insert_after("core.function.generateMediumLevelIL", [APPLY_ACTIVITY_NAME]);
     function_meta_workflow.register().unwrap();
 
     let old_module_meta_workflow = Workflow::instance("core.module.metaAnalysis");