-
Notifications
You must be signed in to change notification settings - Fork 1
/
middleware.go
103 lines (97 loc) · 2.87 KB
/
middleware.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
package main
import (
"fmt"
"net/http"
"net/http/httptest"
"runtime/pprof"
"time"
"github.com/go-chi/chi/v5/middleware"
"github.com/rs/zerolog"
"github.com/subtle-byte/ghloc/internal/util"
)
func NewDebugMiddleware(debugToken string) func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
if debugToken == "" {
return http.HandlerFunc(http.NotFound)
}
fn := func(w http.ResponseWriter, r *http.Request) {
if r.FormValue("debug_token") == debugToken {
w.Header().Set("Content-Type", "application/octet-stream")
w.Header().Set("Content-Disposition", `attachment; filename="profile"`)
if err := pprof.StartCPUProfile(w); err != nil {
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
w.Header().Del("Content-Disposition")
w.Header().Set("X-Go-Pprof", "1")
w.WriteHeader(http.StatusInternalServerError)
fmt.Fprintf(w, "Could not enable CPU profiling: %s\n", err)
return
}
rr := httptest.ResponseRecorder{}
next.ServeHTTP(&rr, r)
pprof.StopCPUProfile()
} else {
http.NotFound(w, r)
}
}
return http.HandlerFunc(fn)
}
}
func NewLoggerMiddleware(logger zerolog.Logger) func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
fn := func(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
reqID := middleware.GetReqID(ctx)
logger := logger.With().Str("requestId", reqID).Logger()
ctx = logger.WithContext(ctx)
r = r.WithContext(ctx)
next.ServeHTTP(w, r)
}
return http.HandlerFunc(fn)
}
}
func NewRequestLoggerMiddleware() func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
fn := func(w http.ResponseWriter, r *http.Request) {
logger := *zerolog.Ctx(r.Context())
logger = logger.With().
Str("method", r.Method).
Bool("tls", r.TLS != nil).
Str("host", r.Host).
Str("url", r.RequestURI).
Str("protocol", r.Proto).
Str("from", r.RemoteAddr).
Str("origin", r.Header.Get("Origin")).
Logger()
logger.Info().Msg("New request")
start := time.Now()
ww := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
defer func() {
logger.Info().
Float64("durationSec", time.Since(start).Seconds()).
Int("status", ww.Status()).
Int("responseBytes", ww.BytesWritten()).
Msg("Request finished")
}()
next.ServeHTTP(ww, r)
}
return http.HandlerFunc(fn)
}
}
func NewRecoverMiddleware() func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
fn := func(w http.ResponseWriter, r *http.Request) {
defer func() {
if err := recover(); err != nil {
stack := util.GetStack(1)
zerolog.Ctx(r.Context()).Error().
Any("stack", stack).
Any("panicValue", err).
Msg("Panic recovered")
w.WriteHeader(http.StatusInternalServerError)
}
}()
next.ServeHTTP(w, r)
}
return http.HandlerFunc(fn)
}
}