vaultik/internal/log/tty_handler.go

package log //nolint:revive,nolintlint // stdlib log unused here; see #76

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode"

	"github.com/dustin/go-humanize"
)

// groupSeparator joins an open group path to an attribute key. This
// format has no nesting, so a group becomes a dotted key prefix:
// slog.New(h).WithGroup("db").With("rows", 3) renders "db.rows=3".
const groupSeparator = "."

// bytesAttrKey is the attribute key whose int64 value is rendered as a
// human-readable byte count rather than a bare number. Keys reaching
// writeAttr are group-qualified, so the match is made against the final
// dot-separated segment: without that, a "bytes" attribute logged under
// an open group would arrive as "transfer.bytes" and silently lose its
// formatting.
const bytesAttrKey = "bytes"

// isBytesAttr reports whether a group-qualified attribute key names the
// byte-count attribute, i.e. whether its last segment is bytesAttrKey.
func isBytesAttr(key string) bool {
	if idx := strings.LastIndex(key, groupSeparator); idx >= 0 {
		key = key[idx+len(groupSeparator):]
	}

	return key == bytesAttrKey
}

// ANSI color codes
const (
	colorReset  = "\033[0m"
	colorRed    = "\033[31m"
	colorYellow = "\033[33m"
	colorBlue   = "\033[34m"
	colorGray   = "\033[90m"
	colorGreen  = "\033[32m"
	colorCyan   = "\033[36m"
	colorBold   = "\033[1m"
)

// TTYHandler is a custom slog handler for TTY output with colors.
//
// A handler and the handlers derived from it via WithAttrs/WithGroup
// all write to the same stream, so they share one mutex; that is why mu
// is a pointer. A value mutex would give every derived handler its own
// lock and stop serializing writes to the stream they have in common.
type TTYHandler struct {
	opts slog.HandlerOptions
	mu   *sync.Mutex
	out  io.Writer

	// color reports whether records are written with ANSI color. It is
	// false when NO_COLOR is set, matching internal/ui's shouldColor
	// test; the caller already chose this handler for a terminal, so
	// TTY-ness is not re-checked here.
	color bool

	// attrs are the attributes accumulated through WithAttrs, emitted
	// ahead of each record's own attributes. Their keys already carry
	// the group path that was open when they were added, so no
	// qualification happens at write time.
	attrs []slog.Attr

	// groups is the group path opened by WithGroup, applied as a key
	// prefix to attributes that arrive later — both on a record and
	// through a further WithAttrs.
	groups []string
}

// NewTTYHandler creates a new TTY handler. Color is enabled unless the
// NO_COLOR environment variable is set to a non-empty value, the same
// test internal/ui applies to status lines.
func NewTTYHandler(out io.Writer, opts *slog.HandlerOptions) *TTYHandler {
	if opts == nil {
		opts = &slog.HandlerOptions{}
	}

	return &TTYHandler{
		out:   out,
		opts:  *opts,
		mu:    &sync.Mutex{},
		color: os.Getenv("NO_COLOR") == "",
	}
}

// Enabled reports whether the handler handles records at the given level.
func (h *TTYHandler) Enabled(_ context.Context, level slog.Level) bool {
	return level >= h.opts.Level.Level()
}

// Handle writes the log record to the output with color formatting.
func (h *TTYHandler) Handle(_ context.Context, r slog.Record) error {
	h.mu.Lock()
	defer h.mu.Unlock()

	// Format timestamp
	timestamp := r.Time.Format("15:04:05")

	// Level and color
	level := r.Level.String()

	var levelColor string

	switch r.Level {
	case slog.LevelDebug:
		levelColor = colorGray
		level = "DEBUG"
	case slog.LevelInfo:
		levelColor = colorGreen
		level = "INFO "
	case slog.LevelWarn:
		levelColor = colorYellow
		level = "WARN "
	case slog.LevelError:
		levelColor = colorRed
		level = "ERROR"
	default:
		levelColor = colorReset
	}

	// Print main message. The message is escaped before the colour codes
	// are written around it: it can carry text from an untrusted source
	// (a storage error, for one), and a raw control character would
	// otherwise reach the terminal.
	_, _ = fmt.Fprintf(h.out, "%s %s %s",
		h.paint(colorGray, timestamp),
		h.paint(levelColor, level),
		h.paint(colorBold, sanitize(r.Message)))

	// Attributes carried by the handler come first, then the record's
	// own. Handler attributes were qualified when they were added; the
	// record's are qualified now, against whatever group path is open.
	for _, a := range h.attrs {
		h.writeAttr(a)
	}

	prefix := strings.Join(h.groups, groupSeparator)

	r.Attrs(func(a slog.Attr) bool {
		for _, flat := range appendAttr(nil, prefix, a) {
			h.writeAttr(flat)
		}

		return true
	})

	_, _ = fmt.Fprintln(h.out)

	return nil
}

// appendAttr flattens a into dst, folding prefix into its key and
// expanding group values into further dotted keys. Following the
// slog.Handler contract: an empty Attr is dropped, a group with no
// attributes is dropped, and a group with an empty key is inlined into
// its parent rather than contributing a level.
func appendAttr(dst []slog.Attr, prefix string, a slog.Attr) []slog.Attr {
	a.Value = a.Value.Resolve()

	if a.Equal(slog.Attr{}) {
		return dst
	}

	key := a.Key

	switch {
	case prefix == "":
		// key stands alone.
	case key == "":
		key = prefix
	default:
		key = prefix + groupSeparator + key
	}

	if a.Value.Kind() != slog.KindGroup {
		return append(dst, slog.Attr{Key: key, Value: a.Value})
	}

	for _, member := range a.Value.Group() {
		dst = appendAttr(dst, key, member)
	}

	return dst
}

// WithAttrs returns a new handler that emits attrs on every record it
// handles, in addition to whatever the handler already carried. Keys
// are qualified by the group path open at the time of the call, so
// WithGroup("db").WithAttrs(rows=3) later renders "db.rows=3".
//
// The receiver is not modified.
func (h *TTYHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	if len(attrs) == 0 {
		return h
	}

	prefix := strings.Join(h.groups, groupSeparator)
	next := h.clone()

	for _, a := range attrs {
		next.attrs = appendAttr(next.attrs, prefix, a)
	}

	return next
}

// WithGroup returns a new handler that qualifies every subsequent
// attribute key with name. This format is a single line with nowhere to
// nest, so grouping is rendered as a dotted key prefix: after
// WithGroup("db"), an attribute "rows" is emitted as "db.rows".
//
// An empty name returns the receiver unchanged, per the slog.Handler
// contract. The receiver is not modified.
func (h *TTYHandler) WithGroup(name string) slog.Handler {
	if name == "" {
		return h
	}

	next := h.clone()
	next.groups = append(next.groups, name)

	return next
}

// clone returns a copy of h that shares its output stream and mutex but
// owns its attribute and group slices.
//
// The slices are copied rather than resliced on purpose. slog permits
// one handler to be derived from concurrently, and two derivations that
// appended into a shared backing array would each overwrite the other's
// attribute — a data race with a silent wrong-output failure mode.
func (h *TTYHandler) clone() *TTYHandler {
	next := &TTYHandler{
		opts:   h.opts,
		mu:     h.mu,
		out:    h.out,
		color:  h.color,
		attrs:  make([]slog.Attr, len(h.attrs), len(h.attrs)+1),
		groups: make([]string, len(h.groups), len(h.groups)+1),
	}

	copy(next.attrs, h.attrs)
	copy(next.groups, h.groups)

	return next
}

// writeAttr renders one already-flattened, already-qualified attribute
// as " key=value". Callers hold h.mu.
func (h *TTYHandler) writeAttr(a slog.Attr) {
	value := a.Value.String()
	// Special handling for certain attribute types
	switch a.Value.Kind() {
	case slog.KindDuration:
		if d, ok := a.Value.Any().(time.Duration); ok {
			value = formatDuration(d)
		}
	case slog.KindInt64:
		if isBytesAttr(a.Key) {
			value = formatBytes(a.Value.Int64())
		}
	case slog.KindAny, slog.KindBool, slog.KindFloat64, slog.KindString,
		slog.KindTime, slog.KindUint64, slog.KindGroup, slog.KindLogValuer:
		// Plain string form above is already correct for these kinds.
	default:
		// Future kinds also use the plain string form.
	}

	// Escape the key and value before the colour codes are written around
	// them. Both can carry text from an untrusted source — a manifest
	// timestamp, a storage error, a path or symlink target read back from
	// the snapshot database — so a control character in one of them must
	// be rendered as an escape sequence rather than reaching the terminal,
	// where it could move the cursor or inject its own colours.
	_, _ = fmt.Fprintf(h.out, " %s=%s",
		h.paint(colorCyan, sanitize(a.Key)),
		h.paint(colorBlue, sanitize(value)))
}

// paint wraps s in the given ANSI color when color is enabled, and
// returns s unchanged otherwise.
func (h *TTYHandler) paint(color, s string) string {
	if !h.color {
		return s
	}

	return color + s + colorReset
}

// sanitize returns s unchanged when every rune in it is printable, and a
// double-quoted, backslash-escaped form (\n, \x1b, …) otherwise. It is
// applied to untrusted text before any colour code is written, so a
// control character can never reach the terminal raw.
func sanitize(s string) string {
	for _, r := range s {
		if !unicode.IsPrint(r) {
			return strconv.Quote(s)
		}
	}

	return s
}

// formatDuration formats a duration in a human-readable way
func formatDuration(d time.Duration) string {
	switch {
	case d < time.Millisecond:
		return fmt.Sprintf("%dµs", d.Microseconds())
	case d < time.Second:
		return fmt.Sprintf("%dms", d.Milliseconds())
	case d < time.Minute:
		return fmt.Sprintf("%.1fs", d.Seconds())
	default:
		return d.String()
	}
}

// formatBytes formats a byte count with humanize.Bytes, as every command
// prints one. humanize.Bytes takes no negative count, so one is printed
// as it is.
func formatBytes(b int64) string {
	if b < 0 {
		return fmt.Sprintf("%d B", b)
	}

	return humanize.Bytes(uint64(b))
}