fix(SSHChannelReader): Read output completely before checking command exit.

fleetctl/journal.go

@@ -54,21 +54,17 @@ func journalAction(c *cli.Context) {
 	}
 
 	// check if the job is running on this machine
-	var channel *ssh.Channel
-	var err error
 	if machine.IsLocalMachineState(js.MachineState) {
-		channel = runLocalCommand(cmd)
+		runLocalCommand(cmd)
 	} else {
-		channel, err = runRemoteCommand(cmd, js.MachineState.PublicIP)
+		err := runRemoteCommand(cmd, js.MachineState.PublicIP)
 		if err != nil {
 			log.Fatalf("Unable to run command over SSH: %v", err)
 		}
 	}
-
-	readSSHChannel(channel)
 }
 
-func runLocalCommand(cmd string) *ssh.Channel {
+func runLocalCommand(cmd string) {
 	cmdSlice := strings.Split(cmd, " ")
 	osCmd := exec.Command(cmdSlice[0], cmdSlice[1:]...)
 	stdout, _ := osCmd.StdoutPipe()
@@ -86,10 +82,10 @@ func runLocalCommand(cmd string) *ssh.Channel {
 		channel.Exit <- err
 	}()
 
-	return channel
+	readSSHChannel(channel)
 }
 
-func runRemoteCommand(cmd string, ip string) (*ssh.Channel, error) {
+func runRemoteCommand(cmd string, ip string) error {
 	addr := fmt.Sprintf("%s:22", ip)
 
 	var sshClient *ssh.SSHForwardingClient
@@ -100,14 +96,16 @@ func runRemoteCommand(cmd string, ip string) (*ssh.Channel, error) {
 		sshClient, err = ssh.NewSSHClient("core", addr, getChecker(), false)
 	}
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	defer sshClient.Close()
 
 	channel, err := ssh.Execute(sshClient, cmd)
 	if err != nil {
-		return nil, err
+		return err
 	}
-	return channel, nil
+
+	readSSHChannel(channel)
+	return nil
 }

fleetctl/ssh.go

@@ -1,9 +1,9 @@
 package main
 
 import (
+	"bufio"
 	"errors"
 	"fmt"
-	"io"
 	"log"
 	"os"
 	"os/exec"
@@ -159,42 +159,34 @@ func findAddressInRunningUnits(lookup string) (string, bool) {
 }
 
 func readSSHChannel(channel *ssh.Channel) {
+	readSSHChannelOutput(channel.Stdout)
+
+	exitErr := <-channel.Exit
+	if exitErr == nil {
+		return
+	}
+
+	readSSHChannelOutput(channel.Stderr)
+
+	exitStatus := -1
+	switch exitError := exitErr.(type) {
+	case *gossh.ExitError:
+		exitStatus = exitError.ExitStatus()
+	case *exec.ExitError:
+		status := exitError.Sys().(syscall.WaitStatus)
+		exitStatus = status.ExitStatus()
+	}
+
+	os.Exit(exitStatus)
+}
+
+func readSSHChannelOutput(o *bufio.Reader) {
 	for {
-		select {
-		case exitErr := <-channel.Exit:
-			if exitErr == nil {
-				return
-			}
-
-			for {
-				bytes, err := channel.Stderr.ReadBytes('\n')
-				if err != nil {
-					break
-				}
-
-				fmt.Print(string(bytes))
-			}
-
-			exitStatus := -1
-			switch exitError := exitErr.(type) {
-			case *gossh.ExitError:
-				exitStatus = exitError.ExitStatus()
-			case *exec.ExitError:
-				status := exitError.Sys().(syscall.WaitStatus)
-				exitStatus = status.ExitStatus()
-			}
-
-			os.Exit(exitStatus)
-		default:
-			bytes, err := channel.Stdout.ReadBytes('\n')
-			if err != nil {
-				if err != io.EOF {
-					return // unknown error reading stdout, finishing.
-				}
-				break // let the exit channel to finish properly.
-			}
-
-			fmt.Print(string(bytes))
+		bytes, err := o.ReadBytes('\n')
+		if err != nil {
+			break
 		}
+
+		fmt.Print(string(bytes))
 	}
 }