From 0e971cde331a162b9e50b41920750dcb25a05ca3 Mon Sep 17 00:00:00 2001
From: Keith Rarick <kr@xph.us>
Date: Wed, 29 May 2013 16:06:18 +1000
Subject: [PATCH] ssh: add Output and CombinedOutput helpers

R=golang-dev, dave
CC=golang-dev
https://golang.org/cl/9711043
---
 ssh/session.go      | 46 ++++++++++++++++++++++++++---
 ssh/session_test.go | 72 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+), 4 deletions(-)

diff --git a/ssh/session.go b/ssh/session.go
index 640d322..2576024 100644
--- a/ssh/session.go
+++ b/ssh/session.go
@@ -13,6 +13,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"sync"
 )
 
 type Signal string
@@ -276,7 +277,8 @@ func (s *Session) Start(cmd string) error {
 
 // Run runs cmd on the remote host. Typically, the remote
 // server passes cmd to the shell for interpretation.
-// A Session only accepts one call to Run, Start or Shell.
+// A Session only accepts one call to Run, Start, Shell, Output,
+// or CombinedOutput.
 //
 // The returned error is nil if the command runs, has no problems
 // copying stdin, stdout, and stderr, and exits with a zero exit
@@ -293,8 +295,46 @@ func (s *Session) Run(cmd string) error {
 	return s.Wait()
 }
 
+// Output runs cmd on the remote host and returns its standard output.
+func (s *Session) Output(cmd string) ([]byte, error) {
+	if s.Stdout != nil {
+		return nil, errors.New("ssh: Stdout already set")
+	}
+	var b bytes.Buffer
+	s.Stdout = &b
+	err := s.Run(cmd)
+	return b.Bytes(), err
+}
+
+type singleWriter struct {
+	b  bytes.Buffer
+	mu sync.Mutex
+}
+
+func (w *singleWriter) Write(p []byte) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.b.Write(p)
+}
+
+// CombinedOutput runs cmd on the remote host and returns its combined
+// standard output and standard error.
+func (s *Session) CombinedOutput(cmd string) ([]byte, error) {
+	if s.Stdout != nil {
+		return nil, errors.New("ssh: Stdout already set")
+	}
+	if s.Stderr != nil {
+		return nil, errors.New("ssh: Stderr already set")
+	}
+	var b singleWriter
+	s.Stdout = &b
+	s.Stderr = &b
+	err := s.Run(cmd)
+	return b.b.Bytes(), err
+}
+
 // Shell starts a login shell on the remote host. A Session only
-// accepts one call to Run, Start or Shell.
+// accepts one call to Run, Start, Shell, Output, or CombinedOutput.
 func (s *Session) Shell() error {
 	if s.started {
 		return errors.New("ssh: session already started")
@@ -521,8 +561,6 @@ func (s *Session) StderrPipe() (io.Reader, error) {
 	return s.clientChan.stderr, nil
 }
 
-// TODO(dfc) add Output and CombinedOutput helpers
-
 // NewSession returns a new interactive session on the remote host.
 func (c *ClientConn) NewSession() (*Session, error) {
 	ch := c.newChan(c.transport)
diff --git a/ssh/session_test.go b/ssh/session_test.go
index 7453c8a..9588265 100644
--- a/ssh/session_test.go
+++ b/ssh/session_test.go
@@ -138,6 +138,54 @@ func TestSessionStdoutPipe(t *testing.T) {
 	}
 }
 
+// Test that a simple string is returned via the Output helper,
+// and that stderr is discarded.
+func TestSessionOutput(t *testing.T) {
+	conn := dial(fixedOutputHandler, t)
+	defer conn.Close()
+	session, err := conn.NewSession()
+	if err != nil {
+		t.Fatalf("Unable to request new session: %v", err)
+	}
+	defer session.Close()
+
+	buf, err := session.Output("") // cmd is ignored by fixedOutputHandler
+	if err != nil {
+		t.Error("Remote command did not exit cleanly:", err)
+	}
+	w := "this-is-stdout."
+	g := string(buf)
+	if g != w {
+		t.Error("Remote command did not return expected string:")
+		t.Logf("want %q", w)
+		t.Logf("got  %q", g)
+	}
+}
+
+// Test that both stdout and stderr are returned
+// via the CombinedOutput helper.
+func TestSessionCombinedOutput(t *testing.T) {
+	conn := dial(fixedOutputHandler, t)
+	defer conn.Close()
+	session, err := conn.NewSession()
+	if err != nil {
+		t.Fatalf("Unable to request new session: %v", err)
+	}
+	defer session.Close()
+
+	buf, err := session.CombinedOutput("") // cmd is ignored by fixedOutputHandler
+	if err != nil {
+		t.Error("Remote command did not exit cleanly:", err)
+	}
+	w := "this-is-stdout.this-is-stderr."
+	g := string(buf)
+	if g != w {
+		t.Error("Remote command did not return expected string:")
+		t.Logf("want %q", w)
+		t.Logf("got  %q", g)
+	}
+}
+
 // Test non-0 exit status is returned correctly.
 func TestExitStatusNonZero(t *testing.T) {
 	conn := dial(exitStatusNonZeroHandler, t)
@@ -586,6 +634,30 @@ func shellHandler(ch *serverChan, t *testing.T) {
 	sendStatus(0, ch, t)
 }
 
+// Ignores the command, writes fixed strings to stderr and stdout.
+// Strings are "this-is-stdout." and "this-is-stderr.".
+func fixedOutputHandler(ch *serverChan, t *testing.T) {
+	defer ch.Close()
+
+	_, err := ch.Read(make([]byte, 0))
+	if _, ok := err.(ChannelRequest); !ok {
+		t.Fatalf("error: expected channel request, got: %#v", err)
+		return
+	}
+	// ignore request, always send some text
+	ch.AckRequest(true)
+
+	_, err = io.WriteString(ch, "this-is-stdout.")
+	if err != nil {
+		t.Fatalf("error writing on server: %v", err)
+	}
+	_, err = io.WriteString(ch.Stderr(), "this-is-stderr.")
+	if err != nil {
+		t.Fatalf("error writing on server: %v", err)
+	}
+	sendStatus(0, ch, t)
+}
+
 func readLine(shell *ServerTerminal, t *testing.T) {
 	if _, err := shell.ReadLine(); err != nil && err != io.EOF {
 		t.Errorf("unable to read line: %v", err)