angr/tests/analyses/test_rust_utils.py
Vedant Soni 1808e7fadd
Migrate rust demangling to pydemumble (#6663)
* use pydemumble to demangle rust symbols

* update tests

* remove rust_demangle and unpin pydemumble from pyproject

* pin pydemumble to latest version
2026-07-23 17:32:00 -07:00

380 lines
14 KiB
Python

from __future__ import annotations
import angr
from angr.ailment import Block
from angr.ailment.expression import (
BinaryOp,
Call,
Const,
UnaryOp,
VirtualVariable,
VirtualVariableCategory,
)
from angr.ailment.statement import Assignment
from angr.rust.optimization_passes.utils import (
extract_callee,
extract_str,
extract_str_from_addr,
replace_argument_pairs,
)
from angr.rust.utils.ail import (
deref_vvar_and_offset,
extract_vvar_and_offset,
find_call,
get_terminal_call,
has_call,
unwrap_combo_reg_vvar_reference,
unwrap_stack_vvar_reference,
unwrap_stack_vvar_reference_with_offset,
)
from angr.rust.utils.demangler import _is_rust_hash, demangle, normalize
def _stack_vvar(vvar_id: int = 0, offset: int = 16) -> VirtualVariable:
return VirtualVariable(vvar_id, vvar_id, 64, VirtualVariableCategory.STACK, oident=offset)
def _reg_vvar(vvar_id: int = 0, offset: int = 16) -> VirtualVariable:
return VirtualVariable(vvar_id, vvar_id, 64, VirtualVariableCategory.REGISTER, oident=offset)
def _combo_vvar(vvar_id: int = 0) -> VirtualVariable:
return VirtualVariable(vvar_id, vvar_id, 128, VirtualVariableCategory.COMBO_REGISTER, oident=(16, 24))
def _ref(operand) -> UnaryOp:
return UnaryOp(0, "Reference", operand)
def _const(value: int, bits: int = 64) -> Const:
return Const(0, value, bits)
def test_unwrap_stack_vvar_reference_returns_underlying_vvar():
vvar = _stack_vvar(offset=8)
expr = _ref(vvar)
assert unwrap_stack_vvar_reference(expr) == vvar
def test_unwrap_stack_vvar_reference_rejects_register_vvar():
expr = _ref(_reg_vvar(offset=8))
assert unwrap_stack_vvar_reference(expr) is None
def test_unwrap_stack_vvar_reference_rejects_non_reference_unary_op():
expr = UnaryOp(0, "Neg", _stack_vvar())
assert unwrap_stack_vvar_reference(expr) is None
def test_unwrap_combo_reg_vvar_reference_returns_underlying_vvar_only_for_combo():
combo = _combo_vvar()
assert unwrap_combo_reg_vvar_reference(_ref(combo)) == combo
assert unwrap_combo_reg_vvar_reference(_ref(_reg_vvar())) is None
def test_deref_vvar_and_offset_handles_zero_and_offset_cases():
vvar = _stack_vvar(offset=8)
# Pure reference: offset 0
assert deref_vvar_and_offset(_ref(vvar)) == (vvar, 0)
# &vvar + 12 → (vvar, 12)
add = BinaryOp(0, "Add", [_ref(vvar), _const(12)])
assert deref_vvar_and_offset(add) == (vvar, 12)
def test_deref_vvar_and_offset_rejects_other_shapes():
# Add of two vvars (no Reference): not the pattern.
bogus = BinaryOp(0, "Add", [_stack_vvar(0), _stack_vvar(1)])
assert deref_vvar_and_offset(bogus) == (None, None)
# Sub of &vvar and const: also not the pattern.
sub = BinaryOp(0, "Sub", [_ref(_stack_vvar()), _const(4)])
assert deref_vvar_and_offset(sub) == (None, None)
def test_extract_vvar_and_offset_strips_constant_offset():
vvar = _reg_vvar(offset=24)
assert extract_vvar_and_offset(vvar) == (vvar, 0)
assert extract_vvar_and_offset(BinaryOp(0, "Add", [vvar, _const(7)])) == (vvar, 7)
assert extract_vvar_and_offset(_const(0)) == (None, None)
def test_unwrap_stack_vvar_reference_with_offset_returns_zero_for_plain_reference():
vvar = _stack_vvar(offset=4)
assert unwrap_stack_vvar_reference_with_offset(_ref(vvar)) == (vvar, 0)
def test_unwrap_stack_vvar_reference_with_offset_rejects_combo_register():
assert unwrap_stack_vvar_reference_with_offset(_ref(_combo_vvar())) == (None, None)
def test_find_call_finds_call_inside_block_statement_assignment():
target_call = Call(0, "callee_a", args=[])
block = Block(
0x4000,
0,
statements=[Assignment(0, _stack_vvar(0), target_call)],
)
found = find_call(block)
assert found == target_call
assert has_call(block) is True
def test_find_call_returns_none_when_no_call():
block = Block(0x4010, 0, statements=[Assignment(0, _stack_vvar(0), _const(1))])
assert find_call(block) is None
assert has_call(block) is False
def test_get_terminal_call_returns_call_at_block_end():
final_call = Call(1, "terminal_callee", args=[])
# get_terminal_call defensively handles raw Call expressions sitting at the end of a block.
block = Block(0x4020, 0, statements=[final_call]) # pyright: ignore[reportArgumentType]
assert get_terminal_call(block) is final_call
def test_get_terminal_call_descends_into_assignment_with_call_rhs():
inner_call = Call(2, "inner_callee", args=[])
block = Block(
0x4030,
0,
statements=[Assignment(0, _stack_vvar(0), inner_call)],
)
assert get_terminal_call(block) == inner_call
def test_get_terminal_call_returns_none_for_empty_block():
block = Block(0x4040, 0, statements=[])
assert get_terminal_call(block) is None
def test_replace_argument_pairs_no_args_short_circuits():
call = Call(0, "callee", args=[])
assert replace_argument_pairs(call, lambda a, b: (False, ())) is call
def test_replace_argument_pairs_collapses_consecutive_int_pairs():
args = [_const(1), _const(2), _const(3), _const(4), _const(5)]
call = Call(0, "callee", args=args)
def merge_pair(a, b):
if isinstance(a, Const) and isinstance(b, Const):
return True, [_const(int(a.value) + int(b.value))]
return False, ()
new_call = replace_argument_pairs(call, merge_pair)
assert new_call is not call
assert new_call.args is not None
new_args = list(new_call.args)
# Pairs are consumed left-to-right without reusing replacements:
# (1,2)→3, (3,4)→7, leaving 5 unpaired ⇒ [3, 7, 5].
values = []
for a in new_args:
assert isinstance(a, Const)
values.append(a.value)
assert values == [3, 7, 5]
def test_replace_argument_pairs_returns_original_when_no_replacement_happens():
args = [_const(1), _const(2), _const(3)]
call = Call(0, "callee", args=args)
def never_replace(a, b):
del a, b
return False, ()
assert replace_argument_pairs(call, never_replace) is call
def _project_with_synthetic_function(addr: int, name: str):
project = angr.load_shellcode(b"\xc3" * 16, arch="amd64")
project.kb.functions.function(addr=addr, name=name, create=True)
return project
def test_extract_callee_returns_function_for_call_with_const_target():
project = _project_with_synthetic_function(0x0, "callee")
call = Call(0, Const(0, 0x0, 64), args=[])
result = extract_callee(call, project.kb)
assert result is project.kb.functions[0x0]
def test_extract_callee_walks_block_to_terminal_call():
project = _project_with_synthetic_function(0x0, "callee")
call = Call(0, Const(0, 0x0, 64), args=[])
block = Block(0x100, 0, statements=[call]) # pyright: ignore[reportArgumentType]
assert extract_callee(block, project.kb) is project.kb.functions[0x0]
def test_extract_callee_returns_none_for_unknown_target():
project = _project_with_synthetic_function(0x0, "callee")
unknown = Call(0, Const(0, 0xDEADBEEF, 64), args=[])
assert extract_callee(unknown, project.kb) is None
def test_extract_callee_returns_none_for_non_block_non_call():
project = _project_with_synthetic_function(0x0, "callee")
assert extract_callee(Const(0, 0, 64), project.kb) is None
def test_extract_str_short_circuits_for_zero_length():
project = angr.load_shellcode(b"\x90", arch="amd64")
assert extract_str(project, str_ptr=0xDEAD, str_len=0) == ""
def test_extract_str_rejects_negative_pointer():
project = angr.load_shellcode(b"\x90", arch="amd64")
assert extract_str(project, str_ptr=-1, str_len=4) is None
def test_extract_str_returns_none_when_no_readable_section_covers_pointer():
# load_shellcode produces a Blob with no sections, so find_section_containing
# always returns None and the helper falls through to None.
project = angr.load_shellcode(b"\x90", arch="amd64")
assert extract_str(project, str_ptr=0x0, str_len=4) is None
def test_extract_str_from_addr_rejects_negative_address():
project = angr.load_shellcode(b"\x90", arch="amd64")
assert extract_str_from_addr(project, addr=-1) is None
def test_extract_str_from_addr_returns_none_when_no_section_covers_address():
project = angr.load_shellcode(b"\x90", arch="amd64")
assert extract_str_from_addr(project, addr=0x0) is None
def test_is_rust_hash_accepts_only_17_char_hex_with_h_prefix():
assert _is_rust_hash("h0123456789abcdef") is True
assert _is_rust_hash("h12") is False # too short
assert _is_rust_hash("x0123456789abcdef") is False # wrong prefix
assert _is_rust_hash("h0123456789abcdez") is False # non-hex digit
def test_demangle_strips_trailing_rust_hash_segment():
mangled = "_ZN4core3fmt9Formatter9write_str17h0123456789abcdefE"
assert demangle(mangled) == "core::fmt::Formatter::write_str"
def test_demangle_falls_back_to_original_string_when_unrecognized():
raw = "not_a_rust_symbol"
assert demangle(raw) == raw
def test_demangle_survives_malformed_rust_v0_symbol():
# Regression for angr#6598: a garbage name that merely *looks* like a Rust v0 symbol
# ("_R" prefix). pydemumble rejects it (returns an empty string) and demangle() must
# fall back to the original string.
raw = "_RBOJyX"
assert demangle(raw) == raw
def test_demangle_survives_assorted_malformed_v0_symbols():
# Various truncated/garbage "_R..." names that crashed the previously-used third-party
# demangler. Every one must fall through to a non-empty string without raising.
for raw in ("_R", "_RNvC", "_RNvNtCs", "_RINvNtCs", "_RNvMC0", "_RNCNvC0"):
result = demangle(raw)
assert isinstance(result, str)
assert result
def test_normalize_strips_generic_type_arguments_when_monopolizing():
assert normalize("alloc::vec::Vec<u8>::push") == "alloc::vec::Vec::push"
def test_normalize_keeps_concrete_type_or_trait_name_per_flag():
name = "<core::option::Option<T> as core::fmt::Debug>::fmt"
assert normalize(name) == "core::option::Option::fmt"
assert normalize(name, use_trait_name=True) == "core::fmt::Debug::fmt"
def test_normalize_handles_impl_as_pattern_using_trait_name():
name = "<impl alloc::vec::Vec<u8> as core::iter::Iterator>::next"
assert normalize(name) == "core::iter::Iterator::next"
def test_normalize_concise_returns_only_last_path_segment():
assert normalize("foo::bar::baz", concise=True) == "baz"
def test_is_rust_hash_rejects_empty_and_short_inputs():
assert _is_rust_hash("") is False
assert _is_rust_hash("h") is False
def test_is_rust_hash_rejects_uppercase_hex_digits():
# Cargo's hash suffix is always lowercase hex; uppercase must not match.
assert _is_rust_hash("hABCDEF0123456789") is False
def test_demangle_handles_empty_string_without_raising():
assert demangle("") == ""
def test_demangle_returns_input_for_legacy_symbol_with_no_hash_suffix():
# A well-formed legacy symbol whose final segment isn't a 17-char h-hex hash:
# the demangler still parses it, but our hash-stripping branch must not fire.
mangled = "_ZN4core3fmt9Formatter9write_strE"
out = demangle(mangled)
# The wrapper must not silently truncate the tail.
assert out.endswith("write_str")
def test_demangle_passes_through_non_rust_input():
# Inputs that aren't Rust mangled names must fall through unchanged. In particular,
# _Z3fooi (itanium C++) and MSVC names would demangle under pydemumble, but the
# wrapper is Rust-only and must not touch them.
for raw in ("plain_c_symbol", "main", "?something@msvc@@", "_Z3fooi"):
assert demangle(raw) == raw
def test_demangle_decodes_legacy_dollar_escapes():
mangled = "_ZN42_$LT$$RF$T$u20$as$u20$core..fmt..Debug$GT$3fmt17h517074eb1cb2b995E"
assert demangle(mangled) == "<&T as core::fmt::Debug>::fmt"
def test_demangle_decodes_legacy_closure_escapes():
mangled = "_ZN4core3ptr42drop_in_place$LT$alloc..string..String$GT$17h5d180b0b0e91564fE"
assert demangle(mangled) == "core::ptr::drop_in_place<alloc::string::String>"
def test_demangle_drops_llvm_suffix():
# Compiler-generated ".llvm.<id>" suffixes are dropped, matching rustc-demangle.
mangled = (
"_ZN3std2rt19lang_start_internal28_$u7b$$u7b$closure$u7d$$u7d$17hf421b6f6b8a4a2f4E.llvm.9325873435131735662"
)
assert demangle(mangled) == "std::rt::lang_start_internal::{{closure}}"
def test_demangle_keeps_numeric_instantiation_suffix_attached():
# A trailing ".<n>" (e.g. from symbol versioning/instantiation) stays attached to the
# demangled name, not rendered in demumble's " (.n)" style.
mangled = "_RNvNtNtCsjrHSEGnQ3l9_3std2io5stdio19OUTPUT_CAPTURE_USED.0"
assert demangle(mangled) == "std::io::stdio::OUTPUT_CAPTURE_USED.0"
def test_demangle_v0_symbol_with_generics():
assert demangle("_RINvNtC3std3mem8align_ofjE") == "std::mem::align_of::<usize>"
def test_normalize_strips_nested_generic_brackets():
# Multi-level generics must collapse fully under monopolize.
assert normalize("alloc::vec::Vec<Box<u8>>::push") == "alloc::vec::Vec::push"
def test_normalize_monopolize_false_preserves_generic_arguments():
name = "alloc::vec::Vec<u8>::push"
assert normalize(name, monopolize=False) == name
def test_normalize_concise_with_trait_uses_last_segment_after_trait_substitution():
name = "<core::option::Option<T> as core::fmt::Debug>::fmt"
assert normalize(name, use_trait_name=True, concise=True) == "fmt"
def test_normalize_idempotent_under_repeated_calls():
name = "<impl alloc::vec::Vec<u8> as core::iter::Iterator>::next"
once = normalize(name)
twice = normalize(once)
assert once == twice