From 536aaaab142ddde2fe3569a504434c3738835d51 Mon Sep 17 00:00:00 2001
From: Kayra <kayra.gemalmaz@canonical.com>
Date: Tue, 12 Nov 2024 15:51:56 +0100
Subject: [PATCH] feat: implement response controller for logging middleware
 (#101)

---
 internal/server/middleware.go          | 20 -----------------
 internal/server/response_controller.go | 30 ++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 20 deletions(-)
 create mode 100644 internal/server/response_controller.go

diff --git a/internal/server/middleware.go b/internal/server/middleware.go
index fb681e0..8fef959 100644
--- a/internal/server/middleware.go
+++ b/internal/server/middleware.go
@@ -28,26 +28,6 @@ type middlewareContext struct {
 	jwtSecret          []byte
 }
 
-// The statusRecorder struct wraps the http.ResponseWriter struct, and extracts the status
-// code of the response writer for the middleware to read
-type statusRecorder struct {
-	http.ResponseWriter
-	statusCode int
-}
-
-// newResponseWriter returns a new ResponseWriterCloner struct
-// it returns http.StatusOK by default because the http.ResponseWriter defaults to that header
-// if the WriteHeader() function is never called.
-func newResponseWriter(w http.ResponseWriter) *statusRecorder {
-	return &statusRecorder{w, http.StatusOK}
-}
-
-// WriteHeader overrides the ResponseWriter method to duplicate the status code into the wrapper struct
-func (rwc *statusRecorder) WriteHeader(code int) {
-	rwc.statusCode = code
-	rwc.ResponseWriter.WriteHeader(code)
-}
-
 // createMiddlewareStack chains the given middleware functions to wrap the api.
 // Each middleware functions calls next.ServeHTTP in order to resume the chain of execution.
 // The order the middleware functions are given to createMiddlewareStack matters.
diff --git a/internal/server/response_controller.go b/internal/server/response_controller.go
new file mode 100644
index 0000000..7b7abd8
--- /dev/null
+++ b/internal/server/response_controller.go
@@ -0,0 +1,30 @@
+package server
+
+import "net/http"
+
+type responseWriter struct {
+	http.ResponseWriter
+	statusCode int
+}
+
+func newResponseWriter(w http.ResponseWriter) *responseWriter {
+	return &responseWriter{
+		ResponseWriter: w,
+		statusCode:     0,
+	}
+}
+
+func (mw *responseWriter) WriteHeader(statusCode int) {
+	mw.ResponseWriter.WriteHeader(statusCode)
+	if mw.statusCode == 0 {
+		mw.statusCode = statusCode
+	}
+}
+
+func (mw *responseWriter) Write(b []byte) (int, error) {
+	return mw.ResponseWriter.Write(b)
+}
+
+func (mw *responseWriter) Unwrap() http.ResponseWriter {
+	return mw.ResponseWriter
+}