diff --git a/mythril/disassembler/disassembly.py b/mythril/disassembler/disassembly.py index 48682f5a..640c56cb 100644 --- a/mythril/disassembler/disassembly.py +++ b/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 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"