✨ Add selector clashes detection to proxy detector

Author: Aragahara

tests/detectors_sources/proxy_contract.sol tests/detectors_sources/proxy_contract_selector_clashes.sol

@@ -7,7 +7,7 @@ contract Slots {
     bytes32 internal constant _OWNER_SLOT = 0x02016836a56b71f0d02689e69e326f4f4c1b9057164ef592671cf0d37c8040c0;
 }
 
-abstract contract Legit_Proxy_1 {
+abstract contract AProxy {
     function _implementation() internal view virtual returns (address);
 
     function _delegate(address implementation) internal virtual {
@@ -37,15 +37,19 @@ abstract contract Legit_Proxy_1 {
     function _beforeFallback() internal virtual {}
 }
 
-contract Bug_Proxy_1 is Legit_Proxy_1, Slots {
+contract Proxy is AProxy, Slots {
     function _implementation() internal view override returns (address implementation_) {
         assembly {
             implementation_ := sload(_IMPLEMENTATION_SLOT)
         }
     }
+
+    function bug_clash(address addr) public {
+        int a = 1;
+    }
 }
 
-abstract contract Legit_Proxy_2 {
+abstract contract AProxy2 {
     function _implementation() internal view virtual returns (address);
 
     function _delegate(address implementation) internal virtual {
@@ -248,17 +252,48 @@ library Address {
     }
 }
 
-contract Bug_Impl_1 is Slots {
+contract Proxy2 is AProxy, Slots {
+    function _implementation() internal view override returns (address implementation_) {
+        assembly {
+            implementation_ := sload(_IMPLEMENTATION_SLOT)
+        }
+    }
+
+    function legit_clash(address addr) public {
+        int a = 1;
+    }
+}
+
+contract Impl1 is Slots {
     function _setImplementation(address newImplementation) private {
         require(Address.isContract(newImplementation), "ERC1967: new implementation is not a contract");
         StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value = newImplementation;
     }
+
+    function bug_clash(address addr) public {
+        _setImplementation(addr);
+    }
 }
 
-contract Legit_Impl_2 is Slots {
+contract Impl2 is Slots {
+    function _setImplementation(address newImplementation) private {
+        require(Address.isContract(newImplementation), "ERC1967: new implementation is not a contract");
+        StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value = newImplementation;
+    }
+
+    function legit_clash(address addr, uint a) public {
+        _setImplementation(addr);
+    }
+}
+
+contract Impl3 is Slots {
     function _setImplementation(address newImplementation) private {
         require(Address.isContract(newImplementation), "ERC1967: new implementation is not a contract");
         StorageSlot.getAddressSlot(_OWNER_SLOT).value = newImplementation;
     }
+
+    function legit_clash(address addr) public {
+        _setImplementation(addr);
+    }
 }
 

tests/test_detectors.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 import shutil
 import tempfile
 from pathlib import Path, PurePath
@@ -104,20 +105,24 @@ def check_detected_in_functions(config, source_path):
     detections = detect(config, source_units)
 
     detections_fns = [_get_function_def(det.result.ir_node) for det in detections]
-    detections_fn_names = [fn.name for fn in detections_fns if fn is not None]
+    detections_fn_names = [fn.canonical_name for fn in detections_fns if fn is not None]
 
+    detected_something = False
     for contract in source_file.contracts:
         for fn in contract.functions:
             if fn.name.startswith("legit_"):
-                assert fn.name not in detections_fn_names
+                assert fn.canonical_name not in detections_fn_names
             elif fn.name.startswith("bug_"):
-                assert fn.name in detections_fn_names
+                detected_something = True
+                assert fn.canonical_name in detections_fn_names
 
     for fn in source_file.functions:
         if fn.name.startswith("legit_"):
-            assert fn.name not in detections_fn_names
+            assert fn.canonical_name not in detections_fn_names
         elif fn.name.startswith("bug_"):
-            assert fn.name in detections_fn_names
+            detected_something = True
+            assert fn.canonical_name in detections_fn_names
+    assert detected_something
 
 
 def check_detected_in_contracts(config, source_path):
@@ -130,11 +135,14 @@ def check_detected_in_contracts(config, source_path):
         contract.name for contract in detections_contracts if contract is not None
     ]
 
+    detected_something = False
     for cf in source_file.contracts:
         if cf.name.startswith("Legit_"):
             assert cf.name not in detections_contract_names
         elif cf.name.startswith("Bug_"):
+            detected_something = True
             assert cf.name in detections_contract_names
+    assert detected_something
 
 
 class TestNoReturnDetector:
@@ -214,16 +222,15 @@ def test_sources(self, config, tmp_path):
         test_file = "not_used.sol"
         test_source_path = tmp_path / test_file
         shutil.copyfile(SOURCES_PATH / test_file, test_source_path)
-        check_detected_in_functions(config, test_source_path)
         check_detected_in_contracts(config, test_source_path)
 
 
-class TestProxyContract:
+class TestProxyContractSelectorClashes:
     @pytest.fixture
     def config(self, tmp_path) -> WokeConfig:
         config_dict = {
             "compiler": {"solc": {"include_paths": ["./node_modules"]}},
-            "detectors": {"only": {"proxy-contract"}},
+            "detectors": {"only": {"proxy-contract-selector-clashes"}},
         }
         return WokeConfig.fromdict(
             config_dict,
@@ -232,7 +239,16 @@ def config(self, tmp_path) -> WokeConfig:
         )
 
     def test_sources(self, config, tmp_path):
-        test_file = "proxy_contract.sol"
+        test_file = "proxy_contract_selector_clashes.sol"
         test_source_path = tmp_path / test_file
         shutil.copyfile(SOURCES_PATH / test_file, test_source_path)
-        check_detected_in_contracts(config, test_source_path)
+
+        _, source_units = compile_project(test_source_path, config)
+        detections = detect(config, source_units)
+        logging.error(detections)
+        detections_fns = [_get_function_def(det.result.ir_node) for det in detections]
+        detections_fn_names = [
+            fn.canonical_name for fn in detections_fns if fn is not None
+        ]
+
+        assert "Proxy.bug_clash" in detections_fn_names

woke/analysis/detectors/__init__.py

@@ -14,7 +14,7 @@
 from .overflow_calldata_tuple_reencoding_bug import (
     OverflowCalldataTupleReencodingBugDetector,
 )
-from .proxy_contract import ProxyContractDetector
+from .proxy_contract_selector_clashes import ProxyContractSelectorClashDetector
 from .reentrancy import ReentrancyDetector
 from .unchecked_return_value import UncheckedFunctionReturnValueDetector
 from .unsafe_delegatecall import UnsafeDelegatecallDetector

woke/analysis/detectors/proxy_contract.py woke/analysis/detectors/proxy_contract_selector_clashes.py

@@ -1,12 +1,12 @@
 from functools import lru_cache
-from typing import List, Optional, Set, Tuple, Union
+from typing import Dict, List, Optional, Set, Tuple, Union
 
 from woke.analysis.detectors import DetectorAbc, DetectorResult, detector
 from woke.analysis.detectors.utils import (
     get_function_implementations,
     pair_function_call_arguments,
 )
-from woke.ast.enums import FunctionKind, GlobalSymbolsEnum, LiteralKind
+from woke.ast.enums import FunctionKind, GlobalSymbolsEnum, LiteralKind, Visibility
 from woke.ast.ir.abc import IrAbc
 from woke.ast.ir.declaration.abc import DeclarationAbc
 from woke.ast.ir.declaration.contract_definition import ContractDefinition
@@ -398,25 +398,70 @@ def detect_slot_usage(fn: FunctionDefinition, visited=None) -> List[DetectorResu
     return dets
 
 
-@detector(-1030, "proxy-contract")
-class ProxyContractDetector(DetectorAbc):
+def detect_selector_clashes(
+    proxy_contract: ContractDefinition,
+    impl_contract: ContractDefinition,
+    proxy_detection: DetectorResult,
+    impl_detection: DetectorResult,
+) -> List[DetectorResult]:
+    fn_whitelist = ["implementation"]
+
+    proxy_selectors = {}
+    for c in proxy_contract.linearized_base_contracts:
+        for f in c.functions + c.declared_variables:
+            if f.function_selector is not None:
+                proxy_selectors[f.function_selector] = f
+
+    impl_selectors = {}
+    for c in impl_contract.linearized_base_contracts:
+        for f in c.functions + c.declared_variables:
+            if f.function_selector is not None:
+                impl_selectors[f.function_selector] = f
+
+    clashes = []
+    for proxy_sel, proxy_fn in proxy_selectors.items():
+        if isinstance(proxy_fn, FunctionDefinition) and proxy_fn.name in fn_whitelist:
+            continue
+        if proxy_sel in impl_selectors:
+            clashes.append(
+                DetectorResult(
+                    proxy_fn,
+                    "Detected selector clash with implementation contract",
+                    related_info=(
+                        DetectorResult(
+                            impl_selectors[proxy_sel],
+                            "Implementation function with same selector",
+                            related_info=(impl_detection,),
+                        ),
+                        proxy_detection,
+                    ),
+                )
+            )
+    return clashes
+
+
+@detector(-1030, "proxy-contract-selector-clashes")
+class ProxyContractSelectorClashDetector(DetectorAbc):
     """
-    Detects proxy contracts based on fallback function and usage of slot variables and
+    Detects selector clashes in proxy and implementation contracts.
+    Proxy contracts are detected based on fallback function and usage of slot variables and
     implementation contracts that use same slots as proxy contracts
     """
 
     _proxy_detections: Set[Tuple[ContractDefinition, DetectorResult]]
     _proxy_associated_contracts: Set[ContractDefinition]
     _implementation_slots_detections: Set[Tuple[ContractDefinition, DetectorResult]]
-    _implementation_slots: Set[VariableDeclaration]
+    _implementation_slots: Dict[VariableDeclaration, List[ContractDefinition]]
 
     def __init__(self):
         self._proxy_detections: Set[Tuple[ContractDefinition, DetectorResult]] = set()
         self._proxy_associated_contracts: Set[ContractDefinition] = set()
         self._implementation_slots_detections: Set[
             Tuple[ContractDefinition, DetectorResult]
         ] = set()
-        self._implementation_slots: Set[VariableDeclaration] = set()
+        self._implementation_slots: Dict[
+            VariableDeclaration, List[ContractDefinition]
+        ] = {}
 
     def report(self) -> List[DetectorResult]:
         detections = []
@@ -428,9 +473,10 @@ def report(self) -> List[DetectorResult]:
                     continue
                 base_proxy_contracts.add(c)
 
+        proxy_detections = {}
         for (contract, det) in self._proxy_detections:
             if contract not in base_proxy_contracts:
-                detections.append(det)
+                proxy_detections[contract] = det
 
         base_impl_contracts = set()
         for (contract, _) in self._implementation_slots_detections:
@@ -439,22 +485,37 @@ def report(self) -> List[DetectorResult]:
                     continue
                 base_impl_contracts.add(c)
 
+        checked_pairs = set()
         impl_detections_contracts = set()
         for (contract, det) in self._implementation_slots_detections:
+            impl_slot = get_last_detection_node(det).parent
             if (
                 contract not in self._proxy_associated_contracts
                 and contract not in base_impl_contracts
-                and get_last_detection_node(det).parent in self._implementation_slots
+                and impl_slot in self._implementation_slots
                 and contract not in impl_detections_contracts
             ):
-                impl_detections_contracts.add(contract)
-                detections.append(
-                    DetectorResult(
+                for proxy_contract in self._implementation_slots[impl_slot]:
+                    if (
+                        proxy_contract,
+                        contract,
+                    ) in checked_pairs or proxy_contract not in proxy_detections:
+                        continue
+                    checked_pairs.add((proxy_contract, contract))
+                    impl_det = DetectorResult(
                         contract,
-                        "Detected implementation contract with slot used in proxy contract",
+                        f"Detected implementation contract with slot used in proxy contract {proxy_contract.name}",
                         related_info=(det,),
                     )
-                )
+
+                    dets = detect_selector_clashes(
+                        proxy_contract,
+                        contract,
+                        proxy_detections[proxy_contract],
+                        impl_det,
+                    )
+                    if len(dets) > 0:
+                        detections.extend(dets)
         return list(detections)
 
     def visit_contract_definition(self, node: ContractDefinition):
@@ -478,10 +539,22 @@ def visit_contract_definition(self, node: ContractDefinition):
                             )
                             for det in dets:
                                 last_det_node = get_last_detection_node(det)
-                                if isinstance(
-                                    last_det_node.parent, VariableDeclaration
+                                if (
+                                    isinstance(
+                                        last_det_node.parent, VariableDeclaration
+                                    )
+                                    and last_det_node.parent is not None
                                 ):
-                                    self._implementation_slots.add(last_det_node.parent)
+                                    if (
+                                        last_det_node.parent
+                                        not in self._implementation_slots
+                                    ):
+                                        self._implementation_slots[
+                                            last_det_node.parent
+                                        ] = []
+                                    self._implementation_slots[
+                                        last_det_node.parent
+                                    ].append(node)
                                 self._proxy_associated_contracts.add(node)
                                 for bc in node.linearized_base_contracts:
                                     self._proxy_associated_contracts.add(bc)