Merge pull request #6 from FuzzingLabs/feat/cfg

Implement Control-Flow Graph

README.md

@@ -30,6 +30,19 @@ For a colourless output :
 cargo run --bin sierra-decompiler <sierra file> --no-color
 ```
 
+#### Print the contract's Control-Flow Graph
+
+```
+cargo run ./examples/sierra/fib_array.sierra --cfg  
+
+# Output the Control-Flow Graph to a custom folder (default is ./output_cfg)
+cargo run ./tests/sierra_files/fib_array.sierra --cfg --cfg-output ./test 
+```
+
+<p align="center">
+	<img src="/doc/images/cfg-output.png" height="400px"/>
+</p>
+
 #### Use it as a library 
 
 It is also possible to use the `sierra-analyzer-lib` library to decompile serialised or unserialised Sierra files.
@@ -39,6 +52,6 @@ Examples can be found [here](/lib/examples/).
 #### Features
 
 - [x] Decompiler
-- [ ] Control-Flow Graph
+- [x] Control-Flow Graph
 - [ ] Call Graph
 - [ ] Security detectors

lib/Cargo.toml

@@ -8,6 +8,7 @@ edition = "2021"
 [dependencies]
 cairo-lang-sierra = "~2.6.3"
 colored = "2.1.0"
+graphviz-rust = "0.9.0"
 
 [dev-dependencies]
 serde_json = "1.0.116"

lib/examples/generate_cfg.rs

@@ -0,0 +1,16 @@
+use sierra_analyzer_lib::sierra_program::SierraProgram;
+
+fn main() {
+    let content = include_str!("../../examples/sierra/fib.sierra").to_string();
+
+    // Init a new SierraProgram with the .sierra file content
+    let program = SierraProgram::new(content);
+
+    // Decompile the Sierra programs
+    let mut decompiler = program.decompiler();
+    decompiler.decompile(false);
+
+    // Generate & print the dot graph
+    let cfg = decompiler.generate_cfg();
+    println!("{}", cfg)
+}

lib/src/config.rs

@@ -0,0 +1,32 @@
+pub struct GraphConfig;
+
+#[allow(dead_code)]
+impl GraphConfig {
+    // Node attributes for CFG
+    pub const CFG_NODE_ATTR_STYLE: &'static str = "filled, solid";
+    pub const CFG_NODE_ATTR_SHAPE: &'static str = "rect, plaintext";
+    pub const CFG_NODE_ATTR_COLOR: &'static str = "#9E9E9E";
+    pub const CFG_NODE_ATTR_FILLCOLOR: &'static str = "#F5F5F5";
+    pub const CFG_NODE_ATTR_FONTNAME: &'static str = "Helvetica,Arial,sans-serif";
+    pub const CFG_NODE_ATTR_MARGIN: &'static str = "0.2";
+
+    // Graph attributes for CFG
+    pub const CFG_GRAPH_ATTR_OVERLAP: &'static str = "scale";
+    pub const CFG_GRAPH_ATTR_FONTNAME: &'static str = "Helvetica,Arial,sans-serif";
+    pub const CFG_GRAPH_ATTR_FONTSIZE: &'static str = "20";
+    pub const CFG_GRAPH_ATTR_LAYOUT: &'static str = "dot";
+    pub const CFG_GRAPH_ATTR_NEWRANK: &'static str = "true";
+
+    // Edge attributes for CFG
+    pub const CFG_EDGE_ATTR_ARROWSIZE: &'static str = "0.5";
+    pub const CFG_EDGE_ATTR_FONTNAME: &'static str = "Helvetica,Arial,sans-serif";
+    pub const CFG_EDGE_ATTR_LABELDISTANCE: &'static str = "3";
+    pub const CFG_EDGE_ATTR_LABELFONTCOLOR: &'static str = "#00000080";
+    pub const CFG_EDGE_ATTR_PENWIDTH: &'static str = "2";
+
+    // Edge colors
+    pub const EDGE_CONDITIONAL_TRUE_COLOR: &'static str = "#8BC34A";
+    pub const EDGE_CONDITIONAL_FALSE_COLOR: &'static str = "#C62828";
+    pub const EDGE_UNCONDITIONAL_COLOR: &'static str = "#0D47A1";
+    pub const EDGE_FALLTHROUGH_COLOR: &'static str = "#212121";
+}

lib/src/decompiler/cfg.rs

@@ -4,6 +4,7 @@ use std::collections::HashSet;
 use cairo_lang_sierra::program::BranchTarget;
 use cairo_lang_sierra::program::GenStatement;
 
+use crate::config::GraphConfig;
 use crate::decompiler::function::SierraStatement;
 
 /// A struct representing a control flow graph (CFG) for a function
@@ -14,6 +15,8 @@ use crate::decompiler::function::SierraStatement;
 /// - Edges denote control flow transfers (conditional branches, jumps, fallthroughs)
 #[derive(Debug, Clone)]
 pub struct ControlFlowGraph {
+    /// Function name or ID
+    function_name: String,
     /// List of statements in the function
     statements: Vec<SierraStatement>,
     /// List of basic blocks in the CFG
@@ -22,8 +25,9 @@ pub struct ControlFlowGraph {
 
 impl<'a> ControlFlowGraph {
     /// Creates a new `ControlFlowGraph` instance
-    pub fn new(statements: Vec<SierraStatement>) -> Self {
+    pub fn new(function_name: String, statements: Vec<SierraStatement>) -> Self {
         Self {
+            function_name,
             statements,
             basic_blocks: Vec::new(),
         }
@@ -181,6 +185,76 @@ impl<'a> ControlFlowGraph {
         }
         parents
     }
+
+    /// Generates the DOT format subgraph for function CFG
+    pub fn generate_dot_graph(&self) -> String {
+        let mut dot_graph = format!("\tsubgraph \"cluster_{}\" {{\n", self.function_name);
+        dot_graph += &format!("\t\tlabel=\"{}\"\n", self.function_name);
+        dot_graph += &format!(
+            "\t\tfontname=\"{}\";\n",
+            GraphConfig::CFG_GRAPH_ATTR_FONTNAME
+        );
+        dot_graph += &format!("\t\tfontsize={};\n", GraphConfig::CFG_GRAPH_ATTR_FONTSIZE);
+
+        // Iterate over each basic block to create nodes
+        for block in &self.basic_blocks {
+            let mut label_instruction = String::new();
+            for statement in &block.statements {
+                label_instruction +=
+                    &format!("{} : {}\t\t\\l", statement.offset, statement.raw_statement());
+            }
+            dot_graph += &format!(
+                "\t\t\"{}\" [label=\"{}\" shape=\"box\" style=\"{}\" fillcolor=\"{}\" color=\"{}\" fontname=\"{}\" margin=\"{}\"];\n",
+                block.name,
+                label_instruction,
+                GraphConfig::CFG_NODE_ATTR_STYLE,
+                GraphConfig::CFG_NODE_ATTR_FILLCOLOR,
+                GraphConfig::CFG_NODE_ATTR_COLOR,
+                GraphConfig::CFG_NODE_ATTR_FONTNAME,
+                GraphConfig::CFG_NODE_ATTR_MARGIN
+            );
+        }
+
+        // Add edges between nodes
+        for block in &self.basic_blocks {
+            for edge in &block.edges {
+                let color = match edge.edge_type {
+                    EdgeType::ConditionalTrue => GraphConfig::EDGE_CONDITIONAL_TRUE_COLOR,
+                    EdgeType::ConditionalFalse => GraphConfig::EDGE_CONDITIONAL_FALSE_COLOR,
+                    EdgeType::Unconditional => GraphConfig::EDGE_UNCONDITIONAL_COLOR,
+                    EdgeType::Fallthrough => GraphConfig::EDGE_FALLTHROUGH_COLOR,
+                };
+                if self
+                    .basic_blocks
+                    .iter()
+                    .any(|b| b.start_offset == edge.destination)
+                {
+                    dot_graph += &format!(
+                        "\t\t\"{}\" -> \"{}\" [color=\"{}\" arrowsize={} fontname=\"{}\" labeldistance={} labelfontcolor=\"{}\" penwidth={}];\n",
+                        block.name,
+                        self.get_block_name_by_offset(edge.destination),
+                        color,
+                        GraphConfig::CFG_EDGE_ATTR_ARROWSIZE,
+                        GraphConfig::CFG_EDGE_ATTR_FONTNAME,
+                        GraphConfig::CFG_EDGE_ATTR_LABELDISTANCE,
+                        GraphConfig::CFG_EDGE_ATTR_LABELFONTCOLOR,
+                        GraphConfig::CFG_EDGE_ATTR_PENWIDTH
+                    );
+                }
+            }
+        }
+
+        dot_graph += "\t}\n";
+        dot_graph
+    }
+
+    /// Retrieves the name of a basic block based on its start offset
+    fn get_block_name_by_offset(&self, offset: u32) -> String {
+        self.basic_blocks
+            .iter()
+            .find(|&b| b.start_offset == offset)
+            .map_or(String::from("Unknown"), |b| b.name.clone())
+    }
 }
 
 /// Enum representing different types of CFG edges

lib/src/decompiler/decompiler.rs

@@ -6,6 +6,7 @@ use cairo_lang_sierra::program::LibfuncDeclaration;
 use cairo_lang_sierra::program::StatementIdx;
 use cairo_lang_sierra::program::TypeDeclaration;
 
+use crate::config::GraphConfig;
 use crate::decompiler::cfg::BasicBlock;
 use crate::decompiler::cfg::EdgeType;
 use crate::decompiler::function::Function;
@@ -471,4 +472,51 @@ impl<'a> Decompiler<'a> {
 
         decompiled_basic_block
     }
+
+    /// Generates a control flow graph representation (CFG) in DOT format
+    pub fn generate_cfg(&mut self) -> String {
+        let mut dot = String::from("digraph {\n");
+
+        // Global graph configuration
+        dot.push_str(&format!(
+            "\tgraph [fontname=\"{}\" fontsize={} layout={} newrank={} overlap={}];\n",
+            GraphConfig::CFG_GRAPH_ATTR_FONTNAME,
+            GraphConfig::CFG_GRAPH_ATTR_FONTSIZE,
+            GraphConfig::CFG_GRAPH_ATTR_LAYOUT,
+            GraphConfig::CFG_GRAPH_ATTR_NEWRANK,
+            GraphConfig::CFG_GRAPH_ATTR_OVERLAP,
+        ));
+        // Global node configuration
+        dot.push_str(&format!("\tnode [color=\"{}\" fillcolor=\"{}\" fontname=\"{}\" margin={} shape=\"{}\" style=\"{}\"];\n",
+            GraphConfig::CFG_NODE_ATTR_COLOR,
+            GraphConfig::CFG_NODE_ATTR_FILLCOLOR,
+            GraphConfig::CFG_NODE_ATTR_FONTNAME,
+            GraphConfig::CFG_NODE_ATTR_MARGIN,
+            GraphConfig::CFG_NODE_ATTR_SHAPE,
+            GraphConfig::CFG_NODE_ATTR_STYLE,
+        ));
+        // Global edge configuration
+        dot.push_str(&format!("\tedge [arrowsize={} fontname=\"{}\" labeldistance={} labelfontcolor=\"{}\" penwidth={}];\n",
+            GraphConfig::CFG_EDGE_ATTR_ARROWSIZE,
+            GraphConfig::CFG_EDGE_ATTR_FONTNAME,
+            GraphConfig::CFG_EDGE_ATTR_LABELDISTANCE,
+            GraphConfig::CFG_EDGE_ATTR_LABELFONTCOLOR,
+            GraphConfig::CFG_EDGE_ATTR_PENWIDTH,
+        ));
+
+        // Add a CFG representation for each function
+        for function in &mut self.functions {
+            function.create_cfg();
+            if let Some(cfg) = &function.cfg {
+                // Generate function subgraph
+                let subgraph = cfg.generate_dot_graph();
+                dot += &subgraph;
+            }
+        }
+
+        // Add the closing curly braces to the DOT graph representation
+        dot.push_str("}\n");
+
+        dot
+    }
 }

lib/src/decompiler/function.rs

@@ -8,7 +8,7 @@ use cairo_lang_sierra::program::StatementIdx;
 use crate::decompiler::cfg::ControlFlowGraph;
 use crate::decompiler::cfg::SierraConditionalBranch;
 use crate::extract_parameters;
-use crate::parse_libfunc_name;
+use crate::parse_element_name;
 
 /// A struct representing a statement
 #[derive(Debug, Clone)]
@@ -61,7 +61,7 @@ impl SierraStatement {
             // Function calls & variables assignments
             GenStatement::Invocation(invocation) => {
                 // Function name in blue
-                let libfunc_id_str = parse_libfunc_name!(invocation.libfunc_id).blue();
+                let libfunc_id_str = parse_element_name!(invocation.libfunc_id).blue();
 
                 // Function parameters
                 let parameters = extract_parameters!(invocation.args);
@@ -92,6 +92,13 @@ impl SierraStatement {
         }
     }
 
+    /// Return the raw statement, as in the original sierra file
+    /// Used in the CFG
+    #[inline]
+    pub fn raw_statement(&self) -> String {
+        self.statement.to_string()
+    }
+
     /// Returns a reference to this statement as a conditional branch if it is one
     pub fn as_conditional_branch(&self) -> Option<SierraConditionalBranch> {
         if self.is_conditional_branch {
@@ -100,7 +107,7 @@ impl SierraStatement {
                 let statement = self.statement.clone();
 
                 // Function name
-                let libfunc_id_str = parse_libfunc_name!(invocation.libfunc_id);
+                let libfunc_id_str = parse_element_name!(invocation.libfunc_id);
 
                 // Parameters
                 let parameters = extract_parameters!(invocation.args);
@@ -201,7 +208,10 @@ impl<'a> Function<'a> {
     /// Initializes the control flow graph (CFG) for the function
     pub fn create_cfg(&mut self) {
         // Create a new control flow graph instance
-        let mut cfg = ControlFlowGraph::new(self.statements.clone());
+        let mut cfg = ControlFlowGraph::new(
+            parse_element_name!(self.function.id.clone()),
+            self.statements.clone(),
+        );
 
         // Generate the CFG basic blocks
         cfg.generate_basic_blocks();

lib/src/decompiler/macros.rs

@@ -1,7 +1,7 @@
-/// Macro to parse the debug name from a libfunc_id,
-/// using the debug_name if present or falling back to the id field
+/// Macro to parse the debug name from a libfunc or a function ID,
+/// using the debug_name if present or falling back to the ID field
 #[macro_export]
-macro_rules! parse_libfunc_name {
+macro_rules! parse_element_name {
     ($libfunc_id:expr) => {
         if let Some(debug_name) = &$libfunc_id.debug_name {
             debug_name.to_string()
@@ -13,7 +13,7 @@ macro_rules! parse_libfunc_name {
 
 /// Macro to extract parameters from the args field of a GenInvocation object.
 /// It converts each parameter into a String, using the debug_name if available,
-/// otherwise using the id field.
+/// otherwise using the ID field
 #[macro_export]
 macro_rules! extract_parameters {
     ($args:expr) => {

lib/src/graph/graph.rs

@@ -0,0 +1,26 @@
+use graphviz_rust::cmd::Format;
+use graphviz_rust::dot_structures::*;
+use graphviz_rust::exec;
+use graphviz_rust::parse;
+use graphviz_rust::printer::PrinterContext;
+use std::fs::File;
+use std::io::{self, Write};
+
+/// Converts a DOT graph provided as a string to SVG format and saves it to a file
+pub fn save_svg_graph_to_file(filename: &str, graph: String) -> io::Result<()> {
+    // Parse the graph from the string input
+    let parsed_graph: Graph = parse(&graph).unwrap();
+
+    // Generate SVG output
+    let svg_data = exec(
+        parsed_graph,
+        &mut PrinterContext::default(),
+        vec![Format::Svg.into()],
+    )?;
+
+    // Create the output file and write the SVG data to it
+    let mut file = File::create(filename)?;
+    file.write_all(&svg_data)?;
+
+    Ok(())
+}

lib/src/graph/mod.rs

@@ -0,0 +1 @@
+pub mod graph;

lib/src/lib.rs

@@ -1,2 +1,4 @@
+mod config;
 mod decompiler;
+pub mod graph;
 pub mod sierra_program;

lib/tests/test_decompiler.rs

@@ -2,7 +2,7 @@ use sierra_analyzer_lib::sierra_program::SierraProgram;
 
 #[test]
 fn test_decompiler_output() {
-	// Read file content
+    // Read file content
     let content = include_str!("../../examples/sierra/fib.sierra").to_string();
 
     // Init a new SierraProgram with the .sierra file content