From c4c13cb2797cd11a738bddf5956af333b817df28 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arturo=20Filast=C3=B2?= <arturo@filasto.net>
Date: Fri, 29 Jun 2018 16:38:13 +0200
Subject: [PATCH] Properly compute the padding by fixing the ansi escape regexp

---
 internal/log/handlers/cli/cli.go | 12 +++++++-----
 internal/util/util.go            |  9 ++++++---
 internal/util/util_test.go       | 19 +++++++++++++++++++
 3 files changed, 32 insertions(+), 8 deletions(-)
 create mode 100644 internal/util/util_test.go

diff --git a/internal/log/handlers/cli/cli.go b/internal/log/handlers/cli/cli.go
index d6a6b15..0eaaf68 100644
--- a/internal/log/handlers/cli/cli.go
+++ b/internal/log/handlers/cli/cli.go
@@ -7,7 +7,6 @@ import (
 	"strings"
 	"sync"
 	"time"
-	"unicode/utf8"
 
 	"github.com/apex/log"
 	"github.com/fatih/color"
@@ -104,7 +103,7 @@ func logTable(w io.Writer, f log.Fields) error {
 }
 
 var bar *progress.Bar
-var lastBarChars int64
+var lastBarChars int
 
 // TypedLog is used for handling special "typed" logs to the CLI
 func (h *Handler) TypedLog(t string, e *log.Entry) error {
@@ -116,7 +115,9 @@ func (h *Handler) TypedLog(t string, e *log.Entry) error {
 		}
 		bar.Value(e.Fields.Get("percentage").(float64))
 		bar.Text(e.Message)
-		lastBarChars, err = bar.WriteTo(h.Writer)
+		barStr := bar.String()
+		fmt.Fprintf(h.Writer, "\r   %s", barStr)
+		lastBarChars = util.EscapeAwareRuneCountInString(barStr) + 3
 		return err
 	case "table":
 		return logTable(h.Writer, e.Fields)
@@ -149,10 +150,11 @@ func (h *Handler) DefaultLog(e *log.Entry) error {
 		// We need to move the cursor back to the begging of the line and add some
 		// padding to the end of the string to delete the previous line written to
 		// the console.
-		sChars := int64(utf8.RuneCountInString(s))
+		sChars := util.EscapeAwareRuneCountInString(s)
 		fmt.Fprintf(h.Writer,
-			fmt.Sprintf("\r%s%s", s, strings.Repeat(" ", int(lastBarChars-sChars))),
+			fmt.Sprintf("\r%s%s", s, strings.Repeat(" ", lastBarChars-sChars)),
 		)
+		//fmt.Fprintf(h.Writer, "\n(%d,%d)\n", lastBarChars, sChars)
 		fmt.Fprintln(h.Writer)
 		bar.WriteTo(h.Writer)
 	} else {
diff --git a/internal/util/util.go b/internal/util/util.go
index 5deb009..d647138 100644
--- a/internal/util/util.go
+++ b/internal/util/util.go
@@ -23,12 +23,15 @@ func Fatal(err error) {
 	os.Exit(1)
 }
 
-// Finds the control character sequences (like colors)
-var ctrlFinder = regexp.MustCompile("\x1b\x5b[0-9]+\x6d")
+// Finds the ansi escape sequences (like colors)
+// Taken from: https://github.com/chalk/ansi-regex/blob/d9d806ecb45d899cf43408906a4440060c5c50e5/index.js
+var ansiEscapes = regexp.MustCompile(`[\x1B\x9B][[\]()#;?]*` +
+	`(?:(?:(?:[a-zA-Z\d]*(?:;[a-zA-Z\\d]*)*)?\x07)` +
+	`|(?:(?:\d{1,4}(?:;\d{0,4})*)?[\dA-PRZcf-ntqry=><~]))`)
 
 func EscapeAwareRuneCountInString(s string) int {
 	n := utf8.RuneCountInString(s)
-	for _, sm := range ctrlFinder.FindAllString(s, -1) {
+	for _, sm := range ansiEscapes.FindAllString(s, -1) {
 		n -= utf8.RuneCountInString(sm)
 	}
 	return n
diff --git a/internal/util/util_test.go b/internal/util/util_test.go
new file mode 100644
index 0000000..4a08074
--- /dev/null
+++ b/internal/util/util_test.go
@@ -0,0 +1,19 @@
+package util
+
+import (
+	"testing"
+
+	"github.com/fatih/color"
+	ocolor "github.com/ooni/probe-cli/internal/colors"
+)
+
+func TestEscapeAwareRuneCountInString(t *testing.T) {
+	var bold = color.New(color.Bold)
+	var myColor = color.New(color.FgBlue)
+
+	s := myColor.Sprintf("•ABC%s%s", bold.Sprintf("DEF"), ocolor.Red("GHI"))
+	count := EscapeAwareRuneCountInString(s)
+	if count != 10 {
+		t.Errorf("Count was incorrect, got: %d, want: %d.", count, 10)
+	}
+}