logging: Add traceID field to access logs when tracing is active (#5507)

Co-authored-by: Francis Lavoie <lavofr@gmail.com>

7 files changed, 236 insertions, 53 deletions

diff --git a/go.mod b/go.mod
index ea8a045..eaff3a5 100644
--- a/go.mod
+++ b/go.mod
@@ -23,6 +23,7 @@ require (
 	github.com/smallstep/truststore v0.12.1
 	github.com/spf13/cobra v1.6.1
 	github.com/spf13/pflag v1.0.5
+	github.com/stretchr/testify v1.8.2
 	github.com/tailscale/tscert v0.0.0-20230124224810-c6dc1f4049b2
 	github.com/yuin/goldmark v1.5.4
 	github.com/yuin/goldmark-highlighting/v2 v2.0.0-20220924101305-151362477c87
@@ -50,6 +51,7 @@ require (
 	github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.7.0 // indirect
 	github.com/onsi/ginkgo/v2 v2.2.0 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/quic-go/qpack v0.4.0 // indirect
 	github.com/quic-go/qtls-go1-19 v0.2.1 // indirect
 	github.com/quic-go/qtls-go1-20 v0.1.1 // indirect
@@ -128,7 +130,7 @@ require (
 	go.opentelemetry.io/otel/exporters/otlp/internal/retry v1.14.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.14.0 // indirect
 	go.opentelemetry.io/otel/metric v0.37.0 // indirect
-	go.opentelemetry.io/otel/trace v1.14.0 // indirect
+	go.opentelemetry.io/otel/trace v1.14.0
 	go.opentelemetry.io/proto/otlp v0.19.0 // indirect
 	go.step.sm/cli-utils v0.7.5 // indirect
 	go.step.sm/crypto v0.23.2
diff --git a/go.sum b/go.sum
index 6e8abe6..c19dec2 100644
--- a/go.sum
+++ b/go.sum
@@ -659,6 +659,7 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
+github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/tailscale/tscert v0.0.0-20230124224810-c6dc1f4049b2 h1:TrgfmCXwtWyFw85UkRGXt9qZRzdzt3nWt2Rerdecn0w=
 github.com/tailscale/tscert v0.0.0-20230124224810-c6dc1f4049b2/go.mod h1:kNGUQ3VESx3VZwRwA9MSCUegIl6+saPL8Noq82ozCaU=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20170815181823-89b8d40f7ca8/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U=
diff --git a/modules/caddyhttp/logging.go b/modules/caddyhttp/logging.go
index 4faaec7..6bfeb51 100644
--- a/modules/caddyhttp/logging.go
+++ b/modules/caddyhttp/logging.go
@@ -20,6 +20,7 @@ import (
 	"net/http"
 	"strings"
 
+	"github.com/caddyserver/caddy/v2"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
@@ -139,6 +140,21 @@ func errLogValues(err error) (status int, msg string, fields []zapcore.Field) {
 	return
 }
 
-// Variable name used to indicate that this request
-// should be omitted from the access logs
-const SkipLogVar = "skip_log"
+// ExtraLogFields is a list of extra fields to log with every request.
+type ExtraLogFields struct {
+	fields []zapcore.Field
+}
+
+// Add adds a field to the list of extra fields to log.
+func (e *ExtraLogFields) Add(field zap.Field) {
+	e.fields = append(e.fields, field)
+}
+
+const (
+	// Variable name used to indicate that this request
+	// should be omitted from the access logs
+	SkipLogVar string = "skip_log"
+
+	// For adding additional fields to the access logs
+	ExtraLogFieldsCtxKey caddy.CtxKey = "extra_log_fields"
+)
diff --git a/modules/caddyhttp/reverseproxy/fastcgi/fastcgi.go b/modules/caddyhttp/reverseproxy/fastcgi/fastcgi.go
index ec194e7..5160067 100644
--- a/modules/caddyhttp/reverseproxy/fastcgi/fastcgi.go
+++ b/modules/caddyhttp/reverseproxy/fastcgi/fastcgi.go
@@ -254,9 +254,7 @@ func (t Transport) buildEnv(r *http.Request) (envVars, error) {
 	// if we didn't get a split result here.
 	// See https://github.com/caddyserver/caddy/issues/3718
 	if pathInfo == "" {
-		if remainder, ok := repl.GetString("http.matchers.file.remainder"); ok {
-			pathInfo = remainder
-		}
+		pathInfo, _ = repl.GetString("http.matchers.file.remainder")
 	}
 
 	// SCRIPT_FILENAME is the absolute path of SCRIPT_NAME
@@ -286,10 +284,7 @@ func (t Transport) buildEnv(r *http.Request) (envVars, error) {
 		reqHost = r.Host
 	}
 
-	authUser := ""
-	if val, ok := repl.Get("http.auth.user.id"); ok {
-		authUser = val.(string)
-	}
+	authUser, _ := repl.GetString("http.auth.user.id")
 
 	// Some variables are unused but cleared explicitly to prevent
 	// the parent environment from interfering.
diff --git a/modules/caddyhttp/server.go b/modules/caddyhttp/server.go
index 9721007..fb71b1b 100644
--- a/modules/caddyhttp/server.go
+++ b/modules/caddyhttp/server.go
@@ -282,46 +282,7 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		// capture the original version of the request
 		accLog := s.accessLogger.With(loggableReq)
 
-		defer func() {
-			// this request may be flagged as omitted from the logs
-			if skipLog, ok := GetVar(r.Context(), SkipLogVar).(bool); ok && skipLog {
-				return
-			}
-
-			repl.Set("http.response.status", wrec.Status()) // will be 0 if no response is written by us (Go will write 200 to client)
-			repl.Set("http.response.size", wrec.Size())
-			repl.Set("http.response.duration", duration)
-			repl.Set("http.response.duration_ms", duration.Seconds()*1e3) // multiply seconds to preserve decimal (see #4666)
-
-			logger := accLog
-			if s.Logs != nil {
-				logger = s.Logs.wrapLogger(logger, r.Host)
-			}
-
-			log := logger.Info
-			if wrec.Status() >= 400 {
-				log = logger.Error
-			}
-
-			userID, _ := repl.GetString("http.auth.user.id")
-
-			reqBodyLength := 0
-			if bodyReader != nil {
-				reqBodyLength = bodyReader.Length
-			}
-
-			log("handled request",
-				zap.Int("bytes_read", reqBodyLength),
-				zap.String("user_id", userID),
-				zap.Duration("duration", duration),
-				zap.Int("size", wrec.Size()),
-				zap.Int("status", wrec.Status()),
-				zap.Object("resp_headers", LoggableHTTPHeader{
-					Header:               wrec.Header(),
-					ShouldLogCredentials: shouldLogCredentials,
-				}),
-			)
-		}()
+		defer s.logRequest(accLog, r, wrec, &duration, repl, bodyReader, shouldLogCredentials)
 	}
 
 	start := time.Now()
@@ -703,6 +664,57 @@ func (s *Server) shouldLogRequest(r *http.Request) bool {
 	return !s.Logs.SkipUnmappedHosts
 }
 
+// logRequest logs the request to access logs, unless skipped.
+func (s *Server) logRequest(
+	accLog *zap.Logger, r *http.Request, wrec ResponseRecorder, duration *time.Duration,
+	repl *caddy.Replacer, bodyReader *lengthReader, shouldLogCredentials bool,
+) {
+	// this request may be flagged as omitted from the logs
+	if skipLog, ok := GetVar(r.Context(), SkipLogVar).(bool); ok && skipLog {
+		return
+	}
+
+	repl.Set("http.response.status", wrec.Status()) // will be 0 if no response is written by us (Go will write 200 to client)
+	repl.Set("http.response.size", wrec.Size())
+	repl.Set("http.response.duration", duration)
+	repl.Set("http.response.duration_ms", duration.Seconds()*1e3) // multiply seconds to preserve decimal (see #4666)
+
+	logger := accLog
+	if s.Logs != nil {
+		logger = s.Logs.wrapLogger(logger, r.Host)
+	}
+
+	log := logger.Info
+	if wrec.Status() >= 400 {
+		log = logger.Error
+	}
+
+	userID, _ := repl.GetString("http.auth.user.id")
+
+	reqBodyLength := 0
+	if bodyReader != nil {
+		reqBodyLength = bodyReader.Length
+	}
+
+	extra := r.Context().Value(ExtraLogFieldsCtxKey).(*ExtraLogFields)
+
+	fieldCount := 6
+	fields := make([]zapcore.Field, 0, fieldCount+len(extra.fields))
+	fields = append(fields,
+		zap.Int("bytes_read", reqBodyLength),
+		zap.String("user_id", userID),
+		zap.Duration("duration", *duration),
+		zap.Int("size", wrec.Size()),
+		zap.Int("status", wrec.Status()),
+		zap.Object("resp_headers", LoggableHTTPHeader{
+			Header:               wrec.Header(),
+			ShouldLogCredentials: shouldLogCredentials,
+		}))
+	fields = append(fields, extra.fields...)
+
+	log("handled request", fields...)
+}
+
 // protocol returns true if the protocol proto is configured/enabled.
 func (s *Server) protocol(proto string) bool {
 	for _, p := range s.Protocols {
@@ -738,6 +750,9 @@ func PrepareRequest(r *http.Request, repl *caddy.Replacer, w http.ResponseWriter
 
 	var url2 url.URL // avoid letting this escape to the heap
 	ctx = context.WithValue(ctx, OriginalRequestCtxKey, originalRequest(r, &url2))
+
+	ctx = context.WithValue(ctx, ExtraLogFieldsCtxKey, new(ExtraLogFields))
+
 	r = r.WithContext(ctx)
 
 	// once the pointer to the request won't change
diff --git a/modules/caddyhttp/server_test.go b/modules/caddyhttp/server_test.go
new file mode 100644
index 0000000..96a241b
--- /dev/null
+++ b/modules/caddyhttp/server_test.go
@@ -0,0 +1,146 @@
+package caddyhttp
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+)
+
+type writeFunc func(p []byte) (int, error)
+
+type nopSyncer writeFunc
+
+func (n nopSyncer) Write(p []byte) (int, error) {
+	return n(p)
+}
+
+func (n nopSyncer) Sync() error {
+	return nil
+}
+
+// testLogger returns a logger and a buffer to which the logger writes. The
+// buffer can be read for asserting log output.
+func testLogger(wf writeFunc) *zap.Logger {
+	ws := nopSyncer(wf)
+	encoderCfg := zapcore.EncoderConfig{
+		MessageKey:     "msg",
+		LevelKey:       "level",
+		NameKey:        "logger",
+		EncodeLevel:    zapcore.LowercaseLevelEncoder,
+		EncodeTime:     zapcore.ISO8601TimeEncoder,
+		EncodeDuration: zapcore.StringDurationEncoder,
+	}
+	core := zapcore.NewCore(zapcore.NewJSONEncoder(encoderCfg), ws, zap.DebugLevel)
+
+	return zap.New(core)
+}
+
+func TestServer_LogRequest(t *testing.T) {
+	s := &Server{}
+
+	ctx := context.Background()
+	ctx = context.WithValue(ctx, ExtraLogFieldsCtxKey, new(ExtraLogFields))
+	req := httptest.NewRequest(http.MethodGet, "/", nil).WithContext(ctx)
+	rec := httptest.NewRecorder()
+	wrec := NewResponseRecorder(rec, nil, nil)
+
+	duration := 50 * time.Millisecond
+	repl := NewTestReplacer(req)
+	bodyReader := &lengthReader{Source: req.Body}
+	shouldLogCredentials := false
+
+	buf := bytes.Buffer{}
+	accLog := testLogger(buf.Write)
+	s.logRequest(accLog, req, wrec, &duration, repl, bodyReader, shouldLogCredentials)
+
+	assert.JSONEq(t, `{
+		"msg":"handled request", "level":"info", "bytes_read":0,
+		"duration":"50ms", "resp_headers": {}, "size":0,
+		"status":0, "user_id":""
+	}`, buf.String())
+}
+
+func TestServer_LogRequest_WithTraceID(t *testing.T) {
+	s := &Server{}
+
+	extra := new(ExtraLogFields)
+	ctx := context.WithValue(context.Background(), ExtraLogFieldsCtxKey, extra)
+	extra.Add(zap.String("traceID", "1234567890abcdef"))
+
+	req := httptest.NewRequest(http.MethodGet, "/", nil).WithContext(ctx)
+	rec := httptest.NewRecorder()
+	wrec := NewResponseRecorder(rec, nil, nil)
+
+	duration := 50 * time.Millisecond
+	repl := NewTestReplacer(req)
+	bodyReader := &lengthReader{Source: req.Body}
+	shouldLogCredentials := false
+
+	buf := bytes.Buffer{}
+	accLog := testLogger(buf.Write)
+	s.logRequest(accLog, req, wrec, &duration, repl, bodyReader, shouldLogCredentials)
+
+	assert.JSONEq(t, `{
+		"msg":"handled request", "level":"info", "bytes_read":0,
+		"duration":"50ms", "resp_headers": {}, "size":0,
+		"status":0, "user_id":"",
+		"traceID":"1234567890abcdef"
+	}`, buf.String())
+}
+
+func BenchmarkServer_LogRequest(b *testing.B) {
+	s := &Server{}
+
+	extra := new(ExtraLogFields)
+	ctx := context.WithValue(context.Background(), ExtraLogFieldsCtxKey, extra)
+
+	req := httptest.NewRequest(http.MethodGet, "/", nil).WithContext(ctx)
+	rec := httptest.NewRecorder()
+	wrec := NewResponseRecorder(rec, nil, nil)
+
+	duration := 50 * time.Millisecond
+	repl := NewTestReplacer(req)
+	bodyReader := &lengthReader{Source: req.Body}
+
+	buf := io.Discard
+	accLog := testLogger(buf.Write)
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		s.logRequest(accLog, req, wrec, &duration, repl, bodyReader, false)
+	}
+}
+
+func BenchmarkServer_LogRequest_WithTraceID(b *testing.B) {
+	s := &Server{}
+
+	extra := new(ExtraLogFields)
+	ctx := context.WithValue(context.Background(), ExtraLogFieldsCtxKey, extra)
+	extra.Add(zap.String("traceID", "1234567890abcdef"))
+
+	req := httptest.NewRequest(http.MethodGet, "/", nil).WithContext(ctx)
+	rec := httptest.NewRecorder()
+	wrec := NewResponseRecorder(rec, nil, nil)
+
+	duration := 50 * time.Millisecond
+	repl := NewTestReplacer(req)
+	bodyReader := &lengthReader{Source: req.Body}
+
+	buf := io.Discard
+	accLog := testLogger(buf.Write)
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		s.logRequest(accLog, req, wrec, &duration, repl, bodyReader, false)
+	}
+}
diff --git a/modules/caddyhttp/tracing/tracer.go b/modules/caddyhttp/tracing/tracer.go
index cb72743..c6dd7aa 100644
--- a/modules/caddyhttp/tracing/tracer.go
+++ b/modules/caddyhttp/tracing/tracer.go
@@ -14,6 +14,7 @@ import (
 	"go.opentelemetry.io/otel/sdk/resource"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.17.0"
+	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
 )
 
@@ -81,8 +82,15 @@ func newOpenTelemetryWrapper(
 
 // serveHTTP injects a tracing context and call the next handler.
 func (ot *openTelemetryWrapper) serveHTTP(w http.ResponseWriter, r *http.Request) {
-	ot.propagators.Inject(r.Context(), propagation.HeaderCarrier(r.Header))
-	next := r.Context().Value(nextCallCtxKey).(*nextCall)
+	ctx := r.Context()
+	ot.propagators.Inject(ctx, propagation.HeaderCarrier(r.Header))
+	spanCtx := trace.SpanContextFromContext(ctx)
+	if spanCtx.IsValid() {
+		if extra, ok := ctx.Value(caddyhttp.ExtraLogFieldsCtxKey).(*caddyhttp.ExtraLogFields); ok {
+			extra.Add(zap.String("traceID", spanCtx.TraceID().String()))
+		}
+	}
+	next := ctx.Value(nextCallCtxKey).(*nextCall)
 	next.err = next.next.ServeHTTP(w, r)
 }