Disable brk() before the snapshot so that malloc()'d blocks are restored

ptracer.cc

@@ -357,10 +357,18 @@ int main(int argc, char **argv) {
     assert(WIFSTOPPED(status));
 
     if (WSTOPSIG(status) == (SIGTRAP | kSysFlag)) {
-      if (!syscall_entry) {
-        struct user_regs_struct regs;
-        rc = ptrace(PTRACE_GETREGS, pid, 0, &regs);
-        assert(rc == 0);
+      struct user_regs_struct regs;
+      rc = ptrace(PTRACE_GETREGS, pid, 0, &regs);
+      assert(rc == 0);
+      if (syscall_entry) {
+        // Disable use of the brk() heap so that we don't have to save
+        // and restore the brk() heap pointer and heap contents.
+        if (regs.orig_rax == __NR_brk) {
+          regs.orig_rax = -1;
+          rc = ptrace(PTRACE_SETREGS, pid, 0, &regs);
+          assert(rc == 0);
+        }
+      } else {
         ptracer.HandleSyscall(&regs);
       }
       syscall_entry = !syscall_entry;

tests/save_restore_tests.cc

@@ -1,7 +1,11 @@
 
 #include <assert.h>
+#include <errno.h>
 #include <signal.h>
+#include <stdint.h>
+#include <stdlib.h>
 #include <sys/mman.h>
+#include <sys/syscall.h>
 #include <unistd.h>
 
 
@@ -38,12 +42,35 @@ int main() {
                       MAP_PRIVATE | MAP_ANON | MAP_FIXED, -1, 0);
   assert(result == addr4);
 
+  // Test that brk() is disabled before taking the snapshot.
+  long break_ptr = syscall(__NR_brk, 0);
+  assert(break_ptr == -1);
+  assert(errno == ENOSYS);
+
+  // Test malloc(), which we expect will try to use brk().
+  int num_blocks = 100;
+  char *blocks[num_blocks];
+  for (int i = 0; i < num_blocks; ++i) {
+    blocks[i] = (char *) malloc(100);
+    assert(blocks[i]);
+    *(int *) blocks[i] = i * 100;
+  }
+
   raise(SIGUSR1);
 
   assert(((char *) addr2)[0] == 'a');
   assert(((char *) addr2)[size * 2] == 'b');
 
   assert(*(int *) addr3 == 0x1234);
 
+  // Currently brk() works after resuming from the snapshot.
+  long break_ptr_after = syscall(__NR_brk, 0);
+  assert(break_ptr_after != -1);
+
+  // Check that the malloc()'d blocks still work.
+  for (int i = 0; i < num_blocks; ++i) {
+    assert(*(int *) blocks[i] == i * 100);
+  }
+
   return 0;
 }