propagate stdout and stderr streams properly to hoop exec cmd (#95)

hoophq · Jan 4, 2023 · 1c623e2 · 1c623e2
1 parent 0cd0dc2
commit 1c623e2
Show file tree

Hide file tree

Showing 7 changed files with 48 additions and 22 deletions.
diff --git a/agent/terminal-exec.go b/agent/terminal-exec.go
@@ -36,17 +36,20 @@ func (a *Agent) doExec(pkt *pb.Packet) {
 	}
 	log.Printf("session=%v, tty=false - executing command=%q", string(sessionID), cmd.String())
 
-	spec := map[string][]byte{pb.SpecGatewaySessionID: sessionID}
-	stdoutWriter := pb.NewStreamWriter(a.client, pb.PacketClientAgentExecOKType, spec)
+	stdoutw := pb.NewStdoutStreamWriter(a.client, pb.PacketClientAgentExecOKType,
+		map[string][]byte{pb.SpecGatewaySessionID: sessionID})
+	stderrw := pb.NewStderrStreamWriter(a.client, pb.PacketClientAgentExecOKType,
+		map[string][]byte{pb.SpecGatewaySessionID: sessionID})
 
 	onExecErr := func(exitCode int, errMsg string, v ...any) {
 		errMsg = fmt.Sprintf(errMsg, v...)
-		spec[pb.SpecClientExecExitCodeKey] = []byte(strconv.Itoa(exitCode))
-		_, _ = pb.NewStreamWriter(a.client, packetErrType, spec).
+		_, _ = pb.NewStreamWriter(a.client, packetErrType, map[string][]byte{
+			pb.SpecGatewaySessionID:      sessionID,
+			pb.SpecClientExecExitCodeKey: []byte(strconv.Itoa(exitCode))}).
 			Write([]byte(errMsg))
 	}
 
-	if err = cmd.Run(stdoutWriter, pkt.Payload, onExecErr); err != nil {
+	if err = cmd.Run(stdoutw, stderrw, pkt.Payload, onExecErr); err != nil {
 		log.Printf("session=%v - err=%v", string(sessionID), err)
 	}
 }
diff --git a/agent/terminal/server-exec.go b/agent/terminal/server-exec.go
@@ -76,7 +76,7 @@ func (c *Command) OnPostExec() error {
 	return nil
 }
 
-func (c *Command) Run(streamWriter io.WriteCloser, stdinInput []byte, onExecErr OnExecErrFn, clientArgs ...string) error {
+func (c *Command) Run(stdoutw, stderrw io.WriteCloser, stdinInput []byte, onExecErr OnExecErrFn, clientArgs ...string) error {
 	pipeStdout, err := c.cmd.StdoutPipe()
 	if err != nil {
 		onExecErr(term.InternalErrorExitCode, "internal error, failed returning stdout pipe")
@@ -106,8 +106,8 @@ func (c *Command) Run(streamWriter io.WriteCloser, stdinInput []byte, onExecErr
 		onExecErr(term.InternalErrorExitCode, "internal error, failed writing input")
 		return err
 	}
-	stdoutCh := copyBuffer(streamWriter, pipeStdout, 1024, "stdout")
-	stderrCh := copyBuffer(streamWriter, pipeStderr, 1024, "stderr")
+	stdoutCh := copyBuffer(stdoutw, pipeStdout, 1024, "stdout")
+	stderrCh := copyBuffer(stderrw, pipeStderr, 1024, "stderr")
 
 	go func() {
 		exitCode = 0

diff --git a/client/cmd/connect.go b/client/cmd/connect.go
@@ -339,9 +339,20 @@ func (c *connect) processPacket(pkt *pb.Packet, config *Config, loader *spinner.
 		}
 		c.processGracefulExit(errors.New("user cancelled the action"))
 	case pb.PacketClientGatewayExecRejectType:
-		c.processGracefulExit(errors.New("task rejected. Sorry"))
+		c.processGracefulExit(errors.New("task rejected"))
 	case pb.PacketClientAgentExecOKType:
-		c.printOutputAndExit(pkt.Payload)
+		loader.Stop() // disables after the first packet arrives
+		stdStreamVal := pkt.Spec[pb.SpecServerExecStdStreamKey]
+		switch string(stdStreamVal) {
+		case pb.StdOut: // stdout
+			os.Stdout.Write(pkt.Payload)
+		case pb.StdErr: // stderr
+			os.Stderr.Write(pkt.Payload)
+		default: // it shouldn't happen
+			if len(pkt.Payload) > 0 {
+				os.Stderr.Write(pkt.Payload)
+			}
+		}
 	case pb.PacketClientAgentExecErrType:
 		if len(pkt.Payload) > 0 {
 			_, _ = os.Stderr.Write([]byte(styles.ClientError(string(pkt.Payload)) + "\n"))
@@ -436,13 +447,8 @@ func (c *connect) printHeader(sessionID string) {
 }
 
 func (c *connect) printErrorAndExit(format string, v ...any) {
-	c.loader.Disable()
+	c.loader.Stop()
 	errOutput := styles.ClientError(fmt.Sprintf(format, v...))
 	fmt.Println(errOutput)
 	os.Exit(1)
 }
-
-func (c *connect) printOutputAndExit(output []byte) {
-	c.loader.Disable()
-	fmt.Print(string(output))
-}
diff --git a/client/cmd/exec.go b/client/cmd/exec.go
@@ -79,11 +79,10 @@ func parseExecInput(c *connect) []byte {
 
 func runExec(args []string) {
 	config := getClientConfig()
-
-	loader := spinner.New(spinner.CharSets[78], 70*time.Millisecond)
+	loader := spinner.New(spinner.CharSets[78], 70*time.Millisecond, spinner.WithWriter(os.Stderr))
 	loader.Color("green")
+	loader.Suffix = " running ..."
 	loader.Start()
-	loader.Suffix = " executing input ..."
 
 	c := newClientConnect(config, loader, args, pb.ClientVerbExec)
 	pkt := &pb.Packet{
@@ -99,8 +98,6 @@ func runExec(args []string) {
 		c.printErrorAndExit("failed executing command, err=%v", err)
 	}
 
-	loader.Stop()
-
 	for {
 		pkt, err := c.client.Recv()
 		c.processGracefulExit(err)

diff --git a/client/go.mod b/client/go.mod
@@ -3,7 +3,9 @@ module github.com/runopsio/hoop/client
 go 1.19
 
 require (
-	github.com/briandowns/spinner v1.19.0
+	// latest version breaks when using the loader to stderr
+	// update to latest version after this https://github.com/briandowns/spinner/pull/136
+	github.com/briandowns/spinner v1.18.0
 	github.com/creack/pty v1.1.18
 	github.com/runopsio/hoop/common v0.0.0-00010101000000-000000000000
 	github.com/spf13/cobra v1.5.0

diff --git a/common/proto/const.go b/common/proto/const.go
@@ -50,6 +50,7 @@ const (
 	SpecClientConnectionID        string = "client.connection_id"
 	SpecClientExecExitCodeKey     string = "terminal.exit_code"
 	SpecClientExecArgsKey         string = "terminal.args"
+	SpecServerExecStdStreamKey    string = "terminal.std_stream"
 	SpecAgentConnectionParamsKey  string = "agent.connection_params"
 	SpecAgentGCPRawCredentialsKey string = "agent.gcp_credentials"
 	SpecTCPServerConnectKey       string = "tcp.server_connect"
@@ -72,6 +73,9 @@ const (
 	ConnectionOriginAgent  = "agent"
 	ConnectionOriginClient = "client"
 
+	StdOut = "stdout"
+	StdErr = "stderr"
+
 	ClientLoginCallbackAddress string = "127.0.0.1:3587"
 
 	ClientVerbConnect = "connect"

diff --git a/common/proto/types.go b/common/proto/types.go
@@ -80,6 +80,20 @@ func NewStreamWriter(client ClientTransport, pktType PacketType, spec map[string
 	return &streamWriter{client: client, packetType: pktType, packetSpec: spec}
 }
 
+func NewStdoutStreamWriter(client ClientTransport, pktType PacketType, spec map[string][]byte) io.WriteCloser {
+	if spec != nil {
+		spec[SpecServerExecStdStreamKey] = []byte(StdOut)
+	}
+	return &streamWriter{client: client, packetType: pktType, packetSpec: spec}
+}
+
+func NewStderrStreamWriter(client ClientTransport, pktType PacketType, spec map[string][]byte) io.WriteCloser {
+	if spec != nil {
+		spec[SpecServerExecStdStreamKey] = []byte(StdErr)
+	}
+	return &streamWriter{client: client, packetType: pktType, packetSpec: spec}
+}
+
 func NewHookStreamWriter(
 	client ClientTransport,
 	pktType PacketType,