From 12adc274b2a25217c27ceab81bbc2d980bd17ecc Mon Sep 17 00:00:00 2001
From: Jatin Naik <jnaik@pivotal.io>
Date: Tue, 23 May 2017 10:57:58 +0100
Subject: [PATCH] Reduce duplication in ssh/connection.go

[#145693923]

Signed-off-by: Will Pragnell <wpragnell@pivotal.io>
---
 ssh/connection.go | 52 +++++++++++++++++++++--------------------------
 1 file changed, 23 insertions(+), 29 deletions(-)

diff --git a/ssh/connection.go b/ssh/connection.go
index d5e61a197..b38c3fd77 100644
--- a/ssh/connection.go
+++ b/ssh/connection.go
@@ -59,52 +59,46 @@ func (c Connection) Run(cmd string) ([]byte, []byte, int, error) {
 }
 
 func (c Connection) Stream(cmd string, writer io.Writer) ([]byte, int, error) {
-	session, err := c.connection.NewSession()
-	if err != nil {
-		return nil, 0, errors.Wrap(err, "ssh.Stream.NewSession failed")
-	}
-
 	errBuffer := bytes.NewBuffer([]byte{})
-	exitCode := 0
-
-	session.Stdout = writer
-	session.Stderr = errBuffer
-	err = session.Run(cmd)
-	if err != nil {
-		exitErr, yes := err.(*ssh.ExitError)
-		if yes {
-			exitCode = exitErr.ExitStatus()
-		} else {
-			return nil, -1, errors.Wrap(err, "ssh.Stream.Run failed")
-		}
 
-	}
+	exitCode, err := c.runInSession(cmd, writer, errBuffer, nil)
 
-	return errBuffer.Bytes(), exitCode, nil
+	return errBuffer.Bytes(), exitCode, err
 }
 
-func (c Connection) StreamStdin(cmd string, reader io.Reader) (stdout, stderr []byte, exitCode int, err error) {
-	session, err := c.connection.NewSession()
-
+func (c Connection) StreamStdin(cmd string, stdinReader io.Reader) (stdout, stderr []byte, exitCode int, err error) {
 	outBuffer := bytes.NewBuffer([]byte{})
 	errBuffer := bytes.NewBuffer([]byte{})
 
-	session.Stdin = reader
-	session.Stdout = outBuffer
-	session.Stderr = errBuffer
+	exitCode, err = c.runInSession(cmd, outBuffer, errBuffer, stdinReader)
 
-	err = session.Run(cmd)
+	return outBuffer.Bytes(), errBuffer.Bytes(), exitCode, err
+}
 
+func (c Connection) runInSession(cmd string, stdout, stderr io.Writer, stdin io.Reader) (int, error) {
+	session, err := c.connection.NewSession()
 	if err != nil {
+		return 0, errors.Wrap(err, "ssh.Stream.NewSession failed")
+	}
+
+	session.Stdin = stdin
+	session.Stdout = stdout
+	session.Stderr = stderr
+
+	var exitCode int
+
+	err = session.Run(cmd)
+	if err == nil {
+		exitCode = 0
+	} else {
 		exitErr, yes := err.(*ssh.ExitError)
 		if yes {
 			exitCode = exitErr.ExitStatus()
 		} else {
-			return nil, nil, -1, errors.Wrap(err, "ssh.StreamStdin.Run failed with non-ExitError failure")
+			return -1, errors.Wrap(err, "ssh.Stream.Run failed")
 		}
 	}
-
-	return outBuffer.Bytes(), errBuffer.Bytes(), exitCode, nil
+	return exitCode, nil
 }
 
 func (c Connection) Username() string {