From e1f06e5d0f24bbbd3808b799800cff68dfae496e Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Thu, 9 Feb 2017 16:41:07 -0500
Subject: [PATCH] Skip upload copy if we know the length

If the source length is known, we can skip copying the file.
---
 communicator/ssh/communicator.go | 80 ++++++++++++++++++++------------
 1 file changed, 51 insertions(+), 29 deletions(-)

diff --git a/communicator/ssh/communicator.go b/communicator/ssh/communicator.go
index 36b698a9a..491a62c3f 100644
--- a/communicator/ssh/communicator.go
+++ b/communicator/ssh/communicator.go
@@ -292,8 +292,25 @@ func (c *Communicator) Upload(path string, input io.Reader) error {
 	// which works for unix and windows
 	targetDir = filepath.ToSlash(targetDir)
 
+	// Skip copying if we can get the file size directly from common io.Readers
+	size := int64(0)
+
+	switch src := input.(type) {
+	case *os.File:
+		fi, err := src.Stat()
+		if err != nil {
+			size = fi.Size()
+		}
+	case *bytes.Buffer:
+		size = int64(src.Len())
+	case *bytes.Reader:
+		size = int64(src.Len())
+	case *strings.Reader:
+		size = int64(src.Len())
+	}
+
 	scpFunc := func(w io.Writer, stdoutR *bufio.Reader) error {
-		return scpUploadFile(targetFile, input, w, stdoutR)
+		return scpUploadFile(targetFile, input, w, stdoutR, size)
 	}
 
 	return c.scpSession("scp -vt "+targetDir, scpFunc)
@@ -490,45 +507,50 @@ func checkSCPStatus(r *bufio.Reader) error {
 	return nil
 }
 
-func scpUploadFile(dst string, src io.Reader, w io.Writer, r *bufio.Reader) error {
-	// Create a temporary file where we can copy the contents of the src
-	// so that we can determine the length, since SCP is length-prefixed.
-	tf, err := ioutil.TempFile("", "terraform-upload")
-	if err != nil {
-		return fmt.Errorf("Error creating temporary file for upload: %s", err)
-	}
-	defer os.Remove(tf.Name())
-	defer tf.Close()
+func scpUploadFile(dst string, src io.Reader, w io.Writer, r *bufio.Reader, size int64) error {
+	if size == 0 {
+		// Create a temporary file where we can copy the contents of the src
+		// so that we can determine the length, since SCP is length-prefixed.
+		tf, err := ioutil.TempFile("", "terraform-upload")
+		if err != nil {
+			return fmt.Errorf("Error creating temporary file for upload: %s", err)
+		}
+		defer os.Remove(tf.Name())
+		defer tf.Close()
 
-	log.Println("Copying input data into temporary file so we can read the length")
-	if _, err := io.Copy(tf, src); err != nil {
-		return err
-	}
+		log.Println("Copying input data into temporary file so we can read the length")
+		if _, err := io.Copy(tf, src); err != nil {
+			return err
+		}
 
-	// Sync the file so that the contents are definitely on disk, then
-	// read the length of it.
-	if err := tf.Sync(); err != nil {
-		return fmt.Errorf("Error creating temporary file for upload: %s", err)
-	}
+		// Sync the file so that the contents are definitely on disk, then
+		// read the length of it.
+		if err := tf.Sync(); err != nil {
+			return fmt.Errorf("Error creating temporary file for upload: %s", err)
+		}
 
-	// Seek the file to the beginning so we can re-read all of it
-	if _, err := tf.Seek(0, 0); err != nil {
-		return fmt.Errorf("Error creating temporary file for upload: %s", err)
-	}
+		// Seek the file to the beginning so we can re-read all of it
+		if _, err := tf.Seek(0, 0); err != nil {
+			return fmt.Errorf("Error creating temporary file for upload: %s", err)
+		}
 
-	fi, err := tf.Stat()
-	if err != nil {
-		return fmt.Errorf("Error creating temporary file for upload: %s", err)
+		fi, err := tf.Stat()
+		if err != nil {
+			return fmt.Errorf("Error creating temporary file for upload: %s", err)
+		}
+
+		src = tf
+		size = fi.Size()
 	}
 
 	// Start the protocol
 	log.Println("Beginning file upload...")
-	fmt.Fprintln(w, "C0644", fi.Size(), dst)
+	fmt.Fprintln(w, "C0644", size, dst)
 	if err := checkSCPStatus(r); err != nil {
 		return err
 	}
 
-	if _, err := io.Copy(w, tf); err != nil {
+	if _, err := io.Copy(w, src); err != nil {
 		return err
 	}
 
@@ -592,7 +614,7 @@ func scpUploadDir(root string, fs []os.FileInfo, w io.Writer, r *bufio.Reader) e
 
 			err = func() error {
 				defer f.Close()
-				return scpUploadFile(fi.Name(), f, w, r)
+				return scpUploadFile(fi.Name(), f, w, r, fi.Size())
 			}()
 
 			if err != nil {