Request ID Middleware and Correlated Logs in Go Services

package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"log/slog"
	"net/http"
	"regexp"
)

type ctxKey int

const (
	idKey ctxKey = iota
	loggerKey
)

const Header = "X-Request-ID"

// validID is the shape we are willing to accept from a caller: short, printable, no spaces or control characters.
var validID = regexp.MustCompile(`^[A-Za-z0-9._-]{8,64}$`)

func newID() string {
	b := make([]byte, 8)
	rand.Read(b) // crypto/rand.Read never returns an error on supported platforms
	return hex.EncodeToString(b)
}

// RequestID assigns an ID to every request, stores it in the context and echoes it in the response.
func RequestID(base *slog.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			id := r.Header.Get(Header)
			if !validID.MatchString(id) {
				id = newID() // missing, or not a value we trust
			}
			w.Header().Set(Header, id) // set before the handler writes, so every response carries it

			logger := base.With("request_id", id)
			ctx := context.WithValue(r.Context(), idKey, id)
			ctx = context.WithValue(ctx, loggerKey, logger)
			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

// IDFrom returns the request ID stored in ctx, or "" outside a request.
func IDFrom(ctx context.Context) string {
	id, _ := ctx.Value(idKey).(string)
	return id
}

// LoggerFrom returns the request's logger, which already carries the request ID.
func LoggerFrom(ctx context.Context) *slog.Logger {
	if l, ok := ctx.Value(loggerKey).(*slog.Logger); ok {
		return l
	}
	return slog.Default()
}

package main

import "net/http"

// propagate copies the request ID from the outgoing request's context onto its header.
type propagate struct{ next http.RoundTripper }

func (p propagate) RoundTrip(req *http.Request) (*http.Response, error) {
	if id := IDFrom(req.Context()); id != "" {
		req = req.Clone(req.Context()) // RoundTrippers must not modify the caller's request
		req.Header.Set(Header, id)
	}
	return p.next.RoundTrip(req)
}

// NewClient returns an HTTP client that forwards the request ID.
func NewClient() *http.Client {
	return &http.Client{Transport: propagate{http.DefaultTransport}}
}

package main

import (
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
)

func main() {
	base := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
		ReplaceAttr: func(groups []string, a slog.Attr) slog.Attr {
			if a.Key == slog.TimeKey {
				return slog.Attr{}
			}
			return a
		},
	}))
	mw := RequestID(base)

	downstream := httptest.NewServer(mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		LoggerFrom(r.Context()).Info("downstream handling")
		fmt.Fprint(w, "downstream ok")
	})))
	defer downstream.Close()

	client := NewClient()
	edge := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		LoggerFrom(r.Context()).Info("edge handling")
		req, _ := http.NewRequestWithContext(r.Context(), "GET", downstream.URL, nil) // the context carries the ID
		resp, err := client.Do(req)
		if err != nil {
			http.Error(w, "upstream failed", http.StatusBadGateway)
			return
		}
		defer resp.Body.Close()
		io.Copy(w, resp.Body)
	}))

	for _, inbound := range []string{"", "trace-from-the-gateway-42", "x\nINFO forged log line user=admin"} {
		req := httptest.NewRequest("GET", "/", nil)
		if inbound != "" {
			req.Header.Set(Header, inbound)
		}
		rec := httptest.NewRecorder()
		edge.ServeHTTP(rec, req)
		fmt.Printf("inbound %-38q -> response ID %q (replaced: %v)\n", inbound, rec.Header().Get(Header), rec.Header().Get(Header) != inbound)
	}
}

level=INFO msg="edge handling" request_id=2e9457064fd70cd6
level=INFO msg="downstream handling" request_id=2e9457064fd70cd6
inbound ""                                     -> response ID "2e9457064fd70cd6" (replaced: true)
level=INFO msg="edge handling" request_id=trace-from-the-gateway-42
level=INFO msg="downstream handling" request_id=trace-from-the-gateway-42
inbound "trace-from-the-gateway-42"            -> response ID "trace-from-the-gateway-42" (replaced: false)
level=INFO msg="edge handling" request_id=7b132326c221ce87
level=INFO msg="downstream handling" request_id=7b132326c221ce87
inbound "x\nINFO forged log line user=admin"   -> response ID "7b132326c221ce87" (replaced: true)

writeProblem(w, http.StatusInternalServerError, "internal error", map[string]any{"request_id": IDFrom(r.Context())})

package main

import (
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func run(inbound string) (ctxID, headerID string) {
	h := RequestID(slog.Default())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctxID = IDFrom(r.Context())
	}))
	req := httptest.NewRequest("GET", "/", nil)
	if inbound != "" {
		req.Header.Set(Header, inbound)
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return ctxID, rec.Header().Get(Header)
}

func TestGeneratesWhenMissing(t *testing.T) {
	ctxID, headerID := run("")
	if len(ctxID) != 16 || ctxID != headerID {
		t.Fatalf("ctx %q, header %q: want a 16 character ID, identical in both", ctxID, headerID)
	}
}

func TestKeepsAWellFormedInboundID(t *testing.T) {
	if ctxID, _ := run("gateway-id-12345"); ctxID != "gateway-id-12345" {
		t.Fatalf("the inbound ID was replaced: %q", ctxID)
	}
}

func TestRejectsAHostileInboundID(t *testing.T) {
	for _, bad := range []string{"short", strings.Repeat("a", 65), "has space in it", "line\nbreak-1234", "emoji-\U0001F600-1234"} {
		if ctxID, _ := run(bad); ctxID == bad {
			t.Errorf("accepted the hostile ID %q", bad)
		}
	}
}

func TestOutboundCallsCarryTheID(t *testing.T) {
	var seen string
	downstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { seen = r.Header.Get(Header) }))
	defer downstream.Close()
	h := RequestID(slog.Default())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		req, _ := http.NewRequestWithContext(r.Context(), "GET", downstream.URL, nil)
		NewClient().Do(req)
	}))
	req := httptest.NewRequest("GET", "/", nil)
	req.Header.Set(Header, "caller-supplied-id")
	h.ServeHTTP(httptest.NewRecorder(), req)
	if seen != "caller-supplied-id" {
		t.Fatalf("the downstream service saw %q", seen)
	}
}

=== RUN   TestGeneratesWhenMissing
--- PASS: TestGeneratesWhenMissing (0.00s)
=== RUN   TestKeepsAWellFormedInboundID
--- PASS: TestKeepsAWellFormedInboundID (0.00s)
=== RUN   TestRejectsAHostileInboundID
--- PASS: TestRejectsAHostileInboundID (0.00s)
=== RUN   TestOutboundCallsCarryTheID
--- PASS: TestOutboundCallsCarryTheID (0.00s)
PASS
ok  	example.test/lab	1.320s