From 3d35eebb583413c7298135bda733381c28c6ad40 Mon Sep 17 00:00:00 2001
From: Pat Gavlin <pat@pulumi.com>
Date: Tue, 29 Oct 2019 10:55:30 -0700
Subject: [PATCH] Fix another colorizer bug. (#3417)

If a maximum length was specified and the string did not end in a
colorization directive, the maximum length was not respected.

Fixes #3374.
---
 CHANGELOG.md                   |  3 +++
 pkg/diag/colors/colors.go      | 17 +++++++++--------
 pkg/diag/colors/colors_test.go | 24 ++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 262cd09b1..8883205a1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,9 @@ CHANGELOG
 - Adds a **preview** of .NET support for Pulumi. This code is an preview state and is subject
   to change at any point.
 
+- Fix another colorizer issue that could cause garbled output for messages that did not end in colorization tags.
+  [#3417](https://github.com/pulumi/pulumi/pull/3417)
+
 ## 1.4.0 (2019-10-24)
 
 - `FileAsset` in the Python SDK now accepts anything implementing `os.PathLike` in addition to `str`.
diff --git a/pkg/diag/colors/colors.go b/pkg/diag/colors/colors.go
index 87ae81cea..06fb0b7ee 100644
--- a/pkg/diag/colors/colors.go
+++ b/pkg/diag/colors/colors.go
@@ -132,7 +132,7 @@ func writeDirective(w io.StringWriter, c Colorization, directive string) {
 func colorizeText(s string, c Colorization, maxLen int) string {
 	var buf bytes.Buffer
 
-	textLen := 0
+	textLen, reset := 0, false
 	for input := s; len(input) > 0; {
 		// Do we have another directive to process?
 		nextDirectiveStart := strings.Index(input, colorLeft)
@@ -146,10 +146,8 @@ func colorizeText(s string, c Colorization, maxLen int) string {
 				return input
 			}
 
-			// Otherwise, write the remaining input into the buffer and terminate.
-			_, err := buf.WriteString(input)
-			contract.IgnoreError(err)
-			break
+			// Otherwise, set the start of the next directive to the end of the string and continue.
+			nextDirectiveStart = len(input)
 		}
 		if buf.Cap() < len(input) {
 			buf.Grow(len(input))
@@ -160,12 +158,13 @@ func colorizeText(s string, c Colorization, maxLen int) string {
 		if maxLen >= 0 && textLen+len(text) > maxLen {
 			_, err := buf.WriteString(text[:maxLen-textLen])
 			contract.IgnoreError(err)
-			writeDirective(&buf, c, Reset)
+			if reset {
+				writeDirective(&buf, c, Reset)
+			}
 			break
 		}
 		_, err := buf.WriteString(text)
 		contract.IgnoreError(err)
-		input = input[nextDirectiveStart+len(colorLeft):]
 		textLen += len(text)
 
 		// If we have a start delimiter but no end delimiter, terminate. The partial command will not be present in the
@@ -175,9 +174,11 @@ func colorizeText(s string, c Colorization, maxLen int) string {
 			break
 		}
 
-		directive := command(input[:nextDirectiveEnd])
+		directive := command(input[nextDirectiveStart+len(colorLeft) : nextDirectiveEnd])
 		writeDirective(&buf, c, directive)
 		input = input[nextDirectiveEnd+len(colorRight):]
+
+		reset = directive != Reset
 	}
 
 	return buf.String()
diff --git a/pkg/diag/colors/colors_test.go b/pkg/diag/colors/colors_test.go
index ceba163d5..816e4c226 100644
--- a/pkg/diag/colors/colors_test.go
+++ b/pkg/diag/colors/colors_test.go
@@ -80,3 +80,27 @@ func TestColorizer(t *testing.T) {
 		})
 	}
 }
+
+// TestTrimColorizedString provides extra coverage for TrimColorizedString.
+func TestTrimColorizedString(t *testing.T) {
+	str := "hello, " + Green + "world" + Reset + "!!"
+
+	actual := TrimColorizedString(str, len("hello"))
+	assert.Equal(t, "hello", actual)
+
+	actual = TrimColorizedString(str, len("hello, wo"))
+	assert.Equal(t, "hello, "+Green+"wo"+Reset, actual)
+
+	actual = TrimColorizedString(str, len("hello, world"))
+	assert.Equal(t, "hello, "+Green+"world"+Reset, actual)
+
+	actual = TrimColorizedString(str, len("hello, world!"))
+	assert.Equal(t, "hello, "+Green+"world"+Reset+"!", actual)
+
+	actual = TrimColorizedString(str, len("hello, world!!"))
+	assert.Equal(t, str, actual)
+
+	plain := "hello, world!!"
+	actual = TrimColorizedString(plain, len("hello"))
+	assert.Equal(t, "hello", actual)
+}