diff --git a/cmd/serve.go b/cmd/serve.go index 77882d3..294e8a8 100644 --- a/cmd/serve.go +++ b/cmd/serve.go @@ -157,7 +157,6 @@ You can also enable the caching functionality to speed things up.`, r.Use(func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - shouldCache := false for _, prefix := range cachePrefixes { if strings.HasPrefix(r.URL.Path, strings.TrimSpace(prefix)) { @@ -168,13 +167,16 @@ You can also enable the caching functionality to speed things up.`, if shouldCache { log.Log.Debugf("Caching request for path: %s", r.URL.Path) - cachedMiddleware(next).ServeHTTP(w, r) + wrapper := customMiddleware.NewResponseWrapper(w) + cachedMiddleware(next).ServeHTTP(wrapper, r) + if wrapper.WasWritten() { + w.Header().Set("X-Cache", "HIT") + } } else { next.ServeHTTP(w, r) } }) }) - } if config.UI { diff --git a/internal/middleware/response_wrapper.go b/internal/middleware/response_wrapper.go new file mode 100644 index 0000000..a80e473 --- /dev/null +++ b/internal/middleware/response_wrapper.go @@ -0,0 +1,26 @@ +package middleware + +import "net/http" + +type ResponseWrapper struct { + http.ResponseWriter + written bool +} + +func NewResponseWrapper(w http.ResponseWriter) *ResponseWrapper { + return &ResponseWrapper{ResponseWriter: w} +} + +func (w *ResponseWrapper) Write(b []byte) (int, error) { + w.written = true + return w.ResponseWriter.Write(b) +} + +func (w *ResponseWrapper) WriteHeader(statusCode int) { + w.written = true + w.ResponseWriter.WriteHeader(statusCode) +} + +func (w *ResponseWrapper) WasWritten() bool { + return w.written +}