Add frontend logic for rust projects

arthurscchan · arthurscchan · commit bb9408bd5d05 · 2024-11-12T19:22:35.000Z
Signed-off-by: Arthur Chan &lt;arthur.chan@adalogics.com&gt;
diff --git a/frontends/rust/process.py b/frontends/rust/process.py
@@ -0,0 +1,154 @@
+# Copyright 2024 Fuzz Introspector Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import subprocess
+import json
+import yaml
+import os
+import sys
+import re
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+def get_rs_files(project_dir: str) -> List[str]:
+    """Recursively find all Rust files in the project directory."""
+    rs_files = []
+    for root, _, files in os.walk(project_dir):
+        for file in files:
+            if file.endswith(".rs"):
+                rs_files.append(os.path.join(root, file))
+    return rs_files
+
+def extract_function_line_info_from_file(file_path: str) -> Dict[Tuple[str, str], Tuple[int, int]]:
+    """Extract function names with their start and end line numbers from a Rust file."""
+    info = {}
+    pattern = re.compile(r"fn\s+(\w+)\s*\(")
+
+    with open(file_path, "r") as f:
+        lines = f.readlines()
+
+    curr_func = None
+    start = 0
+
+    for i, line in enumerate(lines, 1):
+        match = pattern.search(line)
+        if match:
+            if curr_func:
+                info[(curr_func, file_path)] = (start, i - 1)
+
+            curr_func = match.group(1)
+            start = i
+
+    if curr_func:
+        info[(curr_func, file_path)] = (start, len(lines))
+
+    return info
+
+def analyze_project_functions(project_dir: str) -> Dict[Tuple[str, str], Tuple[int, int]]:
+    """Analyze all functions in the Rust project and map their line numbers."""
+    all_functions = {}
+    rs_files = get_rs_files(project_dir)
+
+    for rs_file in rs_files:
+        functions = extract_function_line_info_from_file(rs_file)
+        all_functions.update(functions)
+
+    return all_functions
+
+def run_rust_analysis(target_directory: str) -> List[Dict]:
+    """Run the Rust analysis tool and retrieve JSON results."""
+    try:
+        result = subprocess.run(
+            ["cargo", "run", "--", target_directory],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            cwd="rust_function_analyser",
+            text=True,
+            check=True
+        )
+        return json.loads(result.stdout)
+    except subprocess.CalledProcessError:
+        return []
+    except ValueError:
+        return []
+
+def add_line_data(rust_results: List[Dict]):
+    """Add line data to functions from rust analysis result."""
+    line_info = analyze_project_functions(target_dir)
+
+    for func in rust_results:
+        func_key = (func["name"], func["file"])
+        if func_key in line_info:
+            func["start_line"], func["end_line"] = line_info[func_key]
+        else:
+            func["start_line"], func["end_line"] = 0, 0
+        func["called_functions"] = [f.replace(" ", "") for f in func["called_functions"]]
+
+def create_yaml_output(data: List[Dict], output_file="data.yaml"):
+    """Generate a YAML file with the analysis results."""
+    yaml_data = {
+        "Fuzzer filename": "",
+        "All functions": {
+            "Function list name": "All functions",
+            "Elements": []
+        }
+    }
+
+    for func in data:
+        yaml_data["All functions"]["Elements"].append({
+            "functionName": func["name"],
+            "functionSourceFile": func["file"],
+            "linkageType": "",
+            "functionLinenumber": func["start_line"],
+            "functionLinenumberEnd": func["end_line"],
+            "functionDepth": func["depth"],
+            "returnType": func["return_type"],
+            "argCount": func["arg_count"],
+            "argTypes": func["arg_types"],
+            "constantsTouched": [],
+            "argNames": [],
+            "BBCount": 0,
+            "ICount": 0,
+            "EdgeCount": 0,
+            "CyclomaticComplexity": func["complexity"],
+            "functionsReached": func["called_functions"],
+            "functionUses": 0,
+            "BranchProfiles": [],
+            "Callsites": []
+        })
+
+    with open(output_file, "w") as file:
+        yaml.dump(yaml_data, file, default_flow_style=False)
+
+    print(f"YAML output saved to {output_file}")
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python3 script.py <target_directory>")
+        sys.exit(1)
+
+    target_dir = sys.argv[1]
+    if not Path(target_dir).is_dir():
+        print(f"Error: {target_dir} is not a valid directory")
+        sys.exit(1)
+
+    # Run the rust analysis frontend code
+    rust_analysis_results = run_rust_analysis(target_dir)
+
+    # Manually extract the line info for each function.
+    # This is needed because the rust analysis syn AST approach
+    # cannot retrieve line number info on stable rust and non-nightly build
+    add_line_data(rust_analysis_results)
+
+    create_yaml_output(rust_analysis_results)
diff --git a/frontends/rust/rust_function_analyser/.gitignore b/frontends/rust/rust_function_analyser/.gitignore
@@ -0,0 +1 @@
+/target
diff --git a/frontends/rust/rust_function_analyser/Cargo.toml b/frontends/rust/rust_function_analyser/Cargo.toml
@@ -0,0 +1,11 @@
+[package]
+name = "rust_function_analyser"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+syn = { version = "2.0", features = ["full", "visit"] }
+quote = "1.0"
+walkdir = "2.4"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0" 
diff --git a/frontends/rust/rust_function_analyser/src/analyse.rs b/frontends/rust/rust_function_analyser/src/analyse.rs
@@ -0,0 +1,165 @@
+/* Copyright 2024 Fuzz Introspector Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+use syn::{ItemFn, Stmt, FnArg, ReturnType};
+use std::collections::{HashMap, HashSet};
+use serde::{Serialize, Deserialize};
+use std::fs;
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct FunctionInfo {
+    pub name: String,
+    pub file: String,
+    pub return_type: String,
+    pub arg_count: usize,
+    pub arg_types: Vec<String>,
+    pub complexity: usize,
+    pub called_functions: Vec<String>,
+    pub depth: usize,
+}
+
+pub struct FunctionAnalyser {
+    pub functions: Vec<FunctionInfo>,
+    pub call_stack: HashMap<String, HashSet<String>>,
+}
+
+impl FunctionAnalyser {
+    pub fn new() -> Self {
+        Self {
+            functions: Vec::new(),
+            call_stack: HashMap::new(),
+        }
+    }
+
+    pub fn visit_function(&mut self, node: &ItemFn, file_name: &str) {
+        let function_name = node.sig.ident.to_string();
+        let return_type = match &node.sig.output {
+            ReturnType::Default => "void".to_string(),
+            ReturnType::Type(_, ty) => format!("{}", quote::ToTokens::to_token_stream(&**ty)),
+        }
+        .replace(' ', "");
+
+        let arg_types = node
+            .sig
+            .inputs
+            .iter()
+            .filter_map(|arg| {
+                if let FnArg::Typed(pat) = arg {
+                    Some(format!("{}", quote::ToTokens::to_token_stream(&*pat.ty)).replace(' ', ""))
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>();
+
+        let complexity = calculate_cyclomatic_complexity(node);
+
+        self.functions.push(FunctionInfo {
+            name: function_name.clone(),
+            file: file_name.to_string(),
+            return_type,
+            arg_count: arg_types.len(),
+            arg_types,
+            complexity,
+            called_functions: vec![],
+            depth: 0,
+        });
+
+        self.call_stack
+            .entry(function_name.clone())
+            .or_default();
+    }
+
+    pub fn calculate_depths(&mut self) {
+        let mut depth_map: HashMap<String, usize> = HashMap::new();
+
+        for function in &self.functions {
+            let depth = self.calculate_function_depth(&function.name);
+            depth_map.insert(function.name.clone(), depth);
+        }
+
+        for function in self.functions.iter_mut() {
+            if let Some(&depth) = depth_map.get(&function.name) {
+                function.depth = depth;
+            }
+        }
+    }
+
+    fn calculate_function_depth(
+        &self,
+        function_name: &str,
+    ) -> usize {
+        let mut max_depth = 0;
+        let mut stack = vec![(function_name, 0)];
+
+        while let Some((func, depth)) = stack.pop() {
+            if depth > max_depth {
+                max_depth = depth;
+            }
+
+            if let Some(called) = self.call_stack.get(func) {
+                for callee in called {
+                    if callee != func {
+                        stack.push((callee, depth + 1));
+                    }
+                }
+            }
+        }
+
+        max_depth
+    }
+}
+
+fn calculate_cyclomatic_complexity(node: &ItemFn) -> usize {
+    let mut complexity = 1;
+
+    for stmt in &node.block.stmts {
+        if matches!(stmt, Stmt::Expr(..)) {
+            complexity += 1;
+        }
+    }
+
+    complexity
+}
+
+pub fn analyse_directory(dir: &str, exclude_dirs: &[&str]) -> std::io::Result<String> {
+    let mut analyser = FunctionAnalyser::new();
+
+    for entry in fs::read_dir(dir)? {
+        let entry = entry?;
+        let path = entry.path();
+
+        if path.is_dir() && exclude_dirs.iter().any(|d| path.ends_with(d)) {
+            continue;
+        } else if path.is_dir() {
+            let sub_result = analyse_directory(path.to_str().unwrap(), exclude_dirs)?;
+            let parsed_functions: Vec<FunctionInfo> = serde_json::from_str(&sub_result).unwrap();
+            analyser.functions.extend(parsed_functions);
+        } else if path.extension().and_then(|s| s.to_str()) == Some("rs") {
+            let file_content = fs::read_to_string(&path)?;
+            let syntax = syn::parse_file(&file_content)
+                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
+
+            for item in syntax.items {
+                if let syn::Item::Fn(func) = item {
+                    analyser.visit_function(&func, path.to_str().unwrap());
+                }
+            }
+        }
+    }
+
+    analyser.calculate_depths();
+    Ok(serde_json::to_string(&analyser.functions).unwrap())
+}
diff --git a/frontends/rust/rust_function_analyser/src/main.rs b/frontends/rust/rust_function_analyser/src/main.rs
@@ -0,0 +1,45 @@
+/* Copyright 2024 Fuzz Introspector Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+mod analyse;
+
+use std::io;
+
+fn main() -> io::Result<()> {
+    let exclude_dirs = vec![
+        "target",
+        "node_modules",
+        "aflplusplus",
+        "tests",
+        "examples",
+        "benches",
+        "honggfuzz",
+        "inspector",
+        "libfuzzer",
+    ];
+
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() != 2 {
+        eprintln!("Usage: cargo run -- <source_directory>");
+        std::process::exit(1);
+    }
+    let target_directory = &args[1];
+
+    // Collect all results into a single string and print to stdout
+    let result = analyse::analyse_directory(target_directory, &exclude_dirs)?;
+    println!("{}", result);
+
+    Ok(())
+}
