diff --git a/Makefile b/Makefile
index afed320..779c733 100644
--- a/Makefile
+++ b/Makefile
@@ -1,2 +1,2 @@
 debug:
-	gcc main.c -g -O0 -o build/trace -Wall -Wextra -Wno-unused-variable -Wno-unused-function -Wno-switch
+	gcc main.c -g -O0 -o build/trace -Wall -Wextra -Wno-unused-variable -Wno-unused-function -Wno-switch  # -fsanitize=address
diff --git a/command.c b/command.c
index 094db31..db02d17 100644
--- a/command.c
+++ b/command.c
@@ -32,8 +32,8 @@ command_start(struct mi_process proc)
     u64 main_address = proc.base_address + proc.main_address;
     
     long saved_instruction = ptrace(PTRACE_PEEKDATA, proc.pid, main_address, NULL);
-    long int3_byte = 0x000000cc;
-    long int3_word = (saved_instruction & 0xFFFFFF00) | int3_byte;
+    long int3_byte = 0x000000000000cc;
+    long int3_word = (saved_instruction & 0xFFFFFFFFFFFFFF00) | int3_byte;
     
     /* write 0xcc */
     ptrace(PTRACE_POKEDATA, proc.pid, main_address, int3_word);
@@ -66,6 +66,7 @@ command_next(struct mi_process proc)
         //printf("%#lx\n", regs.rip - proc.base_address);
         next_sp = pc_to_sourcepoint(proc, regs.rip - proc.base_address);
         // TODO: skip call, repXXX, etc
+        if (!next_sp) break;
     } while (next_sp->line == sp->line);
     
     print_sourcepoint(proc, next_sp);
@@ -83,7 +84,6 @@ static void
 command_cont(struct mi_process proc)
 {
     ptrace(PTRACE_CONT, proc.pid, 0, 0);
-    // TODO
 }
 
 static void
diff --git a/main.c b/main.c
index 1e1c849..3ab2a48 100644
--- a/main.c
+++ b/main.c
@@ -4,8 +4,6 @@
 #include "util.c"
 #include "command.c"
 
-// TODO: read the directory table, get full path as comp_dir + directory_table (can be ".." or w/e) + filename
-
 int
 main(int argc, char *argv[])
 {
@@ -34,9 +32,9 @@ main(int argc, char *argv[])
     fflush(stdout);
     
     parse_debug_line(process.elf, 0, 0, 0, &process.debug.sp_count, &process.source_file_count, &process.source_dirs_count);
-    process.debug.sp_table = malloc(process.debug.sp_count * sizeof(struct mi_sourcepoint));
-    process.source_directories = malloc(process.source_dirs_count * sizeof(char *));
-    process.source_files = malloc(process.source_file_count * sizeof(struct mi_sourcefile));
+    process.debug.sp_table = calloc(1, process.debug.sp_count * sizeof(struct mi_sourcepoint));
+    process.source_directories = calloc(1, process.source_dirs_count * sizeof(char *));
+    process.source_files = calloc(1, process.source_file_count * sizeof(struct mi_sourcefile));
     parse_debug_line(process.elf, process.debug.sp_table, process.source_files, process.source_directories, 0, 0, 0);
     
     while ((command_length = getline(&command, &max_command_length, stdin))) {
@@ -45,6 +43,8 @@ main(int argc, char *argv[])
             memcpy(command, last_command, command_length);
         }
         
+        int wait = 0;
+        
         if (0 == strncmp(command, "exit\n", command_length) || 0 == strncmp(command, "q\n", command_length)) {
             break;
         } else if (0 == strncmp(command, "regs\n", command_length)) {
@@ -59,12 +59,43 @@ main(int argc, char *argv[])
             command_list(process);
         } else if (0 == strncmp(command, "cont\n", command_length)) {
             command_cont(process);
+            wait = 1;
         } else if (0 == strncmp(command, "stop\n", command_length)) {
             command_stop(process);
+            wait = 1;
         } else {
             printf("Unknown command: %*s", command_length, command);
         }
         
+        if (wait) {
+            int wstatus;
+            int cpid = waitpid(process.pid, &wstatus, 0);
+            
+            if (cpid != 0) {
+                if (WIFEXITED(wstatus)) {
+                    int status = WEXITSTATUS(wstatus);
+                    printf("Child terminated normally with status %d\n", status);
+                    break;
+                }
+                
+                if (WIFSIGNALED(wstatus)) {
+                    int signal = WTERMSIG(wstatus);
+                    printf("Child terminated by signal %d (%s)\n", signal, strsignal(signal));
+                    break;
+                }
+                
+                if (WIFSTOPPED(wstatus)) {
+                    int signal = WSTOPSIG(wstatus);
+                    printf("Child stopped by signal %d (%s)\n", signal, strsignal(signal));
+                    break;
+                }
+                
+                if (WIFCONTINUED(wstatus)) {
+                    printf("Child continued by SIGCONT\n");
+                }
+            }
+        }
+        
         memset(last_command, 0, max_command_length);
         memcpy(last_command, command, command_length);
         
diff --git a/util.c b/util.c
index 690c675..f05dc16 100644
--- a/util.c
+++ b/util.c
@@ -58,6 +58,11 @@ read_file_mmap(char *path)
 static void
 print_sourcepoint(struct mi_process proc, struct mi_sourcepoint *sp)
 {
+    if (!sp) {
+        printf("Unknown location\n");
+        return;
+    }
+    
     // NOTE: sourcepoint file indices are 1-based
     if (proc.source_files[sp->file - 1].file.data == 0) {
         char *file_filename = proc.source_files[sp->file - 1].filename;
@@ -168,8 +173,18 @@ get_process_registers(struct mi_process proc)
 static void
 set_process_registers(struct mi_process proc, struct mi_registers regs)
 {
+    regs._sys.rax = regs.rax;
+    regs._sys.rip = regs.rip;
+    regs._sys.rbp = regs.rbp;
+    regs._sys.rbx = regs.rbx;
+    regs._sys.rcx = regs.rcx;
+    regs._sys.rdx = regs.rdx;
+    regs._sys.rsi = regs.rsi;
+    regs._sys.rdi = regs.rdi;
+    regs._sys.rsp = regs.rsp;
+    
     struct iovec iov = { &regs._sys, sizeof(regs._sys) };
-    ptrace(PTRACE_GETREGSET, proc.pid, NT_PRSTATUS, &iov);
+    ptrace(PTRACE_SETREGSET, proc.pid, NT_PRSTATUS, &iov);
 }
 
 static void