Fix #2614: distance command works with function symbols now (#3033)

* gdb: suggest &main instead of main (address of symbol) in commands

* remove raise

* Revert "gdb: suggest &main instead of main (address of symbol) in commands"

This reverts commit 64e6d85c8e.

* Fix distance

* Remove todo

pwndbg/commands/distance.py

@@ -10,8 +10,10 @@ from pwndbg.commands import CommandCategory
 parser = argparse.ArgumentParser(
     description="Print the distance between the two arguments, or print the offset to the address's page base."
 )
-parser.add_argument("a", type=int, help="The first address.")
-parser.add_argument("b", nargs="?", default=None, type=int, help="The second address.")
+parser.add_argument("a", type=pwndbg.commands.AddressExpr, help="The first address.")
+parser.add_argument(
+    "b", nargs="?", default=None, type=pwndbg.commands.AddressExpr, help="The second address."
+)
 
 
 @pwndbg.commands.Command(parser, category=CommandCategory.MEMORY)

tests/gdb-tests/tests/test_command_distance.py

@@ -11,7 +11,31 @@ REFERENCE_BINARY = tests.binaries.get("reference-binary.out")
 def test_command_distance(start_binary):
     start_binary(REFERENCE_BINARY)
 
+    # Test against regs
     rsp = pwndbg.aglib.regs.rsp
     result = gdb.execute("distance $rsp $rsp+0x10", to_string=True)
-
     assert result == f"{rsp:#x}->{rsp + 0x10:#x} is 0x10 bytes (0x2 words)\n"
+
+    # Test if it works with symbols
+    rip = pwndbg.aglib.regs.rip
+
+    main = pwndbg.aglib.symbol.lookup_symbol_addr("main")
+    break_here = pwndbg.aglib.symbol.lookup_symbol_addr("break_here")
+
+    diff = break_here - main
+
+    # Test symbol (function address) and its proper &symbol address
+    for sym1 in ("main", "&main"):
+        for sym2 in ("break_here", "&break_here"):
+            result = gdb.execute(f"distance {sym1} {sym2}", to_string=True)
+            assert result == f"{main:#x}->{break_here:#x} is {diff:#x} bytes ({diff//8:#x} words)\n"
+
+    # Test if it works with reg + symbol
+    diff = break_here - rip
+    result = gdb.execute("distance $rip &break_here", to_string=True)
+    assert result == f"{rip:#x}->{break_here:#x} is {diff:#x} bytes ({diff//8:#x} words)\n"
+
+    # Test if it works with symbol + reg
+    diff = rip - break_here
+    result = gdb.execute("distance &break_here $rip", to_string=True)
+    assert result == f"{break_here:#x}->{rip:#x} is {diff:#x} bytes ({diff//8:#x} words)\n"