Files

54 lines
1.3 KiB
Go

package middleware
import (
"encoding/json"
"log"
"net/http"
"runtime/debug"
"time"
)
func Logger(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
rw := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
next.ServeHTTP(rw, r)
log.Printf("%s %s %d %s", r.Method, r.URL.Path, rw.status, time.Since(start))
})
}
func Recover(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer func() {
if rec := recover(); rec != nil {
log.Printf("[panic] %v\n%s", rec, debug.Stack())
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
}
}()
next.ServeHTTP(w, r)
})
}
type statusRecorder struct {
http.ResponseWriter
status int
}
func (r *statusRecorder) WriteHeader(code int) {
r.status = code
r.ResponseWriter.WriteHeader(code)
}
// Flush 透传 http.Flusher,避免 SSE 流式响应被中间件阻断。
func (r *statusRecorder) Flush() {
if f, ok := r.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
func writeJSON(w http.ResponseWriter, status int, body any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(body)
}