issue #137 - added tests for heap set-arena, heap chunks

gef.py

@@ -5508,12 +5508,15 @@ def do_invoke(self, argv):
             err("Invalid location")
             return
 
-        new_arena = Address(to_unsigned_long(new_arena))
-        if new_arena is None or not new_arena.valid:
-            err("Invalid location")
-            return
+        if argv[0].startswith("0x"):
+            new_arena = Address(to_unsigned_long(new_arena))
+            if new_arena is None or not new_arena.valid:
+                err("Invalid location")
+                return
 
-        __gef_default_main_arena__ = "*{:s}".format(format_address(new_arena.value))
+            __gef_default_main_arena__ = "*{:s}".format(format_address(new_arena.value))
+        else:
+            __gef_default_main_arena__ = argv[0]
         return
 
 @register_command

tests/test-runner.py

@@ -140,6 +140,17 @@ def test_command_heap_arenas(self):
         self.assertTrue(b"Arena (base=" in res)
         return
 
+
+    def test_command_heap_set_arena(self):
+        cmd = "heap set-arena main_arena"
+        target = "tests/binaries/heap.out"
+        self.assertFailIfInactiveSession(gdb_run_command(cmd, target=target))
+        res = gdb_run_silent_command(cmd, target=target, after=["heap arenas",])
+        self.assertNoException(res)
+        self.assertTrue(b"Arena (base=" in res)
+        return
+
+
     def test_command_heap_chunk(self):
         cmd = "heap chunk p1"
         target = "tests/binaries/heap.out"
@@ -149,6 +160,15 @@ def test_command_heap_chunk(self):
         self.assertTrue(b"NON_MAIN_ARENA flag: " in res)
         return
 
+    def test_command_heap_chunks(self):
+        cmd = "heap chunks"
+        target = "tests/binaries/heap.out"
+        self.assertFailIfInactiveSession(gdb_run_command(cmd, target=target))
+        res = gdb_run_silent_command(cmd, target=target)
+        self.assertNoException(res)
+        self.assertTrue(b"Chunk(addr=" in res and b"top chunk" in res)
+        return
+
     def test_command_heap_bins_fast(self):
         cmd = "heap bins fast"
         target = "tests/binaries/heap-fastbins.out"