Fix the limit in the search command (#2012)

* Fixed the limit in the search command

* Add search --limit test

* Fix lint issues

pwndbg/search.py

@@ -70,6 +70,9 @@ def search(
         start = vmmap.start
         end = vmmap.end
 
+        if limit and found_count >= limit:
+            break
+
         while True:
             # No point in searching if we can't read the memory
             if not pwndbg.gdblib.memory.peek(start):
@@ -114,7 +117,7 @@ def search(
 
             yield start
             found_count += 1
-            if limit and found_count == limit:
+            if limit and found_count >= limit:
                 break
 
             if step is not None:

tests/gdb-tests/tests/binaries/search_memory.c

@@ -7,9 +7,13 @@
 
 void break_here(void) {}
 
+size_t marker = 0xABCDEF1234567890;
+
 int main(void)
 {
     void *p;
+    size_t *heap_marker;
+    size_t local_marker;
 
     p = malloc(0x100000);
     memset(p, 0x0, 0x100000);
@@ -24,6 +28,10 @@ int main(void)
         *(unsigned int *)(p + i + 0x17) = 0xd00dbeef;
     }
 
+    heap_marker = malloc(8);
+    *heap_marker = marker;
+    local_marker = marker;
+
     break_here();
 
     return 0;

tests/gdb-tests/tests/test_command_search.py

@@ -6,11 +6,12 @@ import tests
 
 SEARCH_BINARY = tests.binaries.get("search_memory.out")
 SEARCH_PATTERN = 0xD00DBEEF
+SEARCH_PATTERN2 = 0xABCDEF1234567890
 
 
-def test_command_search_limit(start_binary):
+def test_command_search_limit_single_page(start_binary):
     """
-    Tests simple search limit
+    Tests simple search limit for single memory page
     """
     start_binary(SEARCH_BINARY)
 
@@ -33,6 +34,29 @@ def test_command_search_limit(start_binary):
     assert result_value == hex(SEARCH_PATTERN)
 
 
+def test_command_search_limit_multiple_pages(start_binary):
+    """
+    Tests simple search limit for multiple memory pages
+    """
+    start_binary(SEARCH_BINARY)
+
+    gdb.execute("break break_here")
+    gdb.execute("run")
+
+    def filter_results(line):
+        return hex(SEARCH_PATTERN2).lower() in line.lower()
+
+    total_entries = 3
+    result_str: str = gdb.execute(f"search -8 {SEARCH_PATTERN2}", to_string=True)
+    result_count = len(list(filter(filter_results, result_str.splitlines())))
+    assert result_count == total_entries
+
+    search_limit = 2
+    result_str = gdb.execute(f"search -8 {SEARCH_PATTERN2} -l {search_limit}", to_string=True)
+    result_count = len(list(filter(filter_results, result_str.splitlines())))
+    assert result_count == search_limit
+
+
 def test_command_search_alignment(start_binary):
     """
     Tests aligned search