2 files changed, 85 insertions, 22 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
 }
diff --git a/util/util_test.go b/util/util_test.go
index 1c4ae39..fd9b7b8 100644
--- a/util/util_test.go
+++ b/util/util_test.go
@@ -1,4 +1,4 @@
-package util_test;
+package util_test
 
 import (
 	"io"
@@ -22,9 +22,9 @@ func (stringReadCloser) Close() error {
 	return nil
 }
 
-func testAnsiStripper(t *testing.T, input string, expected string) {
+func testAnsiFilterReadCloser(t *testing.T, input string, expected string) {
 	reader := newStringReadCloser(input)
-	cleanReader := util.NewAnsiStripper(reader)
+	cleanReader := util.NewAnsiFilterReadCloser(reader)
 	defer cleanReader.Close()
 
 	buf := new(strings.Builder)
@@ -36,15 +36,49 @@ func testAnsiStripper(t *testing.T, input string, expected string) {
 	result := buf.String()
 
 	if n != int64(len(expected)) {
-		t.Errorf("Expected to write %d characters but wrote %d\n", 3, n)
+		t.Errorf("Expected to read %d bytes but read %d\n", len(expected), n)
 	}
 
 	if result != expected {
-		t.Errorf("Expected string `%s` but wrote `%s`", "abc", result)
+		t.Errorf("Expected string `%s` but read `%s`\n", "abc", result)
 	}
 }
 
-func TestStripAnsiColors(t *testing.T) {
-	testAnsiStripper(t, "a\x1b[31mbc", "abc")
-	testAnsiStripper(t, "[\x1b[32minfo\x1b[39m]", "[info]")
+func testAnsiFilterWriter(t *testing.T, input string, expected string) {
+	writer := new(strings.Builder)
+	w := util.NewAnsiFilterWriter(writer)
+
+	n, err := w.Write([]byte(input))
+	if err != nil {
+		t.Fatal("Error while writting input", err)
+	}
+
+	if n != len(expected) {
+		t.Errorf("Expected to write %d bytes but read %d\n", len(expected), n)
+	}
+
+	result := writer.String()
+	if result != expected {
+		t.Errorf("Expected to write `%s` but wrote `%s`\n", expected, result)
+	}
+}
+
+func TestAnsiFilterColorsReadCloser(t *testing.T) {
+	testAnsiFilterReadCloser(t, "a\x1b[31mbc", "abc")
+	testAnsiFilterReadCloser(t, "[\x1b[32minfo\x1b[39m]", "[info]")
+	testAnsiFilterReadCloser(t,
+		"FoundryVTT | 2024-10-06 20:11:14 | [\x1b[32minfo\x1b[39m] Running on Node.js - Version 20.17.0\n"+
+			"FoundryVTT | 2024-10-06 20:11:14 | [\x1b[32minfo\x1b[39m] Foundry Virtual Tabletop - Version 12 Build 331\n",
+		"FoundryVTT | 2024-10-06 20:11:14 | [info] Running on Node.js - Version 20.17.0\n"+
+			"FoundryVTT | 2024-10-06 20:11:14 | [info] Foundry Virtual Tabletop - Version 12 Build 331\n")
+}
+
+func TestAnsiFilterColorsWriter(t *testing.T) {
+	testAnsiFilterWriter(t, "a\x1b[31mbc", "abc")
+	testAnsiFilterWriter(t, "[\x1b[32minfo\x1b[39m]", "[info]")
+	testAnsiFilterWriter(t,
+		"FoundryVTT | 2024-10-06 20:11:14 | [\x1b[32minfo\x1b[39m] Running on Node.js - Version 20.17.0\n"+
+			"FoundryVTT | 2024-10-06 20:11:14 | [\x1b[32minfo\x1b[39m] Foundry Virtual Tabletop - Version 12 Build 331\n",
+		"FoundryVTT | 2024-10-06 20:11:14 | [info] Running on Node.js - Version 20.17.0\n"+
+			"FoundryVTT | 2024-10-06 20:11:14 | [info] Foundry Virtual Tabletop - Version 12 Build 331\n")
 }