Added closest parent method on ASTNode (#229)

Co-authored-by: Alex Roan <alex.roan@hotmail.com>
Cyfrin · Mar 18, 2024 · 45c67da · 45c67da
1 parent 6276a46
commit 45c67da
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 16 deletions.
diff --git a/aderyn_core/src/context/browser/closest_parent.rs b/aderyn_core/src/context/browser/closest_parent.rs
@@ -12,6 +12,16 @@ pub trait GetClosestParentOfTypeX {
     ) -> Option<&'a ASTNode>;
 }
 
+impl GetClosestParentOfTypeX for ASTNode {
+    fn closest_parent_of_type<'a>(
+        &self,
+        context: &'a WorkspaceContext,
+        node_type: NodeType,
+    ) -> Option<&'a ASTNode> {
+        context.get_closest_parent(self.id()?, node_type)
+    }
+}
+
 impl GetClosestParentOfTypeX for Assignment {
     fn closest_parent_of_type<'a>(
         &self,

diff --git a/aderyn_core/src/context/browser/parent.rs b/aderyn_core/src/context/browser/parent.rs
@@ -8,6 +8,12 @@ pub trait GetImmediateParent {
     fn parent<'a>(&self, context: &'a WorkspaceContext) -> Option<&'a ASTNode>;
 }
 
+impl GetImmediateParent for ASTNode {
+    fn parent<'a>(&self, context: &'a WorkspaceContext) -> Option<&'a ASTNode> {
+        context.get_parent(self.id()?)
+    }
+}
+
 impl GetImmediateParent for Assignment {
     fn parent<'a>(&self, context: &'a WorkspaceContext) -> Option<&'a ASTNode> {
         context.get_parent(self.id)

diff --git a/aderyn_core/src/context/browser/parent_chain.rs b/aderyn_core/src/context/browser/parent_chain.rs
@@ -8,6 +8,12 @@ pub trait GetParentChain {
     fn parent_chain<'a>(&self, context: &'a WorkspaceContext) -> Option<Vec<&'a ASTNode>>;
 }
 
+impl GetParentChain for ASTNode {
+    fn parent_chain<'a>(&self, context: &'a WorkspaceContext) -> Option<Vec<&'a ASTNode>> {
+        Some(context.get_parent_chain(self.id()?))
+    }
+}
+
 impl GetParentChain for Assignment {
     fn parent_chain<'a>(&self, context: &'a WorkspaceContext) -> Option<Vec<&'a ASTNode>> {
         Some(context.get_parent_chain(self.id))

diff --git a/aderyn_core/src/context/workspace_context.rs b/aderyn_core/src/context/workspace_context.rs
@@ -3,6 +3,7 @@ use crate::visitor::ast_visitor::*;
 use eyre::Result;
 use std::collections::HashMap;
 
+use super::browser::GetImmediateParent;
 use super::capturable::Capturable;
 
 #[derive(Debug, Clone, PartialEq)]
@@ -61,20 +62,6 @@ pub enum ASTNode {
 }
 
 impl ASTNode {
-    pub fn parent<'a>(&self, context: &'a WorkspaceContext) -> Option<&'a ASTNode> {
-        if let Some(id) = self.id() {
-            return context.get_parent(id);
-        }
-        None
-    }
-
-    pub fn parent_chain<'a>(&self, context: &'a WorkspaceContext) -> Option<Vec<&'a ASTNode>> {
-        if let Some(id) = self.id() {
-            return Some(context.get_parent_chain(id));
-        }
-        None
-    }
-
     pub fn node_type(&self) -> NodeType {
         match self {
             ASTNode::ArrayTypeName(_) => NodeType::ArrayTypeName,

diff --git a/aderyn_core/src/detect/experimental/closest_parent.rs b/aderyn_core/src/detect/experimental/closest_parent.rs
@@ -34,7 +34,7 @@ impl IssueDetector for ClosestParentDemonstrator {
                 capture!(self, context, block);
             }
 
-            if let Some(ASTNode::ForStatement(for_statement)) =
+            if let Some(for_statement) =
                 assignment.closest_parent_of_type(context, NodeType::ForStatement)
             {
                 capture!(self, context, for_statement);

diff --git a/aderyn_core/src/detect/experimental/immediate_parent.rs b/aderyn_core/src/detect/experimental/immediate_parent.rs
@@ -32,7 +32,7 @@ impl IssueDetector for ImmediateParentDemonstrator {
                 if let ASTNode::Block(block) = first_parent {
                     println!("1 {}", block);
                     capture!(self, context, first_parent);
-                    if let Some(second_parent) = block.parent(context) {
+                    if let Some(second_parent) = first_parent.parent(context) {
                         if let ASTNode::ForStatement(for_statement) = second_parent {
                             println!("2 {}", for_statement);
                             capture!(self, context, second_parent);