1 files changed, 43 insertions, 14 deletions
diff --git a/util/util.go b/util/util.go
index deb528b..dc73941 100644
--- a/util/util.go
+++ b/util/util.go
@@ -5,34 +5,63 @@ import (
 	"regexp"
 )
 
+const ANSI_REGEX string = "\x1b[[0-9;]*m"
+
 // Strip ansi control characters from source
 //
 // The source reader is automatically closed when stripper reader is
-type AnsiStripper struct {
-	regex  regexp.Regexp
-	source io.ReadCloser
+type AnsiFilterReadCloser struct {
+	regex      regexp.Regexp
+	readCloser io.ReadCloser
 }
 
-func NewAnsiStripper(source io.ReadCloser) *AnsiStripper {
-	striper := new(AnsiStripper)
-	striper.regex = *regexp.MustCompile(`\x1b[[0-9;]*m`)
-	striper.source = source
+func NewAnsiFilterReadCloser(readCloser io.ReadCloser) *AnsiFilterReadCloser {
+	r := new(AnsiFilterReadCloser)
+	r.regex = *regexp.MustCompile(ANSI_REGEX)
+	r.readCloser = readCloser
 
-	return striper
+	return r
 }
 
-func (rc AnsiStripper) Read(p []byte) (int, error) {
-	n, err := rc.source.Read(p)
+func (rc AnsiFilterReadCloser) Read(p []byte) (int, error) {
+	n, err := rc.readCloser.Read(p)
 	if err != io.EOF && err != nil {
 		return n, err
 	}
 
-	stripped := rc.regex.ReplaceAll(p[:n], []byte(""))
-	n = copy(p, stripped)
+	filtered := rc.regex.ReplaceAll(p[:n], []byte(""))
+	n = copy(p, filtered)
 
 	return n, err
 }
 
-func (rc AnsiStripper) Close() error {
-	return rc.source.Close()
+func (rc AnsiFilterReadCloser) Close() error {
+	return rc.readCloser.Close()
+}
+
+type AnsiFilterWriter struct {
+	regex  regexp.Regexp
+	writer io.Writer
+}
+
+func NewAnsiFilterWriter(writer io.Writer) *AnsiFilterWriter {
+	w := new(AnsiFilterWriter)
+	w.regex = *regexp.MustCompile(ANSI_REGEX)
+	w.writer = writer
+
+	return w
+}
+
+func (w AnsiFilterWriter) Write(p []byte) (int, error) {
+	filtered := w.regex.ReplaceAll(p, []byte(""))
+
+	n, err := w.writer.Write(filtered)
+
+	if n < len(filtered) {
+		return n, io.ErrShortWrite
+	}
+
+	// NOTE: report the original size not the filtered size to indicate to caller
+    //       that input has been consumed
+	return len(p), err
 }