From cfd38899e88824892fe85a4d26b520bf30a7a332 Mon Sep 17 00:00:00 2001
From: Alex Roan <alex@cyfrin.io>
Date: Thu, 6 Mar 2025 10:54:12 +0000
Subject: [PATCH] Entry points Audit Mode

---
 aderyn_core/src/audit/auditor.rs      |   3 +-
 aderyn_core/src/audit/entry_points.rs | 101 ++++++++++++++++++++++++++
 aderyn_core/src/audit/mod.rs          |   1 +
 3 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 aderyn_core/src/audit/entry_points.rs

diff --git a/aderyn_core/src/audit/auditor.rs b/aderyn_core/src/audit/auditor.rs
index 473fcf6cd..c073c77ea 100644
--- a/aderyn_core/src/audit/auditor.rs
+++ b/aderyn_core/src/audit/auditor.rs
@@ -4,7 +4,7 @@ use prettytable::{format, Row, Table};
 
 use crate::{
     audit::{
-        attack_surface::AttackSurfaceDetector,
+        attack_surface::AttackSurfaceDetector, entry_points::EntryPointsDetector,
         public_functions_no_sender::PublicFunctionsNoSenderChecksDetector,
     },
     context::workspace_context::WorkspaceContext,
@@ -14,6 +14,7 @@ pub fn get_auditor_detectors() -> Vec<Box<dyn AuditorDetector>> {
     vec![
         Box::<AttackSurfaceDetector>::default(),
         Box::<PublicFunctionsNoSenderChecksDetector>::default(),
+        Box::<EntryPointsDetector>::default(),
     ]
 }
 
diff --git a/aderyn_core/src/audit/entry_points.rs b/aderyn_core/src/audit/entry_points.rs
new file mode 100644
index 000000000..317dcb8a7
--- /dev/null
+++ b/aderyn_core/src/audit/entry_points.rs
@@ -0,0 +1,101 @@
+use prettytable::{row, Row};
+
+use super::auditor::AuditorDetector;
+use crate::{
+    ast::{FunctionKind, NodeType},
+    context::workspace_context::{ASTNode, WorkspaceContext},
+    detect::helpers::get_implemented_external_and_public_functions,
+};
+use std::{cmp::Ordering, collections::BTreeSet, error::Error};
+
+#[derive(Clone, Eq, PartialEq)]
+pub struct EntryPointsInstance {
+    pub contract_name: String,
+    pub function_name: String,
+    pub function_kind: FunctionKind,
+}
+
+impl Ord for EntryPointsInstance {
+    fn cmp(&self, other: &Self) -> Ordering {
+        let by_contract = self.contract_name.cmp(&other.contract_name);
+        if by_contract == Ordering::Equal {
+            self.function_name.cmp(&other.function_name)
+        } else {
+            by_contract
+        }
+    }
+}
+
+impl PartialOrd for EntryPointsInstance {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+#[derive(Default)]
+pub struct EntryPointsDetector {
+    // contract_name, function_name
+    found_instances: BTreeSet<EntryPointsInstance>,
+}
+
+impl AuditorDetector for EntryPointsDetector {
+    fn detect(&mut self, context: &WorkspaceContext) -> Result<bool, Box<dyn Error>> {
+        let functions = get_implemented_external_and_public_functions(context);
+        functions.for_each(|function_definition| {
+            if let ASTNode::ContractDefinition(contract_definition) = context
+                .get_closest_ancestor(function_definition.id, NodeType::ContractDefinition)
+                .unwrap()
+            {
+                let contract_name = contract_definition.name.clone();
+                self.found_instances.insert(EntryPointsInstance {
+                    contract_name,
+                    function_name: function_definition.name.clone(),
+                    function_kind: function_definition.kind().clone(),
+                });
+            }
+        });
+        Ok(!self.found_instances.is_empty())
+    }
+
+    fn title(&self) -> String {
+        String::from("Contract Entry Points")
+    }
+
+    fn skeletal_clone(&self) -> Box<dyn AuditorDetector> {
+        Box::<EntryPointsDetector>::default()
+    }
+
+    fn table_titles(&self) -> Row {
+        row!["Contract", "Function Kind", "Function Name"]
+    }
+
+    fn table_rows(&self) -> Vec<Row> {
+        self.found_instances
+            .iter()
+            .map(|instance| {
+                row![instance.contract_name, instance.function_kind, instance.function_name,]
+            })
+            .collect()
+    }
+}
+
+#[cfg(test)]
+mod public_functions_no_sender_checks {
+    use crate::{
+        audit::{auditor::AuditorDetector, entry_points::EntryPointsDetector},
+        detect::test_utils::load_solidity_source_unit,
+    };
+
+    #[test]
+    fn test_public_functions_no_sender_checks() {
+        let context = load_solidity_source_unit(
+            "../tests/contract-playground/src/auditor_mode/PublicFunctionsWithoutSenderCheck.sol",
+        );
+
+        let mut detector = EntryPointsDetector::default();
+        let found = detector.detect(&context).unwrap();
+        // assert that the detector found an issue
+        assert!(found);
+        assert!(detector.found_instances.len() == 11);
+    }
+}
diff --git a/aderyn_core/src/audit/mod.rs b/aderyn_core/src/audit/mod.rs
index 57a75cfdd..f4c9b94d1 100644
--- a/aderyn_core/src/audit/mod.rs
+++ b/aderyn_core/src/audit/mod.rs
@@ -1,3 +1,4 @@
 pub mod attack_surface;
 pub mod auditor;
+pub mod entry_points;
 pub mod public_functions_no_sender;