package middleware import ( "bytes" "log" "net" "net/http" "net/http/httptest" "os" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // captureLog redirects the standard 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 flags := log.Flags() log.SetOutput(&buf) log.SetFlags(0) t.Cleanup(func() { log.SetOutput(os.Stderr) log.SetFlags(flags) }) 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(), "panic serving GET /tokens: 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(), "panic serving GET /repos: 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(), "panic rendering the error page for GET /tokens: 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") }