Merge pull request #1522 from crytic/webthethird-dev-upgradeability-c…

…omments

Handle custom upgradeability comments for contracts

slither/core/declarations/contract.py

@@ -88,6 +88,7 @@ def __init__(self, compilation_unit: "SlitherCompilationUnit", scope: "FileScope
 
         self._is_upgradeable: Optional[bool] = None
         self._is_upgradeable_proxy: Optional[bool] = None
+        self._upgradeable_version: Optional[str] = None
 
         self.is_top_level = False  # heavily used, so no @property
 
@@ -1221,6 +1222,10 @@ def is_upgradeable(self) -> bool:
                         break
         return self._is_upgradeable
 
+    @is_upgradeable.setter
+    def is_upgradeable(self, upgradeable: bool):
+        self._is_upgradeable = upgradeable
+
     @property
     def is_upgradeable_proxy(self) -> bool:
         from slither.core.cfg.node import NodeType
@@ -1246,6 +1251,18 @@ def is_upgradeable_proxy(self) -> bool:
                                     return self._is_upgradeable_proxy
         return self._is_upgradeable_proxy
 
+    @is_upgradeable_proxy.setter
+    def is_upgradeable_proxy(self, upgradeable_proxy: bool):
+        self._is_upgradeable_proxy = upgradeable_proxy
+
+    @property
+    def upgradeable_version(self) -> Optional[str]:
+        return self._upgradeable_version
+
+    @upgradeable_version.setter
+    def upgradeable_version(self, version_name: str):
+        self._upgradeable_version = version_name
+
     # endregion
     ###################################################################################
     ###################################################################################

slither/solc_parsing/declarations/contract.py

@@ -1,4 +1,5 @@
 import logging
+import re
 from typing import List, Dict, Callable, TYPE_CHECKING, Union, Set
 
 from slither.core.declarations import Modifier, Event, EnumContract, StructureContract, Function
@@ -65,8 +66,10 @@ def __init__(self, slither_parser: "SlitherCompilationUnitSolc", contract: Contr
         # Export info
         if self.is_compact_ast:
             self._contract.name = self._data["name"]
+            self._handle_comment(self._data)
         else:
             self._contract.name = self._data["attributes"][self.get_key()]
+            self._handle_comment(self._data["attributes"])
 
         self._contract.id = self._data["id"]
 
@@ -699,6 +702,24 @@ def delete_content(self):
         self._usingForNotParsed = []
         self._customErrorParsed = []
 
+    def _handle_comment(self, attributes: Dict) -> None:
+        if (
+            "documentation" in attributes
+            and attributes["documentation"] is not None
+            and "text" in attributes["documentation"]
+        ):
+            candidates = attributes["documentation"]["text"].replace("\n", ",").split(",")
+
+            for candidate in candidates:
+                if "@custom:security isDelegatecallProxy" in candidate:
+                    self._contract.is_upgradeable_proxy = True
+                if "@custom:security isUpgradeable" in candidate:
+                    self._contract.is_upgradeable = True
+
+                version_name = re.search(r"@custom:version name=([\w-]+)", candidate)
+                if version_name:
+                    self._contract.upgradeable_version = version_name.group(1)
+
     # endregion
     ###################################################################################
     ###################################################################################

tests/custom_comments/upgrade.sol

@@ -0,0 +1,16 @@
+/// @custom:security isDelegatecallProxy
+contract Proxy{
+
+}
+
+/// @custom:security isUpgradeable
+/// @custom:version name=version-0
+contract V0{
+
+}
+
+/// @custom:security isUpgradeable
+/// @custom:version name=version_1
+contract V1{
+
+}

tests/test_features.py

@@ -9,7 +9,7 @@
 from slither.detectors.abstract_detector import AbstractDetector
 
 
-def _run_all_detectors(slither: Slither):
+def _run_all_detectors(slither: Slither) -> None:
     detectors = [getattr(all_detectors, name) for name in dir(all_detectors)]
     detectors = [d for d in detectors if inspect.isclass(d) and issubclass(d, AbstractDetector)]
 
@@ -19,15 +19,15 @@ def _run_all_detectors(slither: Slither):
     slither.run_detectors()
 
 
-def test_node():
+def test_node() -> None:
     # hardhat must have been installed in tests/test_node_modules
     # For the CI its done through the github action config
 
     slither = Slither("./tests/test_node_modules")
     _run_all_detectors(slither)
 
 
-def test_collision():
+def test_collision() -> None:
 
     standard_json = SolcStandardJson()
     standard_json.add_source_file("./tests/collisions/a.sol")
@@ -39,14 +39,33 @@ def test_collision():
     _run_all_detectors(slither)
 
 
-def test_cycle():
+def test_cycle() -> None:
     slither = Slither("./tests/test_cyclic_import/a.sol")
     _run_all_detectors(slither)
 
 
-def test_funcion_id_rec_structure():
+def test_funcion_id_rec_structure() -> None:
     solc_select.switch_global_version("0.8.0", always_install=True)
     slither = Slither("./tests/function_ids/rec_struct-0.8.sol")
     for compilation_unit in slither.compilation_units:
         for function in compilation_unit.functions:
             assert function.solidity_signature
+
+
+def test_upgradeable_comments() -> None:
+    solc_select.switch_global_version("0.8.10", always_install=True)
+    slither = Slither("./tests/custom_comments/upgrade.sol")
+    compilation_unit = slither.compilation_units[0]
+    proxy = compilation_unit.get_contract_from_name("Proxy")[0]
+
+    assert proxy.is_upgradeable_proxy
+
+    v0 = compilation_unit.get_contract_from_name("V0")[0]
+
+    assert v0.is_upgradeable
+    print(v0.upgradeable_version)
+    assert v0.upgradeable_version == "version-0"
+
+    v1 = compilation_unit.get_contract_from_name("V1")[0]
+    assert v0.is_upgradeable
+    assert v1.upgradeable_version == "version_1"