get_function_info tests

mythril/disassembler/disassembly.py

@@ -37,7 +37,7 @@ class Disassembly(object):
         )
 
         for index in jump_table_indices:
-            function_hash, jump_target, function_name = _get_function_info(
+            function_hash, jump_target, function_name = get_function_info(
                 index, self.instruction_list, signatures
             )
             self.func_hashes.append(function_hash)
@@ -52,7 +52,7 @@ class Disassembly(object):
         return asm.instruction_list_to_easm(self.instruction_list)
 
 
-def _get_function_info(index: int, instruction_list: list, signature_database: SignatureDb) -> (str, int, str):
+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

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"