From ea24e614bec1765317ab27e629c1ec611a563546 Mon Sep 17 00:00:00 2001
From: Josselin <josselin@trailofbits.com>
Date: Mon, 4 Mar 2019 09:25:14 +0000
Subject: [PATCH] Improve slither-check-upgradability + add travis unit tests

---
 scripts/travis_test_upgradability.sh          | 12 +++++
 .../contract_initialization.sol               | 49 +++++++++++++++++++
 tests/check-upgradability/test_5.txt          | 13 +++++
 utils/upgradability/__main__.py               |  1 -
 utils/upgradability/check_initialization.py   |  4 +-
 5 files changed, 76 insertions(+), 3 deletions(-)
 create mode 100644 tests/check-upgradability/contract_initialization.sol
 create mode 100644 tests/check-upgradability/test_5.txt

diff --git a/scripts/travis_test_upgradability.sh b/scripts/travis_test_upgradability.sh
index 92bad4490..4a83893f6 100755
--- a/scripts/travis_test_upgradability.sh
+++ b/scripts/travis_test_upgradability.sh
@@ -44,7 +44,19 @@ then
     exit -1
 fi
 
+slither-check-upgradability "$DIR_TESTS/proxy.sol" Proxy "$DIR_TESTS/contract_initialization.sol" Contract_no_bug --solc solc-0.5.0 > test_5.txt 2>&1 
+DIFF=$(diff test_5.txt "$DIR_TESTS/test_5.txt")
+if [  "$DIFF" != "" ] 
+then
+    echo "slither-check-upgradability failed"
+    cat test_5.txt
+    cat "$DIR_TESTS/test_5.txt"
+    exit -1
+fi
+
 rm test_1.txt
 rm test_2.txt
 rm test_3.txt
 rm test_4.txt
+rm test_5.txt
+
diff --git a/tests/check-upgradability/contract_initialization.sol b/tests/check-upgradability/contract_initialization.sol
new file mode 100644
index 000000000..d17125ee9
--- /dev/null
+++ b/tests/check-upgradability/contract_initialization.sol
@@ -0,0 +1,49 @@
+contract Initializable{
+
+    address destination;
+
+    modifier initializer(){
+        _;
+    }
+
+}
+
+contract Contract_no_bug is Initializable{
+
+    function initialize() public initializer{
+
+    }
+
+}
+
+contract Contract_lack_to_call_modifier is Initializable{
+
+    function initialize() public {
+
+    }
+}
+
+contract Contract_not_called_super_init is Contract_no_bug{
+
+    function initialize() public initializer{
+
+    }
+
+}
+
+contract Contract_no_bug_inherits is Contract_no_bug{
+
+    function initialize() public initializer{
+        Contract_no_bug.initialize();
+    }
+
+}
+
+contract Contract_double_call is Contract_no_bug, Contract_no_bug_inherits{
+
+    function initialize() public initializer{
+        Contract_no_bug_inherits.initialize();
+        Contract_no_bug.initialize();
+    }
+
+}
diff --git a/tests/check-upgradability/test_5.txt b/tests/check-upgradability/test_5.txt
new file mode 100644
index 000000000..2d0fc7c1a
--- /dev/null
+++ b/tests/check-upgradability/test_5.txt
@@ -0,0 +1,13 @@
+INFO:CheckInitialization:Contract_lack_to_call_modifier.initialize does not call initializer
+INFO:CheckInitialization:Missing call to Contract_no_bug.initialize in Contract_not_called_super_init
+INFO:CheckInitialization:Contract_no_bug.initialize() is called multiple time in Contract_double_call
+INFO:CheckInitialization:Check the deployement script to ensure that these functions are called:
+Contract_no_bug needs to be initialized by initialize()
+Contract_lack_to_call_modifier needs to be initialized by initialize()
+Contract_not_called_super_init needs to be initialized by initialize()
+Contract_no_bug_inherits needs to be initialized by initialize()
+Contract_double_call needs to be initialized by initialize()
+
+INFO:CompareFunctions:No function id collision found
+INFO:VariablesOrder:Variable in the proxy: destination address
+INFO:VariablesOrder:No error found (variables ordering proxy <-> implementation)
diff --git a/utils/upgradability/__main__.py b/utils/upgradability/__main__.py
index baa276b5f..d372c2739 100644
--- a/utils/upgradability/__main__.py
+++ b/utils/upgradability/__main__.py
@@ -49,7 +49,6 @@ def main():
 
     check_initialization(v1)
 
-
     if not args.new_version:
         compare_function_ids(v1, v1_name, proxy, proxy_name)
         compare_variables_order_proxy(v1, v1_name, proxy, proxy_name)
diff --git a/utils/upgradability/check_initialization.py b/utils/upgradability/check_initialization.py
index 78f0c82b5..2e8f84fbc 100644
--- a/utils/upgradability/check_initialization.py
+++ b/utils/upgradability/check_initialization.py
@@ -11,7 +11,7 @@ class MultipleInitTarget(Exception):
     pass
 
 def _get_initialize_functions(contract):
-    return [f for father in contract.inheritance for f in father.functions_not_inherited if f.name == 'initialize']
+    return [f for father in contract.inheritance + [contract] for f in father.functions_not_inherited if f.name == 'initialize']
 
 def _get_all_internal_calls(function):
     all_ir = function.all_slithir_operations()
@@ -61,7 +61,7 @@ def check_initialization(s):
             for f in missing_calls:
                 logger.info(red(f'Missing call to {f.contract.name}.{f.name} in {contract.name}'))
                 missing_call = True
-            double_calls = [f for f in all_init_functions_called if all_init_functions_called.count(f) > 1]
+            double_calls = list(set([f for f in all_init_functions_called if all_init_functions_called.count(f) > 1]))
             for f in double_calls:
                 logger.info(red(f'{f.contract.name + "." + f.full_name} is called multiple time in {contract.name}'))
                 double_calls_found = True