Files

200 lines
5.4 KiB
Go

package middleware
import (
"bufio"
"crypto/rand"
"errors"
"log/slog"
"net"
"net/http"
"time"
"git.avsdev.uk/Go/job-search/internal/httpx"
)
// responseRecorder captures status and size while staying transparent to
// http.ResponseController (Unwrap) and to code that type-asserts for
// http.Hijacker / http.Flusher (needed by the WebSocket upgrade).
type responseRecorder struct {
http.ResponseWriter
status int
bytes int64
hijacked bool
// beforeHeader, if set, runs once just before headers are written.
beforeHeader func(http.Header)
}
// WriteHeader records the status, running beforeHeader the first time, and
// writes it.
func (w *responseRecorder) WriteHeader(code int) {
if w.status == 0 {
w.status = code
if w.beforeHeader != nil {
w.beforeHeader(w.Header())
}
}
w.ResponseWriter.WriteHeader(code)
}
// Write writes data, writing a 200 status first if none has been, and counts
// the bytes written.
func (w *responseRecorder) Write(data []byte) (int, error) {
if w.status == 0 {
w.WriteHeader(http.StatusOK)
}
written, err := w.ResponseWriter.Write(data)
w.bytes += int64(written)
return written, err
}
// Unwrap returns the wrapped ResponseWriter, for http.ResponseController.
func (w *responseRecorder) Unwrap() http.ResponseWriter {
return w.ResponseWriter
}
// Flush writes a 200 status first if none has been, then flushes the
// wrapped ResponseWriter if it can.
func (w *responseRecorder) Flush() {
if w.status == 0 {
w.WriteHeader(http.StatusOK)
}
http.NewResponseController(w.ResponseWriter).Flush() //nolint:errcheck,gosec // none returned
}
// Hijack takes over the connection, recording it as hijacked with status 101
// (switching protocols) if no status was written.
func (w *responseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
conn, readWriter, err := http.NewResponseController(w.ResponseWriter).Hijack()
if err == nil {
w.hijacked = true
if w.status == 0 {
w.status = http.StatusSwitchingProtocols
}
}
return conn, readWriter, err
}
// RequestID assigns each request an ID, exposed as X-Request-ID. An incoming
// X-Request-ID is only accepted from trusted proxies and only if well formed.
func RequestID(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
id := ""
if Client(r).ViaProxy {
if incoming := r.Header.Get("X-Request-ID"); validRequestID(incoming) {
id = incoming
}
}
if id == "" {
id = rand.Text()
}
w.Header().Set("X-Request-ID", id)
next.ServeHTTP(w, r.WithContext(httpx.WithRequestID(r.Context(), id)))
})
}
// validRequestID reports whether id is an acceptable request ID: 1 to 64
// characters, each a letter, digit, '-', '_' or '.'.
func validRequestID(id string) bool {
if id == "" || len(id) > 64 {
return false
}
for i := 0; i < len(id); i++ {
if !isRequestIDChar(id[i]) {
return false
}
}
return true
}
// isRequestIDChar reports whether char may appear in a request ID.
func isRequestIDChar(char byte) bool {
switch {
case char >= 'a' && char <= 'z',
char >= 'A' && char <= 'Z',
char >= '0' && char <= '9':
return true
}
return char == '-' || char == '_' || char == '.'
}
// Logger logs one line per request: at debug level, or at info level for a
// connection hijack (a WebSocket upgrade), as one is logged per connection
// rather than per request; a 4xx response is logged at warn level and a 5xx
// at error level. The query string is deliberately not logged, as it may
// carry secrets. It logs to slog.Default() unless given WithLogger.
func Logger(opts ...Option) func(http.Handler) http.Handler {
settings := applyOptions(opts)
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
recorder := &responseRecorder{
ResponseWriter: w,
}
defer func() {
status := recorder.status
if status == 0 {
status = http.StatusOK
}
level := slog.LevelDebug
switch {
case status >= 500:
level = slog.LevelError
case status >= 400:
level = slog.LevelWarn
case recorder.hijacked:
level = slog.LevelInfo
}
settings.logger("http").LogAttrs(
r.Context(),
level,
"request",
slog.String("method", r.Method),
slog.String("path", r.URL.Path),
slog.Int("status", status),
slog.Int64("bytes", recorder.bytes),
slog.Duration("duration", time.Since(start)),
slog.String("ip", ClientIP(r)),
slog.String("proto", r.Proto),
slog.String("request_id", httpx.RequestID(r.Context())),
slog.Bool("hijacked", recorder.hijacked),
)
}()
next.ServeHTTP(recorder, r)
})
}
}
// Recoverer turns panics into 500 responses and logs them. It logs to
// slog.Default() unless given WithLogger.
func Recoverer(opts ...Option) func(http.Handler) http.Handler {
settings := applyOptions(opts)
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer func() {
if recovered := recover(); recovered != nil {
err, isErr := recovered.(error)
if isErr && errors.Is(err, http.ErrAbortHandler) {
panic(recovered)
}
settings.logger("http").Error(
"panic",
"error", recovered,
"path", r.URL.Path,
"request_id", httpx.RequestID(r.Context()),
)
httpx.WriteProblem(w, r, http.StatusInternalServerError, "")
}
}()
next.ServeHTTP(w, r)
})
}
}