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("<html>Something went wrong.</html>"))
}
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("<html>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, "<html>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")
}