From 303fb02a83a6244574fface3c3ab16170f55dcc4 Mon Sep 17 00:00:00 2001
From: Mahmood Ali <mahmood@hashicorp.com>
Date: Mon, 16 Dec 2019 09:50:10 -0500
Subject: [PATCH] cli: sequence cli.Ui operations

Fixes a bug where if a command flag parsing errors, the resulting error
and help usage messages get interleaved in unexpected and non-user
friendly way.

The reason is that we have flag parsing library effectively writes to
ui.Error in a goroutine.  This is problematic: first, we lose the sequencing between help
usage and error message; second, cli.Ui methods are not concurrent safe.

Here, we introduce a custom error writer that buffers result and calls
ui.Error() in the write method and in the same goroutine.

For context, we need to wrap ui.Error because it's line-oriented, while
flags library expects a io.Writer which is bytes oriented.
---
 command/helpers.go      | 42 ++++++++++++++++++++++++++++++++++
 command/helpers_test.go | 50 +++++++++++++++++++++++++++++++++++++++++
 command/meta.go         | 15 +------------
 3 files changed, 93 insertions(+), 14 deletions(-)

diff --git a/command/helpers.go b/command/helpers.go
index 38b986870e35..697f49c3098a 100644
--- a/command/helpers.go
+++ b/command/helpers.go
@@ -1,6 +1,7 @@
 package command
 
 import (
+	"bufio"
 	"bytes"
 	"fmt"
 	"io"
@@ -14,6 +15,7 @@ import (
 	"github.com/hashicorp/nomad/api"
 	"github.com/hashicorp/nomad/jobspec"
 	"github.com/kr/text"
+	"github.com/mitchellh/cli"
 	"github.com/posener/complete"
 
 	"github.com/ryanuber/columnize"
@@ -464,3 +466,43 @@ func sanitizeUUIDPrefix(prefix string) string {
 func commandErrorText(cmd NamedCommand) string {
 	return fmt.Sprintf("For additional help try 'nomad %s -help'", cmd.Name())
 }
+
+// uiErrorWriter is a io.Writer that wraps underlying ui.ErrorWriter().
+// ui.ErrorWriter expects full lines as inputs and it emits its own line breaks.
+//
+// uiErrorWriter scans input for individual lines to pass to ui.ErrorWriter. If data
+// doesn't contain a new line, it buffers result until next new line or writer is closed.
+type uiErrorWriter struct {
+	ui  cli.Ui
+	buf bytes.Buffer
+}
+
+func (w *uiErrorWriter) Write(data []byte) (int, error) {
+	read := 0
+	for len(data) != 0 {
+		a, token, err := bufio.ScanLines(data, false)
+		if err != nil {
+			return read, err
+		}
+
+		if a == 0 {
+			r, err := w.buf.Write(data)
+			return read + r, err
+		}
+
+		w.ui.Error(w.buf.String() + string(token))
+		data = data[a:]
+		w.buf.Reset()
+		read += a
+	}
+
+	return read, nil
+}
+
+func (w *uiErrorWriter) Close() error {
+	// emit what's remaining
+	if w.buf.Len() != 0 {
+		w.ui.Error(w.buf.String())
+	}
+	return nil
+}
diff --git a/command/helpers_test.go b/command/helpers_test.go
index d11b252a2471..139acd30755b 100644
--- a/command/helpers_test.go
+++ b/command/helpers_test.go
@@ -1,6 +1,7 @@
 package command
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -16,6 +17,7 @@ import (
 	"github.com/hashicorp/nomad/helper/flatmap"
 	"github.com/kr/pretty"
 	"github.com/mitchellh/cli"
+	"github.com/stretchr/testify/require"
 )
 
 func TestHelpers_FormatKV(t *testing.T) {
@@ -342,3 +344,51 @@ func TestPrettyTimeDiff(t *testing.T) {
 	}
 
 }
+
+// TestUiErrorWriter asserts that writer buffers and
+func TestUiErrorWriter(t *testing.T) {
+	t.Parallel()
+
+	var outBuf, errBuf bytes.Buffer
+	ui := &cli.BasicUi{
+		Writer:      &outBuf,
+		ErrorWriter: &errBuf,
+	}
+
+	w := &uiErrorWriter{ui: ui}
+
+	inputs := []string{
+		"some line\n",
+		"multiple\nlines\r\nhere",
+		" with  followup\nand",
+		" more lines ",
+		" without new line ",
+		"until here\nand then",
+		"some more",
+	}
+
+	partialAcc := ""
+	for _, in := range inputs {
+		n, err := w.Write([]byte(in))
+		require.NoError(t, err)
+		require.Equal(t, len(in), n)
+
+		// assert that writer emits partial result until last new line
+		partialAcc += strings.ReplaceAll(in, "\r\n", "\n")
+		lastNL := strings.LastIndex(partialAcc, "\n")
+		require.Equal(t, partialAcc[:lastNL+1], errBuf.String())
+	}
+
+	require.Empty(t, outBuf.String())
+
+	// note that the \r\n got replaced by \n
+	expectedErr := "some line\nmultiple\nlines\nhere with  followup\nand more lines  without new line until here\n"
+	require.Equal(t, expectedErr, errBuf.String())
+
+	// close emits the final line
+	err := w.Close()
+	require.NoError(t, err)
+
+	expectedErr += "and thensome more\n"
+	require.Equal(t, expectedErr, errBuf.String())
+}
diff --git a/command/meta.go b/command/meta.go
index 6e6f0e1e8c14..363e666ceb98 100644
--- a/command/meta.go
+++ b/command/meta.go
@@ -1,9 +1,7 @@
 package command
 
 import (
-	"bufio"
 	"flag"
-	"io"
 	"os"
 	"strings"
 
@@ -83,18 +81,7 @@ func (m *Meta) FlagSet(n string, fs FlagSetFlags) *flag.FlagSet {
 
 	}
 
-	// Create an io.Writer that writes to our UI properly for errors.
-	// This is kind of a hack, but it does the job. Basically: create
-	// a pipe, use a scanner to break it into lines, and output each line
-	// to the UI. Do this forever.
-	errR, errW := io.Pipe()
-	errScanner := bufio.NewScanner(errR)
-	go func() {
-		for errScanner.Scan() {
-			m.Ui.Error(errScanner.Text())
-		}
-	}()
-	f.SetOutput(errW)
+	f.SetOutput(&uiErrorWriter{ui: m.Ui})
 
 	return f
 }