feat: custom slog handlers for modules (log contextual data)

Author: dunglas

context.go

@@ -21,12 +21,14 @@ import (
 	"log"
 	"log/slog"
 	"reflect"
+	"sync"
 
 	"github.com/caddyserver/certmagic"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/collectors"
 	"go.uber.org/zap"
 	"go.uber.org/zap/exp/zapslog"
+	"go.uber.org/zap/zapcore"
 
 	"github.com/caddyserver/caddy/v2/internal/filesystems"
 )
@@ -583,24 +585,57 @@ func (ctx Context) Logger(module ...Module) *zap.Logger {
 	return ctx.cfg.Logging.Logger(mod)
 }
 
+type slogHandlerFactory func(handler slog.Handler, core zapcore.Core, moduleID string) slog.Handler
+
+var (
+	slogHandlerFactories   []slogHandlerFactory
+	slogHandlerFactoriesMu sync.RWMutex
+)
+
+// RegisterSlogHandlerFactory allows modules to register custom log/slog.Handler,
+// for instance, to add contextual data to the logs.
+func RegisterSlogHandlerFactory(factory slogHandlerFactory) {
+	slogHandlerFactoriesMu.Lock()
+	slogHandlerFactories = append(slogHandlerFactories, factory)
+	slogHandlerFactoriesMu.Unlock()
+}
+
 // Slogger returns a slog logger that is intended for use by
 // the most recent module associated with the context.
 func (ctx Context) Slogger() *slog.Logger {
+	var (
+		handler  slog.Handler
+		core     zapcore.Core
+		moduleID string
+	)
 	if ctx.cfg == nil {
 		// often the case in tests; just use a dev logger
 		l, err := zap.NewDevelopment()
 		if err != nil {
 			panic("config missing, unable to create dev logger: " + err.Error())
 		}
-		return slog.New(zapslog.NewHandler(l.Core()))
+
+		core = l.Core()
+		handler = zapslog.NewHandler(core)
+	} else {
+		mod := ctx.Module()
+		if mod == nil {
+			core = Log().Core()
+			handler = zapslog.NewHandler(core)
+		} else {
+			moduleID = string(mod.CaddyModule().ID)
+			core = ctx.cfg.Logging.Logger(mod).Core()
+			handler = zapslog.NewHandler(core, zapslog.WithName(moduleID))
+		}
 	}
-	mod := ctx.Module()
-	if mod == nil {
-		return slog.New(zapslog.NewHandler(Log().Core()))
+
+	slogHandlerFactoriesMu.RLock()
+	for _, f := range slogHandlerFactories {
+		handler = f(handler, core, moduleID)
 	}
-	return slog.New(zapslog.NewHandler(ctx.cfg.Logging.Logger(mod).Core(),
-		zapslog.WithName(string(mod.CaddyModule().ID)),
-	))
+	slogHandlerFactoriesMu.RUnlock()
+
+	return slog.New(handler)
 }
 
 // Modules returns the lineage of modules that this context provisioned,

modules/caddyhttp/logging.go

@@ -15,18 +15,28 @@
 package caddyhttp
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
+	"log/slog"
 	"net"
 	"net/http"
 	"strings"
+	"sync"
 
 	"go.uber.org/zap"
+	"go.uber.org/zap/exp/zapslog"
 	"go.uber.org/zap/zapcore"
 
 	"github.com/caddyserver/caddy/v2"
 )
 
+func init() {
+	caddy.RegisterSlogHandlerFactory(func(handler slog.Handler, core zapcore.Core, moduleID string) slog.Handler {
+		return &extraFieldsSlogHandler{defaultHandler: handler, core: core, moduleID: moduleID}
+	})
+}
+
 // ServerLogConfig describes a server's logging configuration. If
 // enabled without customization, all requests to this server are
 // logged to the default logger; logger destinations may be
@@ -223,17 +233,21 @@ func errLogValues(err error) (status int, msg string, fields func() []zapcore.Fi
 
 // ExtraLogFields is a list of extra fields to log with every request.
 type ExtraLogFields struct {
-	fields []zapcore.Field
+	fields   []zapcore.Field
+	handlers sync.Map
 }
 
 // Add adds a field to the list of extra fields to log.
 func (e *ExtraLogFields) Add(field zap.Field) {
+	e.handlers.Clear()
 	e.fields = append(e.fields, field)
 }
 
 // Set sets a field in the list of extra fields to log.
 // If the field already exists, it is replaced.
 func (e *ExtraLogFields) Set(field zap.Field) {
+	e.handlers.Clear()
+
 	for i := range e.fields {
 		if e.fields[i].Key == field.Key {
 			e.fields[i] = field
@@ -243,6 +257,29 @@ func (e *ExtraLogFields) Set(field zap.Field) {
 	e.fields = append(e.fields, field)
 }
 
+func (e *ExtraLogFields) getSloggerHandler(handler *extraFieldsSlogHandler) (h slog.Handler) {
+	if existing, ok := e.handlers.Load(handler); ok {
+		return existing.(slog.Handler)
+	}
+
+	if handler.moduleID == "" {
+		h = zapslog.NewHandler(handler.core.With(e.fields))
+	} else {
+		h = zapslog.NewHandler(handler.core.With(e.fields), zapslog.WithName(handler.moduleID))
+	}
+
+	if handler.group != "" {
+		h = h.WithGroup(handler.group)
+	}
+	if handler.attrs != nil {
+		h = h.WithAttrs(handler.attrs)
+	}
+
+	e.handlers.Store(handler, h)
+
+	return h
+}
+
 const (
 	// Variable name used to indicate that this request
 	// should be omitted from the access logs
@@ -254,3 +291,43 @@ const (
 	// Variable name used to indicate the logger to be used
 	AccessLoggerNameVarKey string = "access_logger_names"
 )
+
+type extraFieldsSlogHandler struct {
+	defaultHandler slog.Handler
+	core           zapcore.Core
+	moduleID       string
+	group          string
+	attrs          []slog.Attr
+}
+
+func (e *extraFieldsSlogHandler) Enabled(ctx context.Context, level slog.Level) bool {
+	return e.defaultHandler.Enabled(ctx, level)
+}
+
+func (e *extraFieldsSlogHandler) Handle(ctx context.Context, record slog.Record) error {
+	if elf, ok := ctx.Value(ExtraLogFieldsCtxKey).(*ExtraLogFields); ok {
+		return elf.getSloggerHandler(e).Handle(ctx, record)
+	}
+
+	return e.defaultHandler.Handle(ctx, record)
+}
+
+func (e *extraFieldsSlogHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
+	return &extraFieldsSlogHandler{
+		e.defaultHandler.WithAttrs(attrs),
+		e.core,
+		e.moduleID,
+		e.group,
+		append(e.attrs, attrs...),
+	}
+}
+
+func (e *extraFieldsSlogHandler) WithGroup(name string) slog.Handler {
+	return &extraFieldsSlogHandler{
+		e.defaultHandler.WithGroup(name),
+		e.core,
+		e.moduleID,
+		name,
+		e.attrs,
+	}
+}

modules/caddyhttp/server.go

@@ -793,8 +793,10 @@ func (s *Server) logRequest(
 	accLog *zap.Logger, r *http.Request, wrec ResponseRecorder, duration *time.Duration,
 	repl *caddy.Replacer, bodyReader *lengthReader, shouldLogCredentials bool,
 ) {
+	ctx := r.Context()
+
 	// this request may be flagged as omitted from the logs
-	if skip, ok := GetVar(r.Context(), LogSkipVar).(bool); ok && skip {
+	if skip, ok := GetVar(ctx, LogSkipVar).(bool); ok && skip {
 		return
 	}
 
@@ -812,7 +814,7 @@ func (s *Server) logRequest(
 	}
 
 	message := "handled request"
-	if nop, ok := GetVar(r.Context(), "unhandled").(bool); ok && nop {
+	if nop, ok := GetVar(ctx, "unhandled").(bool); ok && nop {
 		message = "NOP"
 	}
 
@@ -836,7 +838,7 @@ func (s *Server) logRequest(
 				reqBodyLength = bodyReader.Length
 			}
 
-			extra := r.Context().Value(ExtraLogFieldsCtxKey).(*ExtraLogFields)
+			extra := ctx.Value(ExtraLogFieldsCtxKey).(*ExtraLogFields)
 
 			fieldCount := 6
 			fields = make([]zapcore.Field, 0, fieldCount+len(extra.fields))