package middleware import ( "bytes" "log/slog" "net" "net/http" "net/http/httptest" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // captureLog redirects the default slog logger into a buffer for the duration // of one test, both to keep the panic stacks out of the test output and so that // a test can assert on what an operator would have seen. func captureLog(t *testing.T) *bytes.Buffer { t.Helper() var buf bytes.Buffer previous := slog.Default() slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{ Level: slog.LevelDebug, // Drop the timestamp so an assertion can match a whole line. ReplaceAttr: func(_ []string, a slog.Attr) slog.Attr { if a.Key == slog.TimeKey { return slog.Attr{} } return a }, }))) t.Cleanup(func() { slog.SetDefault(previous) }) return &buf } // renderInternal is the shape a service passes in: its own error page, at 500, // ignoring the recovered value. func renderInternal(w http.ResponseWriter, _ *http.Request, _ any) { w.Header().Set("Content-Type", "text/html; charset=utf-8") w.WriteHeader(http.StatusInternalServerError) _, _ = w.Write([]byte("Something went wrong.")) } func TestPrivateCacheSetsBothHeadersBeforeTheHandlerRuns(t *testing.T) { var seen http.Header h := PrivateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { seen = w.Header().Clone() w.WriteHeader(http.StatusOK) })) rec := httptest.NewRecorder() h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/tokens", nil)) require.Equal(t, http.StatusOK, rec.Code) assert.Equal(t, "private, no-store", rec.Header().Get("Cache-Control")) assert.Equal(t, "Cookie, Authorization", rec.Header().Get("Vary")) // Before, not after: a handler that writes its own body must already have // them, or a flush would put the status line on the wire without them. assert.Equal(t, "private, no-store", seen.Get("Cache-Control")) assert.Equal(t, "Cookie, Authorization", seen.Get("Vary")) } func TestPrivateCacheLetsAHandlerOverrideThem(t *testing.T) { h := PrivateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Cache-Control", "public, max-age=31536000, immutable") w.WriteHeader(http.StatusOK) })) rec := httptest.NewRecorder() h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/static/app.abc123.css", nil)) assert.Equal(t, "public, max-age=31536000, immutable", rec.Header().Get("Cache-Control")) } func TestSetPrivateCacheWritesTheHeadersWithoutAHandler(t *testing.T) { rec := httptest.NewRecorder() SetPrivateCache(rec) assert.Equal(t, "private, no-store", rec.Header().Get("Cache-Control")) assert.Equal(t, "Cookie, Authorization", rec.Header().Get("Vary")) } func TestRecoverPanicsPassesANonPanickingHandlerThrough(t *testing.T) { rendered := false h := RecoverPanics(func(http.ResponseWriter, *http.Request, any) { rendered = true })( http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusTeapot) _, _ = w.Write([]byte("fine")) })) rec := httptest.NewRecorder() h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/", nil)) assert.Equal(t, http.StatusTeapot, rec.Code) assert.Equal(t, "fine", rec.Body.String()) assert.False(t, rendered, "the callback is for panics only") } func TestRecoverPanicsRendersTheErrorPageAndAnswers500(t *testing.T) { logged := captureLog(t) var recovered any calls := 0 h := RecoverPanics(func(w http.ResponseWriter, r *http.Request, v any) { calls++ recovered = v renderInternal(w, r, v) })(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { panic("the store is nil") })) rec := httptest.NewRecorder() require.NotPanics(t, func() { h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/tokens", nil)) }) assert.Equal(t, 1, calls) assert.Equal(t, "the store is nil", recovered, "the callback gets the recovered value") assert.Equal(t, http.StatusInternalServerError, rec.Code) assert.Contains(t, rec.Body.String(), "Something went wrong.") // The detail goes to the log, with the request around it and a stack. assert.Contains(t, logged.String(), `msg="panic serving a request"`) assert.Contains(t, logged.String(), "method=GET path=/tokens") assert.Contains(t, logged.String(), `panic="the store is nil"`) assert.Contains(t, logged.String(), "runtime/debug.Stack") assert.NotContains(t, rec.Body.String(), "the store is nil", "never to the viewer") } func TestRecoverPanicsPassesAnErrorValueThroughUnwrapped(t *testing.T) { captureLog(t) boom := &net.AddrError{Err: "boom", Addr: "nowhere"} var recovered any h := RecoverPanics(func(w http.ResponseWriter, r *http.Request, v any) { recovered = v renderInternal(w, r, v) })(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { panic(boom) })) h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/", nil)) assert.Same(t, boom, recovered, "the value arrives as it was thrown, not stringified") } func TestRecoverPanicsAbandonsAResponseThatHasStarted(t *testing.T) { logged := captureLog(t) rendered := false h := RecoverPanics(func(http.ResponseWriter, *http.Request, any) { rendered = true })( http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("half a p")) panic("template died mid-page") })) rec := httptest.NewRecorder() assert.PanicsWithValue(t, http.ErrAbortHandler, func() { h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/repos", nil)) }, "net/http is told to drop the connection") assert.False(t, rendered, "there is no status line left to render a page over") assert.Equal(t, http.StatusOK, rec.Code, "the status already sent is not rewritten") assert.Equal(t, "half a p", rec.Body.String(), "nothing is appended to the truncated body") assert.Contains(t, logged.String(), `msg="panic serving a request"`) assert.Contains(t, logged.String(), "method=GET path=/repos") assert.Contains(t, logged.String(), `panic="template died mid-page"`) } func TestRecoverPanicsCountsAFlushAsAStartedResponse(t *testing.T) { captureLog(t) rendered := false h := RecoverPanics(func(http.ResponseWriter, *http.Request, any) { rendered = true })( http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { flusher, ok := w.(http.Flusher) require.True(t, ok, "the wrapper keeps http.Flusher") flusher.Flush() panic("after the flush") })) assert.PanicsWithValue(t, http.ErrAbortHandler, func() { h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/", nil)) }) assert.False(t, rendered) } func TestRecoverPanicsRepanicsErrAbortHandler(t *testing.T) { rendered := false h := RecoverPanics(func(http.ResponseWriter, *http.Request, any) { rendered = true })( http.HandlerFunc(func(http.ResponseWriter, *http.Request) { panic(http.ErrAbortHandler) })) assert.PanicsWithValue(t, http.ErrAbortHandler, func() { h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/", nil)) }) assert.False(t, rendered, "a deliberately abandoned response is not resurrected") } func TestRecoverPanicsDoesNotLoopWhenTheErrorPagePanics(t *testing.T) { logged := captureLog(t) calls := 0 h := RecoverPanics(func(http.ResponseWriter, *http.Request, any) { calls++ panic("the chrome is broken too") })(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { panic("the store is nil") })) rec := httptest.NewRecorder() assert.PanicsWithValue(t, http.ErrAbortHandler, func() { h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/tokens", nil)) }) assert.Equal(t, 1, calls, "the error page is attempted exactly once") assert.Contains(t, logged.String(), `msg="panic rendering the error page"`) assert.Contains(t, logged.String(), `panic="the chrome is broken too"`) assert.Contains(t, logged.String(), `original_panic="the store is nil"`) } func TestRecoverPanicsForwardsErrAbortHandlerFromTheErrorPage(t *testing.T) { captureLog(t) h := RecoverPanics(func(http.ResponseWriter, *http.Request, any) { panic(http.ErrAbortHandler) })(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { panic("the store is nil") })) assert.PanicsWithValue(t, http.ErrAbortHandler, func() { h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/", nil)) }) } func TestRecoverPanicsRequiresARenderCallback(t *testing.T) { assert.PanicsWithValue(t, "middleware: RecoverPanics needs a render callback", func() { RecoverPanics(nil) }, "a wiring mistake fails at wiring time") } func TestStatusClientClosedRequestIsNotAServerError(t *testing.T) { assert.Equal(t, 499, StatusClientClosedRequest) assert.Less(t, StatusClientClosedRequest, 500, "a viewer who went away must not land in the rate an alert is written against") }