From 59638cc222ab410de0c12c16560641f98a718d01 Mon Sep 17 00:00:00 2001
From: Disconnect3d <dominik.b.czarnota@gmail.com>
Date: Wed, 28 May 2025 01:03:57 +0200
Subject: [PATCH] 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 64e6d85c8ec9d95f6b7cc98d2d4e239236c8afd5.

* Fix distance

* Remove todo
---
 pwndbg/commands/distance.py                   |  6 +++--
 .../gdb-tests/tests/test_command_distance.py  | 26 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/pwndbg/commands/distance.py b/pwndbg/commands/distance.py
index 6e8f22123..abba44415 100644
--- a/pwndbg/commands/distance.py
+++ b/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)
diff --git a/tests/gdb-tests/tests/test_command_distance.py b/tests/gdb-tests/tests/test_command_distance.py
index aee841eb9..af1160b0d 100644
--- a/tests/gdb-tests/tests/test_command_distance.py
+++ b/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"