starlog/redaction.go

package starlog

import (
	"context"
	"fmt"
	"regexp"
	"sync/atomic"

	"b612.me/starlog/internal/redactutil"
)

func cloneRedactRules(source []RedactRule) []RedactRule {
	if len(source) == 0 {
		return nil
	}
	cloned := make([]RedactRule, len(source))
	copy(cloned, source)
	return cloned
}

func normalizeRedactMask(mask string) string {
	return redactutil.NormalizeMask(mask)
}

func maskEntry(entry *Entry, token string) {
	if entry == nil {
		return
	}
	token = normalizeRedactMask(token)
	entry.Message = token
	entry.Err = nil
	if len(entry.Fields) == 0 {
		return
	}
	maskedFields := redactutil.MaskFields(entry.Fields, token)
	if len(maskedFields) == 0 {
		entry.Fields = nil
		return
	}
	entry.Fields = Fields(maskedFields)
}

func (logger *starlog) applyRedaction(snapshot *starlog, entry *Entry) bool {
	if entry == nil {
		return true
	}
	ctx := entry.Context
	if ctx == nil {
		ctx = context.Background()
	}
	for _, rule := range snapshot.redactRules {
		if rule == nil {
			continue
		}
		if _, err := rule.Apply(ctx, entry); err != nil {
			return logger.handleRedactionFailure(snapshot, entry, err)
		}
	}
	if snapshot.redactor != nil {
		if err := snapshot.redactor.Redact(ctx, entry); err != nil {
			return logger.handleRedactionFailure(snapshot, entry, err)
		}
	}
	return true
}

func (logger *starlog) handleRedactionFailure(snapshot *starlog, entry *Entry, err error) bool {
	reportWriteError(fmt.Errorf("%w: %v", ErrRedactionFailed, err), LogData{
		Name: snapshot.name,
		Log:  entry.Message,
	})
	atomic.AddUint64(&logger.redactErrorCount, 1)
	switch snapshot.redactFailMode {
	case RedactFailOpen:
		return true
	case RedactFailDrop:
		return false
	default:
		maskEntry(entry, snapshot.redactMaskToken)
		return true
	}
}

type RuleRedactor struct {
	rules []RedactRule
}

func NewRuleRedactor(rules ...RedactRule) *RuleRedactor {
	return &RuleRedactor{
		rules: cloneRedactRules(rules),
	}
}

func (redactor *RuleRedactor) Redact(ctx context.Context, entry *Entry) error {
	if redactor == nil || entry == nil {
		return nil
	}
	for _, rule := range redactor.rules {
		if rule == nil {
			continue
		}
		if _, err := rule.Apply(ctx, entry); err != nil {
			return err
		}
	}
	return nil
}

type SensitiveFieldRule struct {
	fields map[string]struct{}
	mask   string
}

func NewSensitiveFieldRule(mask string, fields ...string) *SensitiveFieldRule {
	return &SensitiveFieldRule{
		fields: redactutil.BuildFieldSet(fields...),
		mask:   normalizeRedactMask(mask),
	}
}

func (rule *SensitiveFieldRule) Apply(ctx context.Context, entry *Entry) (bool, error) {
	_ = ctx
	if rule == nil || entry == nil || len(entry.Fields) == 0 {
		return false, nil
	}
	changed := false
	for key, value := range entry.Fields {
		lookup := redactutil.LookupFieldKey(key)
		if _, ok := rule.fields[lookup]; !ok {
			continue
		}
		if !redactutil.IsMasked(value, rule.mask) {
			changed = true
		}
		entry.Fields[key] = rule.mask
	}
	return changed, nil
}

type MessageRegexRule struct {
	pattern     *regexp.Regexp
	replacement string
}

func NewMessageRegexRule(pattern *regexp.Regexp, replacement string) *MessageRegexRule {
	if replacement == "" {
		replacement = "[REDACTED]"
	}
	return &MessageRegexRule{
		pattern:     pattern,
		replacement: replacement,
	}
}

func (rule *MessageRegexRule) Apply(ctx context.Context, entry *Entry) (bool, error) {
	_ = ctx
	if rule == nil || rule.pattern == nil || entry == nil || entry.Message == "" {
		return false, nil
	}
	redacted, changed := redactutil.ReplaceRegex(rule.pattern, entry.Message, rule.replacement)
	if !changed {
		return false, nil
	}
	entry.Message = redacted
	return true, nil
}