Merge pull request #8 from bugout-dev/fix-trap-interrupt

bugout trap now respects SIGINT

Author: zomglings

cmd/bugout/trap/command.go

@@ -3,6 +3,7 @@ package trapcmd
 import (
 	"errors"
 	"fmt"
+	"os"
 
 	"github.com/spf13/cobra"
 
@@ -58,6 +59,10 @@ Specify the wrapped command using "--" followed by the command:
 			}
 
 			cmd.ErrOrStderr().Write([]byte(fmt.Sprintf("\n\nBugout entry created at: %s/journals/%s/%s\n", cmdutils.BugoutURL(), journalID, response.Id)))
+
+			if result.ExitCode > 0 {
+				os.Exit(result.ExitCode)
+			}
 			return nil
 		},
 	}

cmd/bugout/trap/invoke.go

@@ -2,9 +2,10 @@ package trapcmd
 
 import (
 	"io"
+	"os"
 	"os/exec"
+	"os/signal"
 	"strings"
-	"sync"
 
 	"github.com/spf13/cobra"
 )
@@ -15,8 +16,35 @@ type InvocationResult struct {
 	Stderr   string
 }
 
+func stream(reader io.Reader, writer io.Writer, doneChan chan<- bool, cancelChan <-chan bool) {
+	b := make([]byte, 1)
+	for {
+		select {
+		case <-cancelChan:
+			doneChan <- true
+			return
+		default:
+			inN, inErr := reader.Read(b)
+			if inN > 0 {
+				_, outErr := writer.Write(b)
+				if outErr != nil {
+					doneChan <- false
+					return
+				}
+			}
+			if inErr == io.EOF {
+				doneChan <- true
+				return
+			}
+			if inErr != nil {
+				doneChan <- false
+				return
+			}
+		}
+	}
+}
+
 func RunWrappedCommand(trapCmd *cobra.Command, invocation []string) (*InvocationResult, error) {
-	var wg sync.WaitGroup
 	cmd := exec.Command(invocation[0], invocation[1:]...)
 	cmd.Stdin = trapCmd.InOrStdin()
 
@@ -33,15 +61,45 @@ func RunWrappedCommand(trapCmd *cobra.Command, invocation []string) (*Invocation
 	stdoutReader := io.TeeReader(outReader, &outBuilder)
 	stderrReader := io.TeeReader(errReader, &errBuilder)
 
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		io.Copy(trapCmd.OutOrStdout(), stdoutReader)
-	}()
-	wg.Add(1)
+	outDoneChannel := make(chan bool, 1)
+	outCancelChannel := make(chan bool, 1)
+	errDoneChannel := make(chan bool, 1)
+	errCancelChannel := make(chan bool, 1)
+
+	go stream(stdoutReader, trapCmd.OutOrStdout(), outDoneChannel, outCancelChannel)
+	go stream(stderrReader, trapCmd.ErrOrStderr(), errDoneChannel, errCancelChannel)
+
+	signalsChannel := make(chan os.Signal, 1)
+	exitChannel := make(chan int, 1)
+	signal.Notify(signalsChannel, os.Interrupt, os.Kill)
+
 	go func() {
-		defer wg.Done()
-		io.Copy(trapCmd.ErrOrStderr(), stderrReader)
+		success := true
+		completed := 0
+		for {
+			select {
+			case outSuccess := <-outDoneChannel:
+				success = success && outSuccess
+				completed += 1
+			case errSuccess := <-errDoneChannel:
+				success = success && errSuccess
+				completed += 1
+			case <-signalsChannel:
+				outCancelChannel <- true
+				errCancelChannel <- true
+				exitChannel <- 130
+				return
+			default:
+				if completed == 2 {
+					if success {
+						exitChannel <- 0
+					} else {
+						exitChannel <- 1
+					}
+					return
+				}
+			}
+		}
 	}()
 
 	exitCode := 0
@@ -58,6 +116,11 @@ func RunWrappedCommand(trapCmd *cobra.Command, invocation []string) (*Invocation
 			err = runErr
 		}
 	}
-	wg.Wait()
+
+	coordinatorExitCode := <-exitChannel
+	if coordinatorExitCode != 0 {
+		exitCode = coordinatorExitCode
+	}
+
 	return &InvocationResult{ExitCode: exitCode, Stdout: outBuilder.String(), Stderr: errBuilder.String()}, err
 }

pkg/version.go

@@ -1,3 +1,3 @@
 package bugout
 
-const Version string = "0.3.0"
+const Version string = "0.3.1"