wufjfwk

Cyfrin · Sep 30, 2024 · 866ee2e · 866ee2e
1 parent 5daabfa
commit 866ee2e
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 21 deletions.
diff --git a/aderyn_core/src/context/flow/display.rs b/aderyn_core/src/context/flow/display.rs
@@ -43,6 +43,12 @@ impl CfgStartNode {
             CfgStartNode::StartUncheckedBlock(ast_id) => {
                 format!("START UNCHECKED BLOCK ({})", ast_id)
             }
+            CfgStartNode::StartModifierBody(ast_id) => {
+                format!("START MODIFIER BODY {}", ast_id)
+            }
+            CfgStartNode::StartFunctionBody(ast_id) => {
+                format!("START FUNCTION BODY {}", ast_id)
+            }
             CfgStartNode::StartIf(ast_id) => format!("START IF ({})", ast_id),
             CfgStartNode::StartIfCond => String::from("START IF COND"),
             CfgStartNode::StartIfTrue => String::from("START IF TRUE BRANCH"),
@@ -68,6 +74,8 @@ impl CfgEndNode {
             CfgEndNode::End => String::from("END"),
             CfgEndNode::EndBlock(ast_id) => format!("END BLOCK ({})", ast_id),
             CfgEndNode::EndUncheckedBlock(ast_id) => format!("END UNCHECKED BLOCK ({})", ast_id),
+            CfgEndNode::EndModifierBody(ast_id) => format!("END MODIFIER BODY ({})", ast_id),
+            CfgEndNode::EndFunctionBody(ast_id) => format!("END FUNCTION BODY ({})", ast_id),
             CfgEndNode::EndIf(ast_id) => format!("END IF ({})", ast_id),
             CfgEndNode::EndIfCond => String::from("END IF COND"),
             CfgEndNode::EndIfTrue => String::from("END IF TRUE BRANCH"),

diff --git a/aderyn_core/src/context/flow/mod.rs b/aderyn_core/src/context/flow/mod.rs
@@ -264,28 +264,27 @@ impl Cfg {
     }
 }
 
-impl Cfg {
-    pub fn accept_block(&mut self, context: &WorkspaceContext, block: &Block) {
-        let start = self.add_start_node();
-        let end = self.add_end_node();
-        let block = self.add_block_node(block);
-
-        self.add_flow_edge(start, block);
-        self.add_flow_edge(block, end);
-
-        while let Some(reduction_candidate) = self.reduction_queue.pop_front() {
-            self.reduce(context, reduction_candidate);
-        }
-    }
-}
-
 #[cfg(test)]
 mod control_flow_tests {
     use super::*;
     use crate::context::flow::visualizer::control_flow_tests::output_graph;
     use crate::detect::test_utils::load_solidity_source_unit;
     use serial_test::serial;
 
+    impl Cfg {
+        pub fn accept_block(&mut self, context: &WorkspaceContext, block: &Block) {
+            let start = self.add_start_node();
+            let end = self.add_end_node();
+            let block = self.add_block_node(block);
+
+            self.add_flow_edge(start, block);
+            self.add_flow_edge(block, end);
+
+            while let Some(reduction_candidate) = self.reduction_queue.pop_front() {
+                self.reduce(context, reduction_candidate);
+            }
+        }
+    }
     #[test]
     #[serial]
     fn simple_program_function1() {

diff --git a/aderyn_core/src/context/flow/voids.rs b/aderyn_core/src/context/flow/voids.rs
@@ -4,28 +4,32 @@ use super::{AstNodeId, Cfg, CfgNodeDescriptor, CfgNodeId};
 #[derive(Debug, Clone)]
 pub enum CfgStartNode {
     Start,
-    StartBlock(AstNodeId),
-    StartUncheckedBlock(AstNodeId),
-    StartIf(AstNodeId),
+    StartFunctionBody(AstNodeId),   // Function Definition ID
+    StartModifierBody(AstNodeId),   // Modifier Definition ID
+    StartBlock(AstNodeId),          // Block Node ID
+    StartUncheckedBlock(AstNodeId), // Unchecked Block ID
+    StartIf(AstNodeId),             // If Statemtnt ID
     StartIfCond,
     StartIfTrue,
     StartIfFalse,
-    StartWhile(AstNodeId),
+    StartWhile(AstNodeId), // While Statement ID
     StartWhileCond,
     StartWhileBody,
-    StartFor(AstNodeId),
+    StartFor(AstNodeId), // For Statement ID
     StartForInitExp,
     StartForCond,
     StartForLoopExp,
     StartForBody,
-    StartDoWhile(AstNodeId),
+    StartDoWhile(AstNodeId), // Do While Statement ID
     StartDoWhileCond,
     StartDoWhileBody,
 }
 
 #[derive(Debug, Clone)]
 pub enum CfgEndNode {
     End,
+    EndFunctionBody(AstNodeId),
+    EndModifierBody(AstNodeId),
     EndBlock(AstNodeId),
     EndUncheckedBlock(AstNodeId),
     EndIf(AstNodeId),
@@ -199,4 +203,24 @@ impl Cfg {
             CfgEndNode::EndUncheckedBlock(unchecked_block),
         )))
     }
+    pub fn add_start_function_body_node(&mut self, function_definition_id: AstNodeId) -> CfgNodeId {
+        self.add_node(CfgNodeDescriptor::Start(Box::new(
+            CfgStartNode::StartFunctionBody(function_definition_id),
+        )))
+    }
+    pub fn add_end_function_body_node(&mut self, function_definition_id: AstNodeId) -> CfgNodeId {
+        self.add_node(CfgNodeDescriptor::End(Box::new(
+            CfgEndNode::EndFunctionBody(function_definition_id),
+        )))
+    }
+    pub fn add_start_modifier_body_node(&mut self, modifier_definition_id: AstNodeId) -> CfgNodeId {
+        self.add_node(CfgNodeDescriptor::Start(Box::new(
+            CfgStartNode::StartModifierBody(modifier_definition_id),
+        )))
+    }
+    pub fn add_end_modifier_body_node(&mut self, modifier_definition_id: AstNodeId) -> CfgNodeId {
+        self.add_node(CfgNodeDescriptor::End(Box::new(
+            CfgEndNode::EndModifierBody(modifier_definition_id),
+        )))
+    }
 }