improved implementation of converting ANSI color to HTML

internal/api/v1/utils/error.go

@@ -4,6 +4,7 @@ import (
 	"net/http"
 
 	E "github.com/yusing/go-proxy/internal/error"
+	"github.com/yusing/go-proxy/internal/logging"
 	"github.com/yusing/go-proxy/internal/utils/strutils/ansi"
 )
 
@@ -30,8 +31,16 @@ func RespondError(w http.ResponseWriter, err error, code ...int) {
 	if len(code) == 0 {
 		code = []int{http.StatusBadRequest}
 	}
-	// strip ANSI color codes added from Error.WithSubject
-	http.Error(w, ansi.StripANSI(err.Error()), code[0])
+	w.Header().Set("Content-Type", "text/html; charset=utf-8")
+	buf := make([]byte, 0, 100)
+	errMsg := err.Error()
+	buf, err = logging.FormatMessageToHTMLBytes(errMsg, buf)
+	if err != nil {
+		http.Error(w, ansi.StripANSI(errMsg), code[0])
+		return
+	}
+	w.WriteHeader(code[0])
+	_, _ = w.Write(buf)
 }
 
 func ErrMissingKey(k string) error {

internal/logging/html.go

@@ -1,99 +1,14 @@
 package logging
 
 import (
-	"bytes"
+	"errors"
+	"fmt"
 	"time"
 
 	"github.com/rs/zerolog"
 	"github.com/yusing/go-proxy/internal/common"
-	ansiPkg "github.com/yusing/go-proxy/internal/utils/strutils/ansi"
 )
 
-func fmtMessageToHTMLBytes(msg string, buf []byte) []byte {
-	buf = append(buf, []byte(`<span class="log-message">`)...)
-	var last byte
-
-	isAnsi := false
-	nAnsi := 0
-	ansi := bytes.NewBuffer(make([]byte, 0, 4))
-	ansiContent := bytes.NewBuffer(make([]byte, 0, 30))
-	style := bytes.NewBuffer(make([]byte, 0, 30))
-
-	for _, r := range msg {
-		if last == '\n' {
-			buf = append(buf, prefixHTML...)
-		}
-		if last == '\x1b' {
-			if r != 'm' {
-				ansi.WriteRune(r)
-				if r == '[' && ansiContent.Len() > 0 {
-					buf = append(buf, []byte(`<span `)...)
-					buf = append(buf, style.Bytes()...)
-					buf = append(buf, []byte(`>`)...)
-					buf = append(buf, ansiContent.Bytes()...)
-					style.Reset()
-					ansiContent.Reset()
-					nAnsi++
-				}
-			} else {
-				ansiCode := ansi.String()
-				switch ansiCode {
-				case "[0": // reset
-					if style.Len() > 0 {
-						buf = append(buf, []byte(`<span `)...)
-						buf = append(buf, style.Bytes()...)
-						buf = append(buf, []byte(`>`)...)
-					}
-					for nAnsi-1 > 0 {
-						buf = append(buf, []byte(`</span>`)...)
-						nAnsi--
-					}
-					nAnsi = 0
-					buf = append(buf, ansiContent.Bytes()...)
-					buf = append(buf, []byte(`</span>`)...)
-					isAnsi = false
-					ansiContent.Reset()
-					style.Reset()
-				case "[1": // bold
-					style.WriteString(`class="log-bold" `)
-				default:
-					className, ok := ansiPkg.ToHTMLClass[ansiCode]
-					if ok {
-						style.WriteString(`class="` + className + `" `)
-					} else {
-						style.WriteString(`class="log-unknown-ansi" `)
-					}
-				}
-				ansi.Reset()
-				last = 0
-			}
-			continue
-		}
-
-		last = byte(r)
-		if r == '\x1b' {
-			isAnsi = true
-			continue
-		}
-		if isAnsi || nAnsi > 0 {
-			if symbol, ok := symbolMapping[r]; ok {
-				ansiContent.Write(symbol)
-			} else {
-				ansiContent.WriteRune(r)
-			}
-		} else {
-			if symbol, ok := symbolMapping[r]; ok {
-				buf = append(buf, symbol...)
-			} else {
-				buf = append(buf, last)
-			}
-		}
-	}
-
-	buf = append(buf, []byte("</span>")...)
-	return buf
-}
-
 var levelHTMLFormats = [][]byte{
 	[]byte(` <span class="log-trace">TRC</span> `),
 	[]byte(` <span class="log-debug">DBG</span> `),
@@ -104,12 +19,124 @@ var levelHTMLFormats = [][]byte{
 	[]byte(` <span class="log-panic">PAN</span> `),
 }
 
-var symbolMapping = map[rune][]byte{
-	'•':  []byte("&middot;"),
-	'>':  []byte("&gt;"),
-	'<':  []byte("&lt;"),
-	'\t': []byte("&ensp;"),
-	'\n': []byte("<br>"),
+var colorToClass = map[string]string{
+	"1":  "log-bold",
+	"3":  "log-italic",
+	"4":  "log-underline",
+	"30": "log-black",
+	"31": "log-red",
+	"32": "log-green",
+	"33": "log-yellow",
+	"34": "log-blue",
+	"35": "log-magenta",
+	"36": "log-cyan",
+	"37": "log-white",
+	"90": "log-bright-black",
+	"91": "log-red",
+	"92": "log-bright-green",
+	"93": "log-bright-yellow",
+	"94": "log-bright-blue",
+	"95": "log-bright-magenta",
+	"96": "log-bright-cyan",
+	"97": "log-bright-white",
+}
+
+// FormatMessageToHTMLBytes converts text with ANSI color codes to HTML with class names.
+// ANSI codes are mapped to classes via a static map, and reset codes ([0m) close all spans.
+// Time complexity is O(n) with minimal allocations.
+func FormatMessageToHTMLBytes(msg string, buf []byte) ([]byte, error) {
+	buf = append(buf, "<span class=\"log-message\">"...)
+	var stack []string
+	lastPos := 0
+
+	for i := 0; i < len(msg); {
+		if msg[i] == '\x1b' && i+1 < len(msg) && msg[i+1] == '[' {
+			if lastPos < i {
+				escapeAndAppend(msg[lastPos:i], &buf)
+			}
+			i += 2 // Skip \x1b[
+
+			start := i
+			for ; i < len(msg) && msg[i] != 'm'; i++ {
+				if !isANSICodeChar(msg[i]) {
+					return nil, fmt.Errorf("invalid ANSI char: %c", msg[i])
+				}
+			}
+
+			if i >= len(msg) {
+				return nil, errors.New("unterminated ANSI sequence")
+			}
+
+			codeStr := msg[start:i]
+			i++ // Skip 'm'
+			lastPos = i
+
+			startPart := 0
+			for j := 0; j <= len(codeStr); j++ {
+				if j == len(codeStr) || codeStr[j] == ';' {
+					part := codeStr[startPart:j]
+					if part == "" {
+						return nil, errors.New("empty code part")
+					}
+
+					if part == "0" {
+						for range stack {
+							buf = append(buf, "</span>"...)
+						}
+						stack = stack[:0]
+					} else {
+						className, ok := colorToClass[part]
+						if !ok {
+							return nil, fmt.Errorf("invalid ANSI code: %s", part)
+						}
+						stack = append(stack, className)
+						buf = append(buf, `<span class="`...)
+						buf = append(buf, className...)
+						buf = append(buf, `">`...)
+					}
+					startPart = j + 1
+				}
+			}
+		} else {
+			i++
+		}
+	}
+
+	if lastPos < len(msg) {
+		escapeAndAppend(msg[lastPos:], &buf)
+	}
+
+	for range stack {
+		buf = append(buf, "</span>"...)
+	}
+
+	buf = append(buf, "</span>"...)
+	return buf, nil
+}
+
+func isANSICodeChar(c byte) bool {
+	return (c >= '0' && c <= '9') || c == ';'
+}
+
+func escapeAndAppend(s string, buf *[]byte) {
+	for i, r := range s {
+		switch r {
+		case '•':
+			*buf = append(*buf, "&middot;"...)
+		case '&':
+			*buf = append(*buf, "&amp;"...)
+		case '<':
+			*buf = append(*buf, "&lt;"...)
+		case '>':
+			*buf = append(*buf, "&gt;"...)
+		case '\t':
+			*buf = append(*buf, "&#9;"...)
+		case '\n':
+			*buf = append(*buf, "<br>"...)
+		default:
+			*buf = append(*buf, s[i])
+		}
+	}
 }
 
 func timeNowHTML() []byte {
@@ -125,7 +152,7 @@ func FormatLogEntryHTML(level zerolog.Level, message string, buf []byte) []byte
 	if level < zerolog.NoLevel {
 		buf = append(buf, levelHTMLFormats[level+1]...)
 	}
-	buf = fmtMessageToHTMLBytes(message, buf)
+	buf, _ = FormatMessageToHTMLBytes(message, buf)
 	buf = append(buf, []byte("</pre>")...)
 	return buf
 }

internal/logging/html_test.go

@@ -16,14 +16,14 @@ func TestFormatHTML(t *testing.T) {
 func TestFormatHTMLANSI(t *testing.T) {
 	buf := make([]byte, 0, 100)
 	buf = FormatLogEntryHTML(zerolog.InfoLevel, "This is \x1b[91m\x1b[1ma test.\x1b[0mOK!.", buf)
-	ExpectEqual(t, string(buf), `<pre class="log-entry">01-01 01:01 <span class="log-info">INF</span> <span class="log-message">This is <span class="log-red" class="log-bold" >a test.</span>OK!.</span></pre>`)
+	ExpectEqual(t, string(buf), `<pre class="log-entry">01-01 01:01 <span class="log-info">INF</span> <span class="log-message">This is <span class="log-red"><span class="log-bold">a test.</span></span>OK!.</span></pre>`)
 	buf = buf[:0]
 	buf = FormatLogEntryHTML(zerolog.InfoLevel, "This is \x1b[91ma \x1b[1mtest.\x1b[0mOK!.", buf)
-	ExpectEqual(t, string(buf), `<pre class="log-entry">01-01 01:01 <span class="log-info">INF</span> <span class="log-message">This is <span class="log-red" >a <span class="log-bold" >test.</span></span>OK!.</span></pre>`)
+	ExpectEqual(t, string(buf), `<pre class="log-entry">01-01 01:01 <span class="log-info">INF</span> <span class="log-message">This is <span class="log-red">a <span class="log-bold">test.</span></span>OK!.</span></pre>`)
 }
 
 func BenchmarkFormatLogEntryHTML(b *testing.B) {
-	buf := make([]byte, 0, 100)
+	buf := make([]byte, 0, 250)
 	for range b.N {
 		FormatLogEntryHTML(zerolog.InfoLevel, "This is \x1b[91ma \x1b[1mtest.\x1b[0mOK!.", buf)
 	}