fix log wrapped incorrectly in WebUI, implement log SSR

cmd/main.go

@@ -9,6 +9,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/rs/zerolog"
 	"github.com/yusing/go-proxy/internal"
 	v1 "github.com/yusing/go-proxy/internal/api/v1"
 	"github.com/yusing/go-proxy/internal/api/v1/auth"
@@ -28,11 +29,12 @@ import (
 var rawLogger = log.New(os.Stdout, "", 0)
 
 func init() {
-	var out io.Writer = os.Stdout
+	var out io.Writer = os.Stderr
 	if common.EnableLogStreaming {
-		out = io.MultiWriter(out, v1.MemLogger())
+		out = zerolog.MultiLevelWriter(out, v1.GetMemLogger())
 	}
 	logging.InitLogger(out)
+	// logging.AddHook(v1.GetMemLogger())
 	internal.InitIconListCache()
 	homepage.InitOverridesConfig()
 	favicon.InitIconCache()

cmd/main_production.go

@@ -2,4 +2,6 @@
 
 package main
 
-func initProfiling() {}
+func initProfiling() {
+	// no profiling in production
+}

internal/api/v1/mem_logger.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/coder/websocket"
+	"github.com/rs/zerolog"
 	"github.com/yusing/go-proxy/internal/api/v1/utils"
 	"github.com/yusing/go-proxy/internal/common"
 	config "github.com/yusing/go-proxy/internal/config/types"
@@ -26,16 +27,36 @@ type memLogger struct {
 	sync.RWMutex
 	notifyLock sync.RWMutex
 	connChans  F.Map[chan *logEntryRange, struct{}]
+
+	bufPool sync.Pool // used in hook mode
+}
+
+type MemLogger interface {
+	io.Writer
+	// TODO: hook does not pass in fields, looking for a workaround to do server side log rendering
+	zerolog.Hook
+}
+
+type buffer struct {
+	data []byte
 }
 
 const (
 	maxMemLogSize         = 16 * 1024
 	truncateSize          = maxMemLogSize / 2
 	initialWriteChunkSize = 4 * 1024
+	hookModeBufSize       = 256
 )
 
 var memLoggerInstance = &memLogger{
 	connChans: F.NewMapOf[chan *logEntryRange, struct{}](),
+	bufPool: sync.Pool{
+		New: func() any {
+			return &buffer{
+				data: make([]byte, 0, hookModeBufSize),
+			}
+		},
+	},
 }
 
 func init() {
@@ -68,27 +89,28 @@ func LogsWS() func(config config.ConfigInstance, w http.ResponseWriter, r *http.
 	return memLoggerInstance.ServeHTTP
 }
 
-func MemLogger() io.Writer {
+func GetMemLogger() MemLogger {
 	return memLoggerInstance
 }
 
-func (m *memLogger) Write(p []byte) (n int, err error) {
+func (m *memLogger) truncateIfNeeded(n int) {
 	m.RLock()
-	if m.Len() > maxMemLogSize {
-		m.Truncate(truncateSize)
-	}
+	needTruncate := m.Len()+n > maxMemLogSize
 	m.RUnlock()
 
-	n = len(p)
-	m.Lock()
-	pos := m.Len()
-	_, err = m.Buffer.Write(p)
-	m.Unlock()
+	if needTruncate {
+		m.Lock()
+		defer m.Unlock()
+		needTruncate = m.Len()+n > maxMemLogSize
+		if !needTruncate {
+			return
+		}
 
-	if err != nil {
-		return
+		m.Truncate(truncateSize)
 	}
+}
 
+func (m *memLogger) notifyWS(pos, n int) {
 	if m.connChans.Size() > 0 {
 		timeout := time.NewTimer(1 * time.Second)
 		defer timeout.Stop()
@@ -106,6 +128,51 @@ func (m *memLogger) Write(p []byte) (n int, err error) {
 		})
 		return
 	}
+}
+
+func (m *memLogger) writeBuf(b []byte) (pos int, err error) {
+	m.Lock()
+	defer m.Unlock()
+	pos = m.Len()
+	_, err = m.Buffer.Write(b)
+	return
+}
+
+// Run implements zerolog.Hook.
+func (m *memLogger) Run(e *zerolog.Event, level zerolog.Level, message string) {
+	bufStruct := m.bufPool.Get().(*buffer)
+	buf := bufStruct.data
+	defer func() {
+		bufStruct.data = bufStruct.data[:0]
+		m.bufPool.Put(bufStruct)
+	}()
+
+	buf = logging.FormatLogEntryHTML(level, message, buf)
+	n := len(buf)
+
+	m.truncateIfNeeded(n)
+
+	pos, err := m.writeBuf(buf)
+	if err != nil {
+		// not logging the error here, it will cause Run to be called again = infinite loop
+		return
+	}
+
+	m.notifyWS(pos, n)
+}
+
+// Write implements io.Writer.
+func (m *memLogger) Write(p []byte) (n int, err error) {
+	n = len(p)
+	m.truncateIfNeeded(n)
+
+	pos, err := m.writeBuf(p)
+	if err != nil {
+		// not logging the error here, it will cause Run to be called again = infinite loop
+		return
+	}
+
+	m.notifyWS(pos, n)
 	return
 }
 

internal/logging/html.go

@@ -0,0 +1,131 @@
+package logging
+
+import (
+	"bytes"
+	"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> `),
+	[]byte(` <span class="log-info">INF</span> `),
+	[]byte(` <span class="log-warn">WRN</span> `),
+	[]byte(` <span class="log-error">ERR</span> `),
+	[]byte(` <span class="log-fatal">FTL</span> `),
+	[]byte(` <span class="log-panic">PAN</span> `),
+}
+
+var symbolMapping = map[rune][]byte{
+	'•':  []byte("&middot;"),
+	'>':  []byte("&gt;"),
+	'<':  []byte("&lt;"),
+	'\t': []byte("&ensp;"),
+	'\n': []byte("<br>"),
+}
+
+func timeNowHTML() []byte {
+	if !common.IsTest {
+		return []byte(time.Now().Format(timeFmt))
+	}
+	return []byte(time.Date(2024, 1, 1, 1, 1, 1, 1, time.UTC).Format(timeFmt))
+}
+
+func FormatLogEntryHTML(level zerolog.Level, message string, buf []byte) []byte {
+	buf = append(buf, []byte(`<pre class="log-entry">`)...)
+	buf = append(buf, timeNowHTML()...)
+	if level < zerolog.NoLevel {
+		buf = append(buf, levelHTMLFormats[level+1]...)
+	}
+	buf = fmtMessageToHTMLBytes(message, buf)
+	buf = append(buf, []byte("</pre>")...)
+	return buf
+}

internal/logging/html_test.go

@@ -0,0 +1,30 @@
+package logging
+
+import (
+	"testing"
+
+	"github.com/rs/zerolog"
+	. "github.com/yusing/go-proxy/internal/utils/testing"
+)
+
+func TestFormatHTML(t *testing.T) {
+	buf := make([]byte, 0, 100)
+	buf = FormatLogEntryHTML(zerolog.InfoLevel, "This is a test.\nThis is a new line.", 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 a test.<br>`+prefix+`This is a new line.</span></pre>`)
+}
+
+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>`)
+	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>`)
+}
+
+func BenchmarkFormatLogEntryHTML(b *testing.B) {
+	buf := make([]byte, 0, 100)
+	for range b.N {
+		FormatLogEntryHTML(zerolog.InfoLevel, "This is \x1b[91ma \x1b[1mtest.\x1b[0mOK!.", buf)
+	}
+}

internal/logging/logging.go

@@ -10,13 +10,15 @@ import (
 	"github.com/yusing/go-proxy/internal/utils/strutils"
 )
 
-var logger zerolog.Logger
-
-func InitLogger(out io.Writer) {
-	var timeFmt string
-	var level zerolog.Level
-	var exclude []string
+var (
+	logger     zerolog.Logger
+	timeFmt    string
+	level      zerolog.Level
+	prefix     string
+	prefixHTML []byte
+)
 
+func init() {
 	switch {
 	case common.IsTrace:
 		timeFmt = "04:05"
@@ -28,27 +30,37 @@ func InitLogger(out io.Writer) {
 		timeFmt = "01-02 15:04"
 		level = zerolog.InfoLevel
 	}
-
 	prefixLength := len(timeFmt) + 5 // level takes 3 + 2 spaces
-	prefix := strings.Repeat(" ", prefixLength)
+	prefix = strings.Repeat(" ", prefixLength)
+	// prefixHTML = []byte(strings.Repeat(" ", prefixLength))
+	prefixHTML = []byte(prefix)
 
-	logger = zerolog.New(
-		zerolog.ConsoleWriter{
-			Out:           out,
-			TimeFormat:    timeFmt,
-			FieldsExclude: exclude,
-			FormatMessage: func(msgI interface{}) string { // pad spaces for each line
-				msg := msgI.(string)
-				lines := strutils.SplitRune(msg, '\n')
-				if len(lines) == 1 {
-					return msg
-				}
-				for i := 1; i < len(lines); i++ {
-					lines[i] = prefix + lines[i]
-				}
-				return strutils.JoinRune(lines, '\n')
-			},
+	if zerolog.TraceLevel != -1 && zerolog.NoLevel != 6 {
+		panic("zerolog implementation changed")
+	}
+}
+
+func fmtMessage(msg string) string {
+	lines := strutils.SplitRune(msg, '\n')
+	if len(lines) == 1 {
+		return msg
+	}
+	for i := 1; i < len(lines); i++ {
+		lines[i] = prefix + lines[i]
+	}
+	return strutils.JoinRune(lines, '\n')
+}
+
+func InitLogger(out io.Writer) {
+	writer := zerolog.ConsoleWriter{
+		Out:        out,
+		TimeFormat: timeFmt,
+		FormatMessage: func(msgI interface{}) string { // pad spaces for each line
+			return fmtMessage(msgI.(string))
 		},
+	}
+	logger = zerolog.New(
+		writer,
 	).Level(level).With().Timestamp().Logger()
 }
 

internal/utils/strutils/ansi/ansi.go

@@ -23,3 +23,11 @@ const (
 func StripANSI(s string) string {
 	return ansiRegexp.ReplaceAllString(s, "")
 }
+
+var ToHTMLClass = map[string]string{
+	"[91": "log-red",
+	"[92": "log-green",
+	"[93": "log-yellow",
+	"[96": "log-cyan",
+	"[97": "log-white",
+}