Launch prog from cli, also exit cleanly

_fixtures/integrationprog.go

+package main
+
+import (
+	"fmt"
+	"time"
+)
+
+func sayhi() {
+	fmt.Println("hi")
+}
+
+func main() {
+	time.Sleep(1 * time.Second)
+	for i := 0; i < 3; i++ {
+		sayhi()
+		time.Sleep(1 * time.Second)
+	}
+}

command/command.go

@@ -5,11 +5,9 @@ package command
 import (
 	"debug/gosym"
 	"fmt"
-	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
-	"syscall"
 
 	"github.com/derekparker/dbg/proctl"
 )
@@ -23,7 +21,6 @@ type Commands struct {
 // Returns a Commands struct with default commands defined.
 func DebugCommands() *Commands {
 	cmds := map[string]cmdfunc{
-		"exit":     exitFunc,
 		"continue": cont,
 		"next":     next,
 		"break":    breakpoint,
@@ -66,16 +63,6 @@ func noCmdAvailable(p *proctl.DebuggedProcess, ars ...string) error {
 	return fmt.Errorf("command not available")
 }
 
-func exitFunc(p *proctl.DebuggedProcess, ars ...string) error {
-	err := syscall.PtraceDetach(p.Pid)
-	if err != nil {
-		return err
-	}
-
-	os.Exit(0)
-	return nil
-}
-
 func nullCommand(p *proctl.DebuggedProcess, ars ...string) error {
 	return nil
 }

main.go

@@ -2,11 +2,13 @@ package main
 
 import (
 	"bufio"
+	"flag"
 	"fmt"
 	"os"
+	"os/exec"
 	"runtime"
-	"strconv"
 	"strings"
+	"syscall"
 
 	"github.com/derekparker/dbg/command"
 	"github.com/derekparker/dbg/proctl"
@@ -23,32 +25,35 @@ func main() {
 	runtime.LockOSThread()
 
 	var (
+		pid  int
+		proc string
 		t    = newTerm()
 		cmds = command.DebugCommands()
 	)
 
-	if len(os.Args) == 1 {
-		die("You must provide a pid\n")
-	}
+	flag.IntVar(&pid, "pid", 0, "Pid of running process to attach to.")
+	flag.StringVar(&proc, "proc", "", "Path to process to run and debug.")
+	flag.Parse()
 
-	pid, err := strconv.Atoi(os.Args[1])
-	if err != nil {
-		die(err)
+	if flag.NFlag() == 0 {
+		flag.Usage()
+		os.Exit(0)
 	}
 
-	dbgproc, err := proctl.NewDebugProcess(pid)
-	if err != nil {
-		die("Could not start debugging process:", err)
-	}
+	dbgproc := beginTrace(pid, proc)
 
 	for {
 		cmdstr, err := t.promptForInput()
 		if err != nil {
-			die("Prompt for input failed.\n")
+			die(1, "Prompt for input failed.\n")
 		}
 
 		cmdstr, args := parseCommand(cmdstr)
 
+		if cmdstr == "exit" {
+			handleExit(t, dbgproc, 0)
+		}
+
 		cmd := cmds.Find(cmdstr)
 		err = cmd(dbgproc, args...)
 		if err != nil {
@@ -57,9 +62,66 @@ func main() {
 	}
 }
 
-func die(args ...interface{}) {
+func beginTrace(pid int, proc string) *proctl.DebuggedProcess {
+	var (
+		err     error
+		dbgproc *proctl.DebuggedProcess
+	)
+
+	if pid != 0 {
+		dbgproc, err = proctl.NewDebugProcess(pid)
+		if err != nil {
+			die(1, "Could not start debugging process:", err)
+		}
+	}
+
+	if proc != "" {
+		proc := exec.Command(proc)
+		proc.Stdout = os.Stdout
+
+		err = proc.Start()
+		if err != nil {
+			die(1, "Could not start process:", err)
+		}
+
+		dbgproc, err = proctl.NewDebugProcess(proc.Process.Pid)
+		if err != nil {
+			die(1, "Could not start debugging process:", err)
+		}
+	}
+
+	return dbgproc
+}
+
+func handleExit(t *term, dbp *proctl.DebuggedProcess, status int) {
+	fmt.Println("Would you like to kill the process? [y/n]")
+	answer, err := t.stdin.ReadString('\n')
+	if err != nil {
+		die(2, err.Error())
+	}
+
+	fmt.Println("Detaching from process...")
+	err = syscall.PtraceDetach(dbp.Process.Pid)
+	if err != nil {
+		die(2, "Could not detach", err)
+	}
+
+	if answer == "y\n" {
+		fmt.Println("Killing process", dbp.Process.Pid)
+
+		err := dbp.Process.Kill()
+		if err != nil {
+			fmt.Println("Could not kill process", err)
+		}
+	}
+
+	die(status, "Hope I was of service hunting your bug!")
+}
+
+func die(status int, args ...interface{}) {
 	fmt.Fprint(os.Stderr, args)
-	os.Exit(1)
+	fmt.Fprint(os.Stderr, "\n")
+	os.Exit(status)
 }
 
 func newTerm() *term {

main_test.go

+package main
+
+import (
+	"bytes"
+	"os"
+	"os/exec"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+)
+
+func buildBinary(t *testing.T) {
+	cmd := exec.Command("go", "build", "-o", "dbg-test")
+
+	err := cmd.Run()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func startDebugger(t *testing.T, pid int) *os.Process {
+	cmd := exec.Command("sudo", "./dbg-test", "-pid", strconv.Itoa(pid))
+	cmd.Stdin = bytes.NewBufferString("exit\ny\n")
+
+	err := cmd.Start()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return cmd.Process
+}
+
+func startTestProg(t *testing.T, proc string) *os.Process {
+	cmd := exec.Command(proc)
+
+	err := cmd.Start()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return cmd.Process
+}
+
+func TestCleanExit(t *testing.T) {
+	buildBinary(t)
+
+	var (
+		waitchan = make(chan *os.ProcessState)
+		testprog = startTestProg(t, "_fixtures/livetestprog")
+	)
+
+	go func() {
+		ps, err := testprog.Wait()
+		if err != nil {
+			t.Fatal(err)
+		}
+		waitchan <- ps
+	}()
+
+	proc := startDebugger(t, testprog.Pid)
+
+	defer func() {
+		testprog.Kill()
+		proc.Kill()
+
+		err := os.Remove("dbg-test")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	timer := time.NewTimer(5 * time.Second)
+	select {
+	case ps := <-waitchan:
+		// Admittedly, this is weird.
+		// There is/was a bug in Go that marked
+		// a process as done whenever `Wait()` was
+		// called on it, I fixed that, but it seems
+		// there may be another connected bug somewhere
+		// where the process is not marked as exited.
+		if strings.Contains(ps.String(), "exited") {
+			t.Fatal("Process has not exited")
+		}
+
+	case <-timer.C:
+		t.Fatal("timeout")
+	}
+}