package analytics import ( "context" "net/http" ) type contextKey struct{} var countKeyContextKey contextKey func WithCountKey(r *http.Request, key string) *http.Request { return r.WithContext(context.WithValue(r.Context(), countKeyContextKey, key)) } func GetCountKey(r *http.Request) (string, bool) { key, ok := r.Context().Value(countKeyContextKey).(string) return key, ok } func CounterMiddleware(counter Counter) func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { rw := newStatusCapturingResponseWriter(w) next.ServeHTTP(rw, r) countKey, ok := GetCountKey(r) if !ok { if rw.status >= 201 { countKey = http.StatusText(rw.status) } else { countKey = r.URL.Path } } counter.Count(r, countKey) }) } } type statusCapturingResponseWriter struct { http.ResponseWriter status int } func newStatusCapturingResponseWriter(w http.ResponseWriter) *statusCapturingResponseWriter { return &statusCapturingResponseWriter{w, http.StatusOK} } func (w *statusCapturingResponseWriter) WriteHeader(code int) { w.status = code w.ResponseWriter.WriteHeader(code) }