wk-backend/pkg/log/buffer.go

package log

import (
	"encoding/json"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

const (
	DefaultBufferLimit = 1000
)

type Entry struct {
	ID      int64          `json:"id"`
	Time    string         `json:"time"`
	Level   string         `json:"level"`
	Source  string         `json:"source"`
	Message string         `json:"message"`
	Logger  string         `json:"logger,omitempty"`
	Caller  string         `json:"caller,omitempty"`
	Fields  map[string]any `json:"fields,omitempty"`
}

type bufferHub struct {
	mu          sync.RWMutex
	limit       int
	nextEntryID int64
	nextSubID   int
	entries     []Entry
	subscribers map[int]chan Entry
}

type memoryCore struct {
	level  zapcore.LevelEnabler
	fields []zap.Field
}

var defaultHub = newBufferHub(DefaultBufferLimit)

func newBufferHub(limit int) *bufferHub {
	return &bufferHub{
		limit:       limit,
		entries:     make([]Entry, 0, limit),
		subscribers: make(map[int]chan Entry),
	}
}

func (h *bufferHub) append(entry Entry) Entry {
	h.mu.Lock()
	h.nextEntryID++
	entry.ID = h.nextEntryID
	if len(h.entries) >= h.limit {
		h.entries = append(h.entries[1:], entry)
	} else {
		h.entries = append(h.entries, entry)
	}

	subscribers := make([]chan Entry, 0, len(h.subscribers))
	for _, ch := range h.subscribers {
		subscribers = append(subscribers, ch)
	}
	h.mu.Unlock()

	for _, ch := range subscribers {
		select {
		case ch <- entry:
		default:
		}
	}

	return entry
}

func (h *bufferHub) snapshot() []Entry {
	h.mu.RLock()
	defer h.mu.RUnlock()

	entries := make([]Entry, len(h.entries))
	copy(entries, h.entries)
	return entries
}

func (h *bufferHub) subscribe() (int, <-chan Entry) {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.nextSubID++
	id := h.nextSubID
	ch := make(chan Entry, 256)
	h.subscribers[id] = ch
	return id, ch
}

func (h *bufferHub) unsubscribe(id int) {
	h.mu.Lock()
	defer h.mu.Unlock()

	ch, ok := h.subscribers[id]
	if !ok {
		return
	}
	delete(h.subscribers, id)
	close(ch)
}

func Entries() []Entry {
	return defaultHub.snapshot()
}

func Subscribe() (int, <-chan Entry) {
	return defaultHub.subscribe()
}

func Unsubscribe(id int) {
	defaultHub.unsubscribe(id)
}

func Capture(level zapcore.Level, source, message string, fields map[string]any) Entry {
	return defaultHub.append(Entry{
		Time:    time.Now().Format(TimeFormatDateTime),
		Level:   strings.ToLower(level.String()),
		Source:  source,
		Message: message,
		Fields:  cloneFields(fields),
	})
}

func NewMemoryCore(level zapcore.LevelEnabler) zapcore.Core {
	return &memoryCore{level: level}
}

func (c *memoryCore) Enabled(level zapcore.Level) bool {
	return c.level.Enabled(level)
}

func (c *memoryCore) With(fields []zap.Field) zapcore.Core {
	merged := make([]zap.Field, 0, len(c.fields)+len(fields))
	merged = append(merged, c.fields...)
	merged = append(merged, fields...)
	return &memoryCore{
		level:  c.level,
		fields: merged,
	}
}

func (c *memoryCore) Check(entry zapcore.Entry, checked *zapcore.CheckedEntry) *zapcore.CheckedEntry {
	if c.Enabled(entry.Level) {
		return checked.AddCore(entry, c)
	}
	return checked
}

func (c *memoryCore) Write(entry zapcore.Entry, fields []zap.Field) error {
	combined := make([]zap.Field, 0, len(c.fields)+len(fields))
	combined = append(combined, c.fields...)
	combined = append(combined, fields...)

	defaultHub.append(Entry{
		Time:    entry.Time.Format(TimeFormatDateTime),
		Level:   strings.ToLower(entry.Level.String()),
		Source:  "app",
		Message: entry.Message,
		Logger:  entry.LoggerName,
		Caller:  entry.Caller.TrimmedPath(),
		Fields:  fieldsToMap(combined),
	})

	return nil
}

func (c *memoryCore) Sync() error {
	return nil
}

func fieldsToMap(fields []zap.Field) map[string]any {
	if len(fields) == 0 {
		return nil
	}

	encoder := zapcore.NewMapObjectEncoder()
	for _, field := range fields {
		field.AddTo(encoder)
	}
	if len(encoder.Fields) == 0 {
		return nil
	}

	return cloneFields(encoder.Fields)
}

func cloneFields(fields map[string]any) map[string]any {
	if len(fields) == 0 {
		return nil
	}

	cloned := make(map[string]any, len(fields))
	for key, value := range fields {
		cloned[key] = value
	}
	return cloned
}

func SanitizeHeaders(headers http.Header) map[string][]string {
	if len(headers) == 0 {
		return nil
	}

	sanitized := make(map[string][]string, len(headers))
	for key, values := range headers {
		copied := append([]string(nil), values...)
		if isSensitiveKey(key) {
			for i := range copied {
				copied[i] = maskValue(copied[i])
			}
		}
		sanitized[key] = copied
	}
	return sanitized
}

func SanitizeBody(contentType, body string) string {
	if body == "" {
		return ""
	}

	switch {
	case strings.Contains(contentType, "application/json"):
		var payload any
		if err := json.Unmarshal([]byte(body), &payload); err == nil {
			maskValueRecursive(payload)
			if b, err := json.Marshal(payload); err == nil {
				return string(b)
			}
		}
	case strings.Contains(contentType, "application/x-www-form-urlencoded"):
		values, err := url.ParseQuery(body)
		if err == nil {
			for key := range values {
				if isSensitiveKey(key) {
					values.Set(key, maskValue(values.Get(key)))
				}
			}
			return values.Encode()
		}
	}

	return body
}

func maskValueRecursive(value any) {
	switch typed := value.(type) {
	case map[string]any:
		for key, item := range typed {
			if isSensitiveKey(key) {
				typed[key] = maskValue("")
				continue
			}
			maskValueRecursive(item)
		}
	case []any:
		for _, item := range typed {
			maskValueRecursive(item)
		}
	}
}

func isSensitiveKey(key string) bool {
	key = strings.ToLower(strings.TrimSpace(key))
	switch key {
	case "authorization", "cookie", "set-cookie", "x-session-id", "password", "token", "code", "session_id":
		return true
	default:
		return false
	}
}

func maskValue(_ string) string {
	return "******"
}