logging update with context

internal/middleware/context.go

@@ -0,0 +1,21 @@
+package middleware
+
+import (
+	"context"
+	"log/slog"
+	"net/http"
+
+	"github.com/google/uuid"
+)
+
+func ContextMiddleware(logger *slog.Logger) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			traceid := uuid.NewString()
+			ctx := context.WithValue(r.Context(), "trace_id", traceid)
+			newReq := r.WithContext(ctx)
+
+			next.ServeHTTP(w, newReq)
+		})
+	}
+}

internal/middleware/logging.go

@@ -1,23 +1,39 @@
 package middleware
 
 import (
-	"context"
 	"log/slog"
 	"net/http"
 )
 
+type LoggingResponseWriter struct {
+	http.ResponseWriter
+	statusCode int
+}
+
+func NewLoggingResponseWriter(w http.ResponseWriter) *LoggingResponseWriter {
+	return &LoggingResponseWriter{w, http.StatusOK}
+}
+
+func (lrw *LoggingResponseWriter) WriteHeader(code int) {
+	lrw.statusCode = code
+	lrw.ResponseWriter.WriteHeader(code)
+}
+
 func LoggingMiddleware(logger *slog.Logger) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			logger.LogAttrs(
-				context.Background(),
-				slog.LevelInfo,
-				"Incoming request",
+			logger.InfoContext(r.Context(), "Incoming request",
 				slog.String("method", r.Method),
 				slog.String("path", r.URL.String()),
 			)
 
-			next.ServeHTTP(w, r)
+			lrw := NewLoggingResponseWriter(w)
+			next.ServeHTTP(lrw, r)
+
+			logger.InfoContext(r.Context(), "Sent response",
+				slog.Int("code", lrw.statusCode),
+				slog.String("message", http.StatusText(lrw.statusCode)),
+			)
 		})
 	}
 }