Repository: ConsenSys/mythril-classic
Pull request #609 (2018-10-23 13:09:15)

Disassembly cleanup, testing and documentation


Patch:
diff --git a/mythril/disassembler/disassembly.py b/mythril/disassembler/disassembly.py
index e28a4d73..1fd98a73 100644
--- a/mythril/disassembler/disassembly.py
+++ b/mythril/disassembler/disassembly.py
@@ -4,57 +4,91 @@ import logging
 
 
 class Disassembly(object):
+    """
+    Disassembly class
 
-    def __init__(self, code, enable_online_lookup=False):
+    Stores bytecode, and its disassembly.
+    Additionally it will gather the following information on the existing functions in the disassembled code:
+     - function hashes
+     - function name to entry point mapping
+     - function entry point to function name mapping
+    """
+    def __init__(self, code: str, enable_online_lookup: bool=False):
+        self.bytecode = code
         self.instruction_list = asm.disassemble(util.safe_decode(code))
+
         self.func_hashes = []
-        self.func_to_addr = {}
-        self.addr_to_func = {}
-        self.bytecode = code
+        self.function_name_to_address = {}
+        self.address_to_function_name = {}
 
-        signatures = SignatureDb(enable_online_lookup=enable_online_lookup)  # control if you want to have online sighash lookups
+        signatures = SignatureDb(
+            enable_online_lookup=enable_online_lookup
+        )  # control if you want to have online signature hash lookups
         try:
             signatures.open()  # open from default locations
         except FileNotFoundError:
-            logging.info("Missing function signature file. Resolving of function names from signature file disabled.")
-
-        # Parse jump table & resolve function names
+            logging.info(
+                "Missing function signature file. Resolving of function names from signature file disabled."
+            )
 
         # Need to take from PUSH1 to PUSH4 because solc seems to remove excess 0s at the beginning for optimizing
-        jmptable_indices = asm.find_opcode_sequence([("PUSH1", "PUSH2", "PUSH3", "PUSH4"), ("EQ",)],
-                                                    self.instruction_list)
-
-        for i in jmptable_indices:
-            func_hash = self.instruction_list[i]['argument']
-
-            # Append with missing 0s at the beginning
-            func_hash = "0x" + func_hash[2:].rjust(8, "0")
-
-            self.func_hashes.append(func_hash)
-            try:
-                # tries local cache, file and optional online lookup
-                # may return more than one function signature. since we cannot probe for the correct one we'll use the first
-                func_names = signatures.get(func_hash)
-                if len(func_names) > 1:
-                    # ambigious result
-                    func_name = "**ambiguous** %s" % func_names[0]  # return first hit but note that result was ambiguous
-                else:
-                    # only one item
-                    func_name = func_names[0]
-            except KeyError:
-                func_name = "_function_" + func_hash
-
-            try:
-                offset = self.instruction_list[i + 2]['argument']
-                jump_target = int(offset, 16)
-
-                self.func_to_addr[func_name] = jump_target
-                self.addr_to_func[jump_target] = func_name
-            except:
-                continue
+        jump_table_indices = asm.find_opcode_sequence(
+            [("PUSH1", "PUSH2", "PUSH3", "PUSH4"), ("EQ",)], self.instruction_list
+        )
+
+        for index in jump_table_indices:
+            function_hash, jump_target, function_name = get_function_info(
+                index, self.instruction_list, signatures
+            )
+            self.func_hashes.append(function_hash)
+
+            if jump_target is not None and function_name is not None:
+                self.function_name_to_address[function_name] = jump_target
+                self.address_to_function_name[jump_target] = function_name
 
         signatures.write()  # store resolved signatures (potentially resolved online)
 
     def get_easm(self):
-        # todo: tintinweb - print funcsig resolved data from self.addr_to_func?
         return asm.instruction_list_to_easm(self.instruction_list)
+
+
+def get_function_info(index: int, instruction_list: list, signature_database: SignatureDb) -> (str, int, str):
+    """
+    Finds the function information for a call table entry
+    Solidity uses the first 4 bytes of the calldata to indicate which function the message call should execute
+    The generated code that directs execution to the correct function looks like this:
+    - PUSH function_hash
+    - EQ
+    - PUSH entry_point
+    - JUMPI
+
+    This function takes an index that points to the first instruction, and from that finds out the function hash,
+    function entry and the function name.
+
+    :param index: Start of the entry pattern
+    :param instruction_list: Instruction list for the contract that is being analyzed
+    :param signature_database: Database used to map function hashes to their respective function names
+    :return: function hash, function entry point, function name
+    """
+
+    # Append with missing 0s at the beginning
+    function_hash = "0x" + instruction_list[index]["argument"][2:].rjust(8, "0")
+
+    function_names = signature_database.get(function_hash)
+    if len(function_names) > 1:
+        # In this case there was an ambiguous result
+        function_name = (
+            "**ambiguous** {}".format(function_names[0])
+        )
+    elif len(function_names) == 1:
+        function_name = function_names[0]
+    else:
+        function_name = "_function_" + function_hash
+
+    try:
+        offset = instruction_list[index + 2]["argument"]
+        entry_point = int(offset, 16)
+    except (KeyError, IndexError):
+        return function_hash, None, None
+
+    return function_hash, entry_point, function_name
diff --git a/mythril/laser/ethereum/svm.py b/mythril/laser/ethereum/svm.py
index ed5f5724..d7809705 100644
--- a/mythril/laser/ethereum/svm.py
+++ b/mythril/laser/ethereum/svm.py
@@ -1,4 +1,5 @@
 import logging
+from mythril.disassembler.disassembly import Disassembly
 from mythril.laser.ethereum.state import WorldState
 from mythril.laser.ethereum.transaction import TransactionStartSignal, TransactionEndSignal, \
     ContractCreationTransaction
@@ -259,10 +260,9 @@ class LaserEVM:
 
         environment = state.environment
         disassembly = environment.code
-        if address in state.environment.code.addr_to_func:
+        if address in disassembly.address_to_function_name:
             # Enter a new function
-
-            environment.active_function_name = disassembly.addr_to_func[address]
+            environment.active_function_name = disassembly.address_to_function_name[address]
             new_node.flags |= NodeFlags.FUNC_ENTRY
 
             logging.debug(
diff --git a/mythril/support/signatures.py b/mythril/support/signatures.py
index 48d60297..a5290088 100644
--- a/mythril/support/signatures.py
+++ b/mythril/support/signatures.py
@@ -164,9 +164,12 @@ class SignatureDb(object):
             except FourByteDirectoryOnlineLookupError as fbdole:
                 self.online_directory_unavailable_until = time.time() + 2 * 60  # wait at least 2 mins to try again
                 logging.warning("online function signature lookup not available. will not try to lookup hash for the next 2 minutes. exception: %r" % fbdole)
+
+        if sighash not in self.signatures:
+            return []
         if type(self.signatures[sighash]) != list:
             return [self.signatures[sighash]]
-        return self.signatures[sighash]  # raise keyerror
+        return self.signatures[sighash]
 
     def __getitem__(self, item):
         """


Test patch:
diff --git a/tests/disassembler/__init__.py b/tests/disassembler/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/disassembler/disassembly.py b/tests/disassembler/disassembly.py
new file mode 100644
index 00000000..49bca31d
--- /dev/null
+++ b/tests/disassembler/disassembly.py
@@ -0,0 +1,61 @@
+from mythril.disassembler.disassembly import *
+
+instruction_list = [
+    {"opcode": "PUSH4", "argument": "0x10203040"},
+    {"opcode": "EQ"},
+    {"opcode": "PUSH4", "argument": "0x40302010"},
+    {"opcode": "JUMPI"},
+]
+
+
+def test_get_function_info(mocker):
+    # Arrange
+    global instruction_list
+
+    signature_database_mock = SignatureDb()
+    mocker.patch.object(signature_database_mock, "get")
+    signature_database_mock.get.return_value = ["function_name"]
+
+    # Act
+    function_hash, entry_point, function_name = get_function_info(
+        0, instruction_list, signature_database_mock
+    )
+
+    # Assert
+    assert function_hash == "0x10203040"
+    assert entry_point == 0x40302010
+    assert function_name == "function_name"
+
+
+def test_get_function_info_multiple_names(mocker):
+    # Arrange
+    global instruction_list
+
+    signature_database_mock = SignatureDb()
+    mocker.patch.object(signature_database_mock, "get")
+    signature_database_mock.get.return_value = ["function_name", "another_name"]
+
+    # Act
+    function_hash, entry_point, function_name = get_function_info(
+        0, instruction_list, signature_database_mock
+    )
+
+    # Assert
+    assert function_name == "**ambiguous** function_name"
+
+
+def test_get_function_info_no_names(mocker):
+    # Arrange
+    global instruction_list
+
+    signature_database_mock = SignatureDb()
+    mocker.patch.object(signature_database_mock, "get")
+    signature_database_mock.get.return_value = []
+
+    # Act
+    function_hash, entry_point, function_name = get_function_info(
+        0, instruction_list, signature_database_mock
+    )
+
+    # Assert
+    assert function_name == "_function_0x10203040"