From a568f341ddb8384f844b15f4e8a671b9c3d5267f Mon Sep 17 00:00:00 2001
From: dadav <33197631+dadav@users.noreply.github.com>
Date: Wed, 25 Dec 2024 11:16:38 +0100
Subject: [PATCH] fix: Improve

---
 cmd/serve.go                            |  8 +++++---
 internal/middleware/response_wrapper.go | 26 +++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)
 create mode 100644 internal/middleware/response_wrapper.go

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
+}