Merge lp:~axwalk/juju-core/sshstorage-stream-input into lp:~go-bot/juju-core/trunk

[Andrew Wilkins]

environs/sshstorage: stream stdin

Change sshstorage's Put to stream to the
ssh command's stdin, rather than reading
into memory and then writing.

Also, drive-by: use `mktemp --tmpdir` rather
than exporting $TMPDIR.

https://codereview.appspot.com/14197043/

environs/sshstorage: stream stdin

Change sshstorage's Put to stream to the
ssh command's stdin, rather than reading
into memory and then writing.

Also, drive-by: use `mktemp --tmpdir` rather
than exporting $TMPDIR.

https://codereview.appspot.com/14197043/

[Andrew Wilkins (axwalk)]

Reviewers: mp+188508_code.launchpad.net,

Message:
Please take a look.

Description:
environs/sshstorage: stream stdin

Change sshstorage's Put to stream to the
ssh command's stdin, rather than reading
into memory and then writing.

Also, drive-by: use `mktemp --tmpdir` rather
than exporting $TMPDIR.

https://code.launchpad.net/~axwalk/juju-core/sshstorage-stream-input/+merge/188508

(do not edit description out of merge proposal)

Please review this at https://codereview.appspot.com/14197043/

Affected files (+15, -16 lines):
   A [revision details]
   M environs/sshstorage/storage.go

Index: [revision details]
=== added file '[revision details]'
--- [revision details] 2012-01-01 00:00:00 +0000
+++ [revision details] 2012-01-01 00:00:00 +0000
@@ -0,0 +1,2 @@
+Old revision: tarmac-20130930183118-qlwtuzv7j6d2iy73
+New revision: <email address hidden>

Index: environs/sshstorage/storage.go
=== modified file 'environs/sshstorage/storage.go'
--- environs/sshstorage/storage.go 2013-09-24 01:55:30 +0000
+++ environs/sshstorage/storage.go 2013-10-01 04:04:41 +0000
@@ -129,10 +129,10 @@

  func (s *SSHStorage) runf(flockmode flockmode, command string,
args ...interface{}) (string, error) {
   command = fmt.Sprintf(command, args...)
- return s.run(flockmode, command, nil)
+ return s.run(flockmode, command, nil, 0)
  }

-func (s *SSHStorage) run(flockmode flockmode, command string, input
[]byte) (string, error) {
+func (s *SSHStorage) run(flockmode flockmode, command string, input
io.Reader, inputlen int64) (string, error) {
   const rcPrefix = "JUJU-RC: "
   command = fmt.Sprintf(
    "SHELL=/bin/bash flock %s %s -c %s",
@@ -140,23 +140,24 @@
    s.remotepath,
    utils.ShQuote(command),
   )
- var encoded string
   if input != nil {
- encoded = base64.StdEncoding.EncodeToString(input)
- command = fmt.Sprintf(
- "head -q -c %d | base64 -d | (%s)",
- len(encoded),
- command,
- )
+ command = fmt.Sprintf("head -q -n 1 | base64 -d | (%s)", command)
   }
   command = fmt.Sprintf("(%s) 2>&1; echo %s$?", command, rcPrefix)
   if _, err := s.stdin.Write([]byte(command + "\n")); err != nil {
    return "", fmt.Errorf("failed to write command: %v", err)
   }
   if input != nil {
- if _, err := s.stdin.Write([]byte(encoded)); err != nil {
+ encoder := base64.NewEncoder(base64.StdEncoding, s.stdin)
+ if _, err := io.CopyN(encoder, input, inputlen); err != nil {
     return "", fmt.Errorf("failed to write input: %v", err)
    }
+ if err := encoder.Close(); err != nil {
+ return "", fmt.Errorf("failed to flush input: %v", err)
+ }
+ if _, err := s.stdin.Write([]byte("\n")); err != nil {
+ return "", fmt.Errorf("failed to terminate input: %v", err)
+ }
   }
   var output []string
   for s.scanner.Scan() {
@@ -259,21 +260,17 @@
   if err != nil {
    return err
   }
- buf := make([]byte, length)
- if _, err := r.Read(buf); err != nil {
- return err
- }
   path = utils.ShQuote(path)
   tmpdir := utils.ShQuote(s.tmpdir)

   // Write to a temporary file ($TMPFILE), then mv atomically.
   command := fmt.Sprintf("mkdir -p `dirname %s` && cat > $TMPFILE", path)
   command = fmt.Sprintf(
- "export TMPDIR=%...

[Andrew Wilkins (axwalk)]

Please take a look.

https://codereview.appspot.com/14197043/

[Andrew Wilkins (axwalk)]

Please take a look.

https://codereview.appspot.com/14197043/

[Roger Peppe (rogpeppe)]

Thanks for doing this. Looks good, with a few suggestions below.

https://codereview.appspot.com/14197043/diff/1003/environs/sshstorage/storage.go
File environs/sshstorage/storage.go (right):

https://codereview.appspot.com/14197043/diff/1003/environs/sshstorage/storage.go#newcode153
environs/sshstorage/storage.go:153: command = fmt.Sprintf("(%s) << EOF",
command)
s/EOF/'@EOF'/

using quotes means that the shell doesn't scan through
for potential variable expansion.

The @ character is just so it can't appear in a normal
base64 output.

https://codereview.appspot.com/14197043/diff/1003/environs/sshstorage/storage.go#newcode159
environs/sshstorage/storage.go:159: wrapper, err :=
utils.NewWrapWriter(s.stdin, base64LineLength)
You should consider wrapping s.stdin in a bufio.Writer here.
(perhaps even wrap it inside the SSHStorage and flushing
when appropriate)

https://codereview.appspot.com/14197043/diff/1003/utils/wrapwriter.go
File utils/wrapwriter.go (right):

https://codereview.appspot.com/14197043/diff/1003/utils/wrapwriter.go#newcode44
utils/wrapwriter.go:44: return n, err
s/n/total/

https://codereview.appspot.com/14197043/diff/1003/utils/wrapwriter.go#newcode44
utils/wrapwriter.go:44: return n, err
s/n/total/

https://codereview.appspot.com/14197043/diff/1003/utils/wrapwriter_test.go
File utils/wrapwriter_test.go (right):

https://codereview.appspot.com/14197043/diff/1003/utils/wrapwriter_test.go#newcode29
utils/wrapwriter_test.go:29: type test struct {
I'd like to see better tests here:

- what happens when writes return errors or part
- does it cope ok with several writes in a row?

For the latter, perhaps take a piece of input, split it up
many different ways (use some simple algorithm or
table to do that), and check that the output is
always the same.

https://codereview.appspot.com/14197043/

[John A Meinel (jameinel)]

https://codereview.appspot.com/14197043/diff/1003/environs/sshstorage/storage.go
File environs/sshstorage/storage.go (right):

https://codereview.appspot.com/14197043/diff/1003/environs/sshstorage/storage.go#newcode153
environs/sshstorage/storage.go:153: command = fmt.Sprintf("(%s) << EOF",
command)
On 2013/10/01 10:27:39, rog wrote:
> s/EOF/'@EOF'/

> using quotes means that the shell doesn't scan through
> for potential variable expansion.

> The @ character is just so it can't appear in a normal
> base64 output.

Sounds interesting, though we should test that it works properly with
/bin/dash.

https://codereview.appspot.com/14197043/

[Roger Peppe (rogpeppe)]

On 1 October 2013 12:09, <email address hidden> wrote:
>
> https://codereview.appspot.com/14197043/diff/1003/environs/sshstorage/storage.go
> File environs/sshstorage/storage.go (right):
>
> https://codereview.appspot.com/14197043/diff/1003/environs/sshstorage/storage.go#newcode153
> environs/sshstorage/storage.go:153: command = fmt.Sprintf("(%s) << EOF",
> command)
> On 2013/10/01 10:27:39, rog wrote:
>>
>> s/EOF/'@EOF'/
>
>
>> using quotes means that the shell doesn't scan through
>> for potential variable expansion.
>
>
>> The @ character is just so it can't appear in a normal
>> base64 output.
>
>
> Sounds interesting, though we should test that it works properly with
> /bin/dash.

It does (it's original Bourne shell behaviour).

[Andrew Wilkins (axwalk)]

Please take a look.

https://codereview.appspot.com/14197043/

[Roger Peppe (rogpeppe)]

LGTM modulo the below suggestions. Thanks!

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/linewrapwriter.go
File environs/sshstorage/linewrapwriter.go (right):

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/linewrapwriter.go#newcode22
environs/sshstorage/linewrapwriter.go:22: // after every "lineLength"
bytes.
// Moreover it gives no consideration to multibyte
// utf-8 characters, which it can split arbitrarily.

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/storage.go
File environs/sshstorage/storage.go (right):

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/storage.go#newcode149
environs/sshstorage/storage.go:149: stdin = bufio.NewWriter(s.stdin)
I'm not sure this half-way house is great.
I'd either make s.stdin a bufio.Writer, or
call NewWriter inside the later "if input != nil" block,
avoiding the ugly dynamic type conversion either way,
and allowing the use of the more natural WriteString
methods.

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/storage.go#newcode154
environs/sshstorage/storage.go:154: // here-document must start on the
outer-most command.
That's not actually the case, I think. Did you try it?

https://codereview.appspot.com/14197043/

[Andrew Wilkins (axwalk)]

Please take a look.

https://codereview.appspot.com/14197043/

[Andrew Wilkins (axwalk)]

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/linewrapwriter.go
File environs/sshstorage/linewrapwriter.go (right):

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/linewrapwriter.go#newcode22
environs/sshstorage/linewrapwriter.go:22: // after every "lineLength"
bytes.
On 2013/10/01 15:28:47, rog wrote:
> // Moreover it gives no consideration to multibyte
> // utf-8 characters, which it can split arbitrarily.

Done.

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/storage.go
File environs/sshstorage/storage.go (right):

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/storage.go#newcode149
environs/sshstorage/storage.go:149: stdin = bufio.NewWriter(s.stdin)
On 2013/10/01 15:28:47, rog wrote:
> I'm not sure this half-way house is great.
> I'd either make s.stdin a bufio.Writer, or
> call NewWriter inside the later "if input != nil" block,
> avoiding the ugly dynamic type conversion either way,
> and allowing the use of the more natural WriteString
> methods.

It's not in the later block because it's used before then. If it were in
there, then we'd have two writes to s.stdin.

I agree about the ugliness, though. I'll wrap s.stdin, but only in this
function so as not to lose the WriteCloser interface.

https://codereview.appspot.com/14197043/diff/14001/environs/sshstorage/storage.go#newcode154
environs/sshstorage/storage.go:154: // here-document must start on the
outer-most command.
On 2013/10/01 15:28:47, rog wrote:
> That's not actually the case, I think. Did you try it?

I did, but I'm not sure what I was doing before to cause it to fail. I
*think* I may have had a bug where something was being wrapped twice,
and ended up as ((command << EOF)).

Anyway, I've put it after 'base64 -d' and it does indeed work fine.
Thanks.

https://codereview.appspot.com/14197043/

[Go Bot (go-bot)]

There are additional revisions which have not been approved in review. Please seek review and approval of these new revisions.

=== added file 'environs/sshstorage/export_test.go'
--- environs/sshstorage/export_test.go	1970-01-01 00:00:00 +0000
+++ environs/sshstorage/export_test.go	2013-10-02 05:50:56 +0000
@@ -0,0 +1,8 @@
+// Copyright 2012, 2013 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package sshstorage
+
+var (
+	NewLineWrapWriter = newLineWrapWriter
+)

=== added file 'environs/sshstorage/linewrapwriter.go'
--- environs/sshstorage/linewrapwriter.go	1970-01-01 00:00:00 +0000
+++ environs/sshstorage/linewrapwriter.go	2013-10-02 05:50:56 +0000
@@ -0,0 +1,57 @@
+// Copyright 2013 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package sshstorage
+
+import (
+	"fmt"
+	"io"
+)
+
+type lineWrapWriter struct {
+	out    io.Writer
+	remain int
+	max    int
+}
+
+// NewLineWrapWriter returns an io.Writer that encloses the given
+// io.Writer, wrapping lines at the the specified line length.
+//
+// Note: there is no special consideration for input that
+// already contains newlines; this will simply add newlines
+// after every "lineLength" bytes. Moreover it gives no
+// consideration to multibyte utf-8 characters, which it can split
+// arbitrarily.
+//
+// This is currently only appropriate for wrapping base64-encoded
+// data, which is why it lives here.
+func newLineWrapWriter(out io.Writer, lineLength int) (io.Writer, error) {
+	if lineLength <= 0 {
+		return nil, fmt.Errorf("line length %d <= 0", lineLength)
+	}
+	return &lineWrapWriter{
+		out:    out,
+		remain: lineLength,
+		max:    lineLength,
+	}, nil
+}
+
+func (w *lineWrapWriter) Write(buf []byte) (int, error) {
+	total := 0
+	for len(buf) >= w.remain {
+		n, err := w.out.Write(buf[0:w.remain])
+		w.remain -= n
+		total += n
+		if err != nil {
+			return total, err
+		}
+		if _, err := w.out.Write([]byte("\n")); err != nil {
+			return total, err
+		}
+		w.remain = w.max
+		buf = buf[n:]
+	}
+	n, err := w.out.Write(buf)
+	w.remain -= n
+	return total + n, err
+}

=== added file 'environs/sshstorage/linewrapwriter_test.go'
--- environs/sshstorage/linewrapwriter_test.go	1970-01-01 00:00:00 +0000
+++ environs/sshstorage/linewrapwriter_test.go	2013-10-02 05:50:56 +0000
@@ -0,0 +1,151 @@
+// Copyright 2013 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package sshstorage_test
+
+import (
+	"bytes"
+	"errors"
+	"io"
+
+	gc "launchpad.net/gocheck"
+
+	"launchpad.net/juju-core/environs/sshstorage"
+)
+
+type wrapWriterSuite struct{}
+
+var _ = gc.Suite(&wrapWriterSuite{})
+
+func (*wrapWriterSuite) TestLineWrapWriterBadLength(c *gc.C) {
+	var buf bytes.Buffer
+	w, err := sshstorage.NewLineWrapWriter(&buf, 0)
+	c.Assert(err, gc.ErrorMatches, "line length 0 <= 0")
+	c.Assert(w, gc.IsNil)
+	w, err = sshstorage.NewLineWrapWriter(&buf, -1)
+	c.Assert(err, gc.ErrorMatches, "line length -1 <= 0")
+}
+
+func (*wrapWriterSuite) TestLineWrapWriter(c *gc.C) {
+	type test struct {
+		input      []string
+		lineLength int
+		expected   string
+	}
+	tests := []test{{
+		input:      []string{""},
+		lineLength: 1,
+		expected:   "",
+	}, {
+		input:      []string{"hi!"},
+		lineLength: 1,
+		expected:   "h\ni\n!\n",
+	}, {
+		input:      []string{"hi!"},
+		lineLength: 2,
+		// Note: no trailing newline.
+		expected: "hi\n!",
+	}, {
+		input:      []string{"", "h", "i!"},
+		lineLength: 2,
+		expected:   "hi\n!",
+	}, {
+		input:      []string{"", "h", "i!"},
+		lineLength: 2,
+		expected:   "hi\n!",
+	}, {
+		input:      []string{"hi", "!!"},
+		lineLength: 2,
+		expected:   "hi\n!!\n",
+	}, {
+		input:      []string{"hi", "!", "!"},
+		lineLength: 2,
+		expected:   "hi\n!!\n",
+	}, {
+		input:      []string{"h", "i", "!!"},
+		lineLength: 2,
+		expected:   "hi\n!!\n",
+	}}
+	for i, t := range tests {
+		c.Logf("test %d: %q, line length %d", i, t.input, t.lineLength)
+		var buf bytes.Buffer
+		w, err := sshstorage.NewLineWrapWriter(&buf, t.lineLength)
+		c.Assert(err, gc.IsNil)
+		c.Assert(w, gc.NotNil)
+		for _, input := range t.input {
+			n, err := w.Write([]byte(input))
+			c.Assert(err, gc.IsNil)
+			c.Assert(n, gc.Equals, len(input))
+		}
+		c.Assert(buf.String(), gc.Equals, t.expected)
+	}
+}
+
+type limitedWriter struct {
+	io.Writer
+	remaining int
+}
+
+var writeLimited = errors.New("write limited")
+
+func (w *limitedWriter) Write(buf []byte) (int, error) {
+	inputlen := len(buf)
+	if len(buf) > w.remaining {
+		buf = buf[:w.remaining]
+	}
+	n, err := w.Writer.Write(buf)
+	w.remaining -= n
+	if n < inputlen && err == nil {
+		err = writeLimited
+	}
+	return n, err
+}
+
+func (*wrapWriterSuite) TestLineWrapWriterErrors(c *gc.C) {
+	// Note: after an error is returned, all bets are off.
+	// In the only place we use this code, we bail out immediately.
+	const lineLength = 3
+	type test struct {
+		input   string
+		output  string
+		limit   int
+		written int
+		err     error
+	}
+	tests := []test{{
+		input:   "abc",
+		output:  "abc",
+		limit:   3, // "\n" will be limited
+		written: 3,
+		err:     writeLimited,
+	}, {
+		input:   "abc",
+		output:  "abc\n",
+		limit:   4,
+		written: 3, // 3/3 bytes of input
+	}, {
+		input:   "abc",
+		output:  "ab",
+		limit:   2,
+		written: 2, // 2/3 bytes of input
+		err:     writeLimited,
+	}, {
+		input:   "abc!",
+		output:  "abc\n",
+		limit:   4,
+		written: 3, // 3/4 bytes of input
+		err:     writeLimited,
+	}}
+	for i, t := range tests {
+		c.Logf("test %d: %q, limit %d", i, t.input, t.limit)
+		var buf bytes.Buffer
+		wrapWriter := &limitedWriter{&buf, t.limit}
+		w, err := sshstorage.NewLineWrapWriter(wrapWriter, lineLength)
+		c.Assert(err, gc.IsNil)
+		c.Assert(w, gc.NotNil)
+		n, err := w.Write([]byte(t.input))
+		c.Assert(n, gc.Equals, t.written)
+		c.Assert(buf.String(), gc.Equals, t.output)
+		c.Assert(err, gc.Equals, t.err)
+	}
+}

=== modified file 'environs/sshstorage/storage.go'
--- environs/sshstorage/storage.go	2013-09-27 05:19:30 +0000
+++ environs/sshstorage/storage.go	2013-10-02 05:50:56 +0000
@@ -22,6 +22,10 @@
 	"launchpad.net/juju-core/utils"
 )
 
+// base64LineLength is the default line length for wrapping
+// output generated by the base64 command line utility.
+const base64LineLength = 76
+
 // SSHStorage implements storage.Storage.
 //
 // The storage is created under sudo, and ownership given over to the
@@ -129,10 +133,10 @@
 
 func (s *SSHStorage) runf(flockmode flockmode, command string, args ...interface{}) (string, error) {
 	command = fmt.Sprintf(command, args...)
-	return s.run(flockmode, command, nil)
+	return s.run(flockmode, command, nil, 0)
 }
 
-func (s *SSHStorage) run(flockmode flockmode, command string, input []byte) (string, error) {
+func (s *SSHStorage) run(flockmode flockmode, command string, input io.Reader, inputlen int64) (string, error) {
 	const rcPrefix = "JUJU-RC: "
 	command = fmt.Sprintf(
 		"SHELL=/bin/bash flock %s %s -c %s",
@@ -140,23 +144,32 @@
 		s.remotepath,
 		utils.ShQuote(command),
 	)
-	var encoded string
+	stdin := bufio.NewWriter(s.stdin)
 	if input != nil {
-		encoded = base64.StdEncoding.EncodeToString(input)
-		command = fmt.Sprintf(
-			"head -q -c %d | base64 -d | (%s)",
-			len(encoded),
-			command,
-		)
+		command = fmt.Sprintf("base64 -d << '@EOF' | (%s)", command)
 	}
 	command = fmt.Sprintf("(%s) 2>&1; echo %s$?", command, rcPrefix)
-	if _, err := s.stdin.Write([]byte(command + "\n")); err != nil {
+	if _, err := stdin.WriteString(command + "\n"); err != nil {
 		return "", fmt.Errorf("failed to write command: %v", err)
 	}
 	if input != nil {
-		if _, err := s.stdin.Write([]byte(encoded)); err != nil {
+		wrapper, err := newLineWrapWriter(stdin, base64LineLength)
+		if err != nil {
+			return "", fmt.Errorf("failed to create split writer: %v", err)
+		}
+		encoder := base64.NewEncoder(base64.StdEncoding, wrapper)
+		if _, err := io.CopyN(encoder, input, inputlen); err != nil {
 			return "", fmt.Errorf("failed to write input: %v", err)
 		}
+		if err := encoder.Close(); err != nil {
+			return "", fmt.Errorf("failed to flush encoder: %v", err)
+		}
+		if _, err := stdin.WriteString("\n@EOF\n"); err != nil {
+			return "", fmt.Errorf("failed to terminate input: %v", err)
+		}
+	}
+	if err := stdin.Flush(); err != nil {
+		return "", fmt.Errorf("failed to flush input: %v", err)
 	}
 	var output []string
 	for s.scanner.Scan() {
@@ -259,21 +272,17 @@
 	if err != nil {
 		return err
 	}
-	buf := make([]byte, length)
-	if _, err := io.ReadFull(r, buf); err != nil {
-		return err
-	}
 	path = utils.ShQuote(path)
 	tmpdir := utils.ShQuote(s.tmpdir)
 
 	// Write to a temporary file ($TMPFILE), then mv atomically.
 	command := fmt.Sprintf("mkdir -p `dirname %s` && cat > $TMPFILE", path)
 	command = fmt.Sprintf(
-		"export TMPDIR=%s && TMPFILE=`mktemp` && ((%s && mv $TMPFILE %s) || rm -f $TMPFILE)",
+		"TMPFILE=`mktemp --tmpdir=%s` && ((%s && mv $TMPFILE %s) || rm -f $TMPFILE)",
 		tmpdir, command, path,
 	)
 
-	_, err = s.run(flockExclusive, command+"\n", buf)
+	_, err = s.run(flockExclusive, command+"\n", r, length)
 	return err
 }
 

=== modified file 'environs/sshstorage/storage_test.go'
--- environs/sshstorage/storage_test.go	2013-09-23 03:07:31 +0000
+++ environs/sshstorage/storage_test.go	2013-10-02 05:50:56 +0000
@@ -13,7 +13,6 @@
 	"path"
 	"path/filepath"
 	"regexp"
-	stdtesting "testing"
 	"time"
 
 	gc "launchpad.net/gocheck"
@@ -31,10 +30,6 @@
 
 var _ = gc.Suite(&storageSuite{})
 
-func Test(t *stdtesting.T) {
-	gc.TestingT(t)
-}
-
 func sshCommandTesting(host string, tty bool, command string) *exec.Cmd {
 	cmd := exec.Command("bash", "-c", command)
 	uid := fmt.Sprint(os.Getuid())

=== added file 'environs/sshstorage/suite_test.go'
--- environs/sshstorage/suite_test.go	1970-01-01 00:00:00 +0000
+++ environs/sshstorage/suite_test.go	2013-10-02 05:50:56 +0000
@@ -0,0 +1,14 @@
+// Copyright 2013 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package sshstorage_test
+
+import (
+	"testing"
+
+	gc "launchpad.net/gocheck"
+)
+
+func Test(t *testing.T) {
+	gc.TestingT(t)
+}