log_output.go

package main

import (
	"context"
	"io/ioutil"
	"regexp"
	"strings"
	"time"

	"github.com/Nitro/sidecar-executor/loghooks"
	"github.com/Shimmur/logtailer/reporter"
	limiter "github.com/sethvargo/go-limiter"
	"github.com/sethvargo/go-limiter/memorystore"
	log "github.com/sirupsen/logrus"
)

// LogFormat represents the type of log format for level extraction
type LogFormat int

const (
	LogFormatGo LogFormat = iota
	LogFormatJSON
	LogFormatTab
	LogFormatRuby
)

// logFormatRegexes maps LogFormat enum values to their compiled regex patterns
// for extracting log level from different log formats.
var logFormatRegexes = map[LogFormat]*regexp.Regexp{
	LogFormatGo:   regexp.MustCompile(`level="?([a-zA-Z]+)"?`),                   // logfmt-style: level=info
	LogFormatJSON: regexp.MustCompile(`"level"\s*:\s*"([a-zA-Z]+)"`),             // JSON: "level":"info"
	LogFormatTab:  regexp.MustCompile(`^\S+\s+([A-Z]+)\s+`),                      // Tab-separated: INFO
	LogFormatRuby: regexp.MustCompile(`\]\s+(INFO|DEBUG|WARN|ERROR|FATAL)\s+--`), // Ruby Logger
}

// logFormatNames maps LogFormat enum values to their string names for logging purposes
var logFormatNames = map[LogFormat]string{
	LogFormatGo:   "go-logfmt",
	LogFormatJSON: "json",
	LogFormatTab:  "tab",
	LogFormatRuby: "ruby",
}

// parseLogFormat converts a string annotation to a LogFormat enum value.
// Empty string defaults to LogFormatGo. Invalid formats log a warning and fall back to LogFormatGo.
func parseLogFormat(logFormat string) LogFormat {
	chosenFormat := LogFormatGo
	switch logFormat {
	case "json":
		chosenFormat = LogFormatJSON
	case "tab":
		chosenFormat = LogFormatTab
	case "ruby":
		chosenFormat = LogFormatRuby
	case "go-logfmt", "":
		chosenFormat = LogFormatGo
	default:
		log.Warnf("Unsupported log format '%s', falling back to 'go-logfmt'", logFormat)
		chosenFormat = LogFormatGo
	}
	return chosenFormat
}

type LogLine struct {
	Text      string // The log line
	Container string // The container name it came from
}

type LogOutput interface {
	Log(line *LogLine)
	Stop()
}

type UDPSyslogger struct {
	syslogger                  *log.Entry
	enableRegexLogLevelParsing bool
	logFormat                  LogFormat
	logRegex                   *regexp.Regexp
}

// selectLogRegex returns the appropriate regex pattern based on the LogFormat enum.
// This should be called once during initialization rather than on every log line.
func selectLogRegex(logFormat LogFormat) *regexp.Regexp {
	return logFormatRegexes[logFormat]
}

// extractLogLevelWithRegex attempts to extract the log level using a pre-selected regex pattern.
// This is the optimized version used internally by UDPSyslogger.
// It returns the level string (e.g., "info", "error", "warning") and a boolean indicating
// whether a level was found.
func extractLogLevelWithRegex(logLine string, regex *regexp.Regexp) (string, bool) {
	matches := regex.FindStringSubmatch(logLine)
	if len(matches) >= 2 {
		return strings.ToLower(matches[1]), true
	}
	return "", false
}

func NewUDPSyslogger(labels map[string]string, address string, enableRegexLogLevelParsing bool, logFormat string) *UDPSyslogger {
	syslogger := log.New()

	// We relay UDP syslog because we don't plan to ship it off the box and
	// because it's simplest since there is no backpressure issue to deal with.
	hook, err := loghooks.NewUDPHook(address)
	if err != nil {
		log.Errorf("Error adding hook: %s", err)
	}

	syslogger.Hooks.Add(hook)
	syslogger.SetFormatter(&log.JSONFormatter{
		FieldMap: log.FieldMap{
			log.FieldKeyTime:  "Timestamp",
			log.FieldKeyLevel: "Level",
			log.FieldKeyMsg:   "Payload",
			log.FieldKeyFunc:  "Func",
		},
	})
	syslogger.SetOutput(ioutil.Discard)

	// Parse the string annotation to enum
	parsedFormat := parseLogFormat(logFormat)

	// Add four to the labels length to account for hostname, etc
	fields := make(log.Fields, len(labels)+4)

	// Loop through the fields we're supposed to pass, and add them
	for field, val := range labels {
		fields[field] = val
	}

	return &UDPSyslogger{
		syslogger:                  syslogger.WithFields(fields),
		enableRegexLogLevelParsing: enableRegexLogLevelParsing,
		logFormat:                  parsedFormat,
		logRegex:                   selectLogRegex(parsedFormat),
	}
}

// relayLogs will watch a container and send the logs to Syslog
func (sysl *UDPSyslogger) Log(line *LogLine) {
	// Log lines all start like:
	// 2022-12-03T16:09:51.741778906Z stdout F

	// Wasn't a K8s log line!
	if len(line.Text) < 41 {
		return
	}

	k8sFields := strings.Split(line.Text[0:40], " ")
	descriptor := k8sFields[1]

	lineTxt := line.Text

	// Strip the K8s logging stuff from the log. Because the timestamp length
	// changes sometimes, we check this. It's cheaper than a split on the full
	// log line.
	if lineTxt[39] == ' ' {
		lineTxt = lineTxt[40:len(lineTxt)]
	} else {
		lineTxt = lineTxt[39:len(lineTxt)]
	}

	fields := log.Fields{
		"Container": line.Container,
	}
	if sysl.enableRegexLogLevelParsing {
		fields["LogFormat"] = logFormatNames[sysl.logFormat]
	}
	logger := sysl.syslogger.WithFields(fields)

	// If regex log level parsing is enabled, try to extract the log level
	// from structured logs (e.g., level=info)
	if sysl.enableRegexLogLevelParsing {
		if level, found := extractLogLevelWithRegex(lineTxt, sysl.logRegex); found {
			// Map to Error, Warn or Info based on severity
			switch level {
			case "panic", "fatal", "error":
				logger.Error(lineTxt)
			case "warning", "warn":
				logger.Warn(lineTxt)
			default:
				// info, debug, trace all go to Info
				logger.Info(lineTxt)
			}
			return
		}
	}

	// Fallback to heuristic detection (a la sidecar-executor)
	// This is used when enhanced extraction is disabled or no structured level was found
	lowerLine := strings.ToLower(lineTxt)
	if descriptor == "stderr" || strings.Contains(lowerLine, "error") {
		logger.Error(lineTxt)
		return
	}

	// Support warning level by line-scraping as well
	if strings.Contains(lowerLine, "warn") {
		logger.Warn(lineTxt)
		return
	}

	logger.Info(lineTxt)
}

// Stop would clean up any resources if we needed to manage any
func (sysl *UDPSyslogger) Stop() { /* noop */ }

// A RateLimitingLogger is a LogOutput that wraps another LogOutput, adding rate limiting
// capability
type RateLimitingLogger struct {
	limitStore    limiter.Store
	limitReporter *reporter.LimitExceededReporter
	output        LogOutput
	limitKey      string
}

func NewRateLimitingLogger(
	limitReporter *reporter.LimitExceededReporter, tokenLimit int,
	reportInterval time.Duration, key string, output LogOutput) *RateLimitingLogger {

	// Set up the rate limiter
	store, err := memorystore.New(&memorystore.Config{
		// Number of tokens allowed per interval.
		Tokens: uint64(tokenLimit),

		// Interval until tokens reset.
		Interval: reportInterval,
	})

	if err != nil {
		log.Errorf("Unable to create memory store: %s", err)
	}

	return &RateLimitingLogger{
		limitStore:    store,
		limitReporter: limitReporter,
		output:        output,
		limitKey:      key,
	}
}

// isRateLimited compares the tracking key to the stored limit and returns
// a boolen value for whether or not it is limited.
func (logger *RateLimitingLogger) isRateLimited() bool {
	// See if we're going to rate limit this
	limit, remaining, reset, ok, err := logger.limitStore.Take(context.Background(), logger.limitKey)
	log.Debugf("Checking rate limit: %d %d %d %t", limit, remaining, reset, ok)
	if err != nil {
		log.Warnf("Unable to fetch rate limit for %v", logger.limitKey)
		return true // Rate limit it since we can't track
	}

	return !ok
}

// Log is a pass-through to the downstream LogOutput, but checks rate limiting status
func (logger *RateLimitingLogger) Log(line *LogLine) {
	if !logger.isRateLimited() {
		logger.output.Log(line)
		return
	}

	logger.limitReporter.Incr()
}

// Stop cleans up our resources on shutdown
func (logger *RateLimitingLogger) Stop() {
	logger.limitStore.Close(context.Background())
}