Cache function list from finfo data

Cargo.toml

@@ -6,7 +6,7 @@ edition = "2021"
 [dependencies]
 r2pipe = "0.6.0"
 serde = { version = "1.0", features = ["derive"] }
-serde_json = "1.0.60"
+serde_json = { version = "1.0.60", features = ["raw_value"] }
 clap = { version = "4.0.29", features = ["derive"] }
 goblin = { version = "0.6.0", optional = true }
 walkdir = "2"
@@ -19,7 +19,8 @@ tokenizers = "0.13.2"
 rust_search = "2.1.0"
 tch = {version = "0.10.3", optional = true}
 mimalloc = { version = "*", default-features = false }
-indicatif = { version = "0.17.3", features = ["rayon"]}
+indicatif = { version = "0.18.0", features = ["rayon"]}
+indicatif-log-bridge = "0.2.3"
 itertools = "0.10.5"
 prettytable = "0.10.0"
 serde-aux = "4.2.0"
@@ -32,6 +33,8 @@ hex = { version = "0.4", features = ["alloc"] }
 md5 = "0.7.0"
 sha1 = "0.10.6"
 sha2 = "0.10.8"
+csv = "1.3"
+once_cell = "1.19"
 
 [dependencies.petgraph]
 version = "0.6.2"

src/afij.rs

@@ -6,9 +6,10 @@ use serde_json::Value;
 #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct AFIJFunctionInfo {
+    #[serde(alias = "addr")]
     pub offset: u64,
     pub name: String,
-    pub size: i128,
+    pub size: u64,
     #[serde(rename = "is-pure")]
     pub is_pure: String,
     pub realsz: u64,
@@ -23,21 +24,23 @@ pub struct AFIJFunctionInfo {
     pub nbbs: u64,
     #[serde(rename = "is-lineal")]
     pub is_lineal: bool,
-    pub ninstrs: i64,
-    pub edges: i64,
+    pub ninstrs: u64,
+    pub edges: u64,
     pub ebbs: u64,
     pub signature: String,
+    #[serde(alias = "minaddr")]
     pub minbound: u64,
-    pub maxbound: i128,
+    #[serde(alias = "maxaddr")]
+    pub maxbound: u64,
     pub callrefs: Option<Vec<Callref>>,
     // TODO: Need to fix this and change to string instead of i64 to get round large random numbers
     pub datarefs: Option<Vec<Dataref>>,
     pub codexrefs: Option<Vec<Codexref>>,
-    pub dataxrefs: Option<Vec<i64>>,
-    pub indegree: Option<i64>,
-    pub outdegree: Option<i64>,
-    pub nlocals: Option<i64>,
-    pub nargs: Option<i64>,
+    pub dataxrefs: Option<Vec<i64>>, // TODO: Check this data type
+    pub indegree: Option<u64>,
+    pub outdegree: Option<u64>,
+    pub nlocals: Option<u64>,
+    pub nargs: Option<u64>,
     pub bpvars: Option<Vec<Bpvar>>,
     // Cannot find a good example of an spvars yet
     pub spvars: Option<Vec<Value>>,
@@ -103,12 +106,12 @@ pub struct Regvar {
 #[serde(rename_all = "camelCase")]
 pub struct AFIJFeatureSubset {
     pub name: String,
-    pub ninstrs: i64,
-    pub edges: i64,
-    pub indegree: i64,
-    pub outdegree: i64,
-    pub nlocals: i64,
-    pub nargs: i64,
+    pub ninstrs: u64,
+    pub edges: u64,
+    pub indegree: u64,
+    pub outdegree: u64,
+    pub nlocals: u64,
+    pub nargs: u64,
     pub signature: String,
 }
 
@@ -130,12 +133,12 @@ impl From<&AFIJFunctionInfo> for AFIJFeatureSubset {
 #[derive(Default, Debug, Clone, PartialEq, Hash, Serialize, Deserialize)]
 pub struct AFIJFeatureSubsetExtended {
     pub name: String,
-    pub ninstrs: i64,
-    pub edges: i64,
-    pub indegree: i64,
-    pub outdegree: i64,
-    pub nlocals: i64,
-    pub nargs: i64,
+    pub ninstrs: u64,
+    pub edges: u64,
+    pub indegree: u64,
+    pub outdegree: u64,
+    pub nlocals: u64,
+    pub nargs: u64,
     pub nbbs: u64,
     pub avg_ins_bb: OrderedFloat<f32>,
 }

src/agcj.rs

@@ -3,7 +3,7 @@ use crate::networkx::{
     CallGraphFuncNameNode, CallGraphFuncWithMetadata, CallGraphTikNibFeatures,
     CallGraphTikNibFinfoFeatures, NetworkxDiGraph,
 };
-use crate::utils::{check_or_create_dir, get_save_file_path};
+use crate::utils::{check_or_create_dir, get_save_file_path, sanitize_filename};
 use itertools::Itertools;
 use petgraph::prelude::Graph;
 use serde::{Deserialize, Serialize};
@@ -14,14 +14,14 @@ use std::path::{Path, PathBuf};
 #[serde(rename_all = "camelCase")]
 pub struct AGCJFunctionCallGraph {
     pub name: String,
-    pub size: i64,
+    pub size: u64,
     pub imports: Option<Vec<String>>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct AGCJParsedObjects {
     pub edge_property: String,
-    pub edges: Vec<Vec<i32>>,
+    pub edges: Vec<Vec<u32>>,
     pub node_holes: Vec<String>,
     pub nodes: Vec<String>,
 }
@@ -45,10 +45,7 @@ impl AGCJFunctionCallGraph {
 
         let mut function_name = self.name.clone();
 
-        // This is a pretty dirty fix and may break things
-        if function_name.chars().count() > 100 {
-            function_name = self.name[..75].to_string();
-        }
+        function_name = sanitize_filename(&function_name);
 
         let filename = format!("{}-{}.json", function_name, type_suffix);
 
@@ -82,11 +79,7 @@ impl AGCJFunctionCallGraph {
         check_or_create_dir(&full_output_path);
 
         let mut function_name = self.name.clone();
-
-        // This is a pretty dirty fix and may break things
-        if function_name.chars().count() > 100 {
-            function_name = self.name[..75].to_string();
-        }
+        function_name = sanitize_filename(&function_name);
 
         let filename = format!(
             "{}/{}-{}.json",
@@ -121,11 +114,7 @@ impl AGCJFunctionCallGraph {
         check_or_create_dir(&full_output_path);
 
         let mut function_name = self.name.clone();
-
-        // This is a pretty dirty fix and may break things
-        if function_name.chars().count() > 100 {
-            function_name = self.name[..75].to_string();
-        }
+        function_name = sanitize_filename(&function_name);
 
         let filename = format!(
             "{}/{}-{}.json",
@@ -159,11 +148,7 @@ impl AGCJFunctionCallGraph {
         check_or_create_dir(&full_output_path);
         debug!("Built Path: {:?}", full_output_path);
         let mut function_name = self.name.clone();
-
-        // This is a pretty dirty fix and may break things
-        if function_name.chars().count() > 100 {
-            function_name = self.name[..75].to_string();
-        }
+        function_name = sanitize_filename(&function_name);
 
         let filename = format!("{}-{}.json", function_name, type_suffix);
         // Normalise string for windows

src/agfj.rs

@@ -35,18 +35,19 @@ struct EdgePair {
     wt: u16,
 }
 
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct AGFJFunc {
     pub name: String,
+    #[serde(alias = "addr")]
+    pub offset: u64,
+    pub size: u64,
+    pub ninstr: u64,
     nargs: u64,
-    ninstr: u64,
     nlocals: u64,
-    offset: u64,
-    size: Option<u64>,
     stack: u64,
     r#type: String,
     pub blocks: Vec<ACFJBlock>,
-    addr_idx: Option<Vec<i64>>,
+    addr_idx: Option<Vec<u64>>,
     pub edge_list: Option<Vec<(u32, u32, u32)>>,
     graph: Option<Graph<String, u32>>,
 }
@@ -143,7 +144,7 @@ impl AGFJFunc {
     pub fn create_bb_edge_list(&mut self, min_blocks: &u16) {
         if self.blocks.len() > <u16 as Into<usize>>::into(*min_blocks) && self.blocks[0].offset != 1
         {
-            let bb_start_addrs: Vec<i64> = self.blocks.iter().map(|x| x.offset).collect::<Vec<_>>();
+            let bb_start_addrs: Vec<u64> = self.blocks.iter().map(|x| x.offset).collect::<Vec<_>>();
             let mut edge_list = Vec::<(u32, u32, u32)>::new();
 
             for bb in &self.blocks {
@@ -423,7 +424,7 @@ impl AGFJFunc {
                     }
                 };
 
-                let bb_start_addrs: Vec<i64> =
+                let bb_start_addrs: Vec<u64> =
                     self.blocks.iter().map(|x| x.offset).collect::<Vec<_>>();
 
                 match feature_type {
@@ -606,7 +607,7 @@ impl AGFJFunc {
     }
 
     // Convert string memory address to hex / string
-    fn str_to_hex_node_idxs(graph: &mut Graph<String, u32>, addr_idxs: &[i64]) {
+    fn str_to_hex_node_idxs(graph: &mut Graph<String, u32>, addr_idxs: &[u64]) {
         for idx in graph.node_indices() {
             let i_idx = idx.index();
             let hex = addr_idxs[i_idx];
@@ -732,6 +733,7 @@ impl From<(&String, Vec<TikNibFeaturesBB>)> for TikNibFunc {
 #[cfg(test)]
 mod tests {
     use crate::bb::FeatureType;
+    use crate::utils::get_default_addr;
     use std::path::PathBuf;
 
     use crate::AGFJFile;
@@ -792,7 +794,7 @@ mod tests {
         assert_eq!(file.functions.as_ref().unwrap()[0][0].ninstr, 78);
         assert_eq!(file.functions.as_ref().unwrap()[0][0].nargs, 2);
         assert_eq!(file.functions.as_ref().unwrap()[0][0].nlocals, 0);
-        assert_eq!(file.functions.as_ref().unwrap()[0][0].size, Some(334));
+        assert_eq!(file.functions.as_ref().unwrap()[0][0].size, 334);
         assert_eq!(file.functions.as_ref().unwrap()[0][0].stack, 56);
         assert!(!file.functions.as_ref().unwrap()[0][0].blocks.is_empty());
 
@@ -811,7 +813,10 @@ mod tests {
         assert!(!file.functions.as_ref().unwrap()[0][0].blocks[0]
             .ops
             .is_empty());
-        assert_eq!(file.functions.as_ref().unwrap()[0][0].blocks[0].fail, -1);
+        assert_eq!(
+            file.functions.as_ref().unwrap()[0][0].blocks[0].fail,
+            get_default_addr()
+        );
 
         assert!(file.functions.as_ref().unwrap()[0][0].blocks[0]
             .switchop

src/bb.rs

@@ -2,6 +2,7 @@ use crate::consts::*;
 #[cfg(feature = "inference")]
 use crate::inference::InferenceJob;
 use crate::normalisation::{normalise_disasm_simple, normalise_esil_simple};
+use crate::utils::get_default_addr;
 use serde::{Deserialize, Serialize};
 use serde_aux::prelude::*;
 use serde_json::Value;
@@ -58,8 +59,9 @@ pub enum InstructionMode {
 
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
 pub struct SwitchOpCase {
-    pub jump: i64,
-    pub offset: i64,
+    pub jump: u64,
+    #[serde(alias = "addr")]
+    pub offset: u64,
     #[serde(deserialize_with = "deserialize_string_from_number")]
     // This will make it challenging to use downstream but has been
     // added because sometimes it is a VERY large int (bigger than i64).
@@ -74,6 +76,7 @@ pub struct SwitchOp {
     pub defval: u16,
     pub maxval: u16,
     pub minval: u16,
+    #[serde(alias = "addr")]
     pub offset: u64,
 }
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -86,6 +89,7 @@ pub struct Op {
     pub fcn_addr: Option<u64>,
     pub fcn_last: Option<u64>,
     pub flags: Option<Vec<String>>,
+    #[serde(alias = "addr")]
     pub offset: u64,
     pub opcode: Option<String>,
     pub ptr: Option<u128>,
@@ -100,26 +104,21 @@ pub struct Op {
     pub val: Option<u64>,
 }
 
-// Function to set offset, jump and fail to default values
-fn return_minus_one() -> i64 {
-    -1
-}
-
 #[serde_as]
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct ACFJBlock {
-    #[serde(default = "return_minus_one")]
-    pub offset: i64,
-    #[serde(default = "return_minus_one")]
+    #[serde(default = "get_default_addr", alias = "addr")]
+    pub offset: u64,
+    #[serde(default = "get_default_addr")]
     #[serde_as(deserialize_as = "DefaultOnError")]
     // This has been added to eliminate an error where
     // the jump address from x86-64 binaries is larger than
     // an i64.
-    pub jump: i64,
-    #[serde(default = "return_minus_one")]
-    pub fail: i64,
+    pub jump: u64,
+    #[serde(default = "get_default_addr")]
+    pub fail: u64,
     pub ops: Vec<Op>,
-    pub size: Option<i64>,
+    pub size: Option<u64>,
     pub switchop: Option<SwitchOp>,
 }
 
@@ -457,18 +456,18 @@ impl ACFJBlock {
         }
         num_offspring
     }
-    pub fn get_block_edges(&self, bb_start_addrs: &[i64], edge_list: &mut Vec<(u32, u32, u32)>) {
+    pub fn get_block_edges(&self, bb_start_addrs: &[u64], edge_list: &mut Vec<(u32, u32, u32)>) {
         let offset_idx = bb_start_addrs.iter().position(|&p| p == self.offset);
 
         if let Some(offset_idx) = offset_idx {
-            if self.jump != -1 {
+            if self.jump != get_default_addr() {
                 let jump_idx = bb_start_addrs.iter().position(|&p| p == self.jump);
                 if let Some(jump_idx) = jump_idx {
                     edge_list.push((offset_idx as u32, jump_idx as u32, 1));
                 }
             }
 
-            if self.fail != -1 {
+            if self.fail != get_default_addr() {
                 let fail_idx = bb_start_addrs.iter().position(|&p| p == self.fail);
                 if let Some(fail_idx) = fail_idx {
                     edge_list.push((offset_idx as u32, fail_idx as u32, 1));

src/combos.rs

@@ -133,11 +133,11 @@ impl ComboJob {
 #[derive(Default, Hash, PartialEq, Clone, Debug, Deserialize, Serialize)]
 pub struct FinfoTiknib {
     pub name: String,
-    pub edges: i64,
-    pub indegree: i64,
-    pub outdegree: i64,
-    pub nlocals: i64,
-    pub nargs: i64,
+    pub edges: u64,
+    pub indegree: u64,
+    pub outdegree: u64,
+    pub nlocals: u64,
+    pub nargs: u64,
     pub avg_arithshift: OrderedFloat<f32>,
     pub avg_compare: OrderedFloat<f32>,
     pub avg_ctransfer: OrderedFloat<f32>,