From 55e3a68db671c090b63004a740cd3bf8b0072fbc Mon Sep 17 00:00:00 2001
From: Ryan D McGuire <ryand_mcguire@sweetwater.com>
Date: Sun, 5 Jan 2025 16:22:15 -0500
Subject: [PATCH] Add logging middleware

---
 pkg/config/types.go |  2 ++
 pkg/srv/http.go     |  5 +++
 pkg/srv/http_log.go | 87 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 94 insertions(+)
 create mode 100644 pkg/srv/http_log.go

diff --git a/pkg/config/types.go b/pkg/config/types.go
index d44b334..907e9c1 100644
--- a/pkg/config/types.go
+++ b/pkg/config/types.go
@@ -15,6 +15,7 @@ var DefaultConfig = &AppConfig{
 	},
 	HTTP: &HTTPConfig{
 		Listen:       "127.0.0.1:8080",
+		LogRequests:  false,
 		ReadTimeout:  "10s",
 		WriteTimeout: "10s",
 		IdleTimeout:  "1m",
@@ -76,6 +77,7 @@ const (
 // HTTP Configuration
 type HTTPConfig struct {
 	Listen       string `yaml:"listen,omitempty" env:"APP_HTTP_LISTEN"`
+	LogRequests  bool   `yaml:"logRequests" env:"APP_HTTP_LOG_REQUESTS"`
 	ReadTimeout  string `yaml:"readTimeout" env:"APP_HTTP_READ_TIMEOUT"`   // Go duration (e.g. 10s)
 	WriteTimeout string `yaml:"writeTimeout" env:"APP_HTTP_WRITE_TIMEOUT"` // Go duration (e.g. 10s)
 	IdleTimeout  string `yaml:"idleTimeout" env:"APP_HTTP_IDLE_TIMEOUT"`   // Go duration (e.g. 10s)
diff --git a/pkg/srv/http.go b/pkg/srv/http.go
index 07e48c8..0b43962 100644
--- a/pkg/srv/http.go
+++ b/pkg/srv/http.go
@@ -89,6 +89,11 @@ func prepHTTPServer(ctx context.Context, handleFuncs []HTTPFunc, hcFuncs ...Heal
 		idleTimeout = *iT
 	}
 
+	// Inject logging middleware
+	if cfg.HTTP.LogRequests {
+		handler = loggingMiddleware(ctx, handler)
+	}
+
 	return &http.Server{
 		Addr:         cfg.HTTP.Listen,
 		ReadTimeout:  readTimeout,
diff --git a/pkg/srv/http_log.go b/pkg/srv/http_log.go
new file mode 100644
index 0000000..b59c7e0
--- /dev/null
+++ b/pkg/srv/http_log.go
@@ -0,0 +1,87 @@
+package srv
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"net/http"
+	"regexp"
+	"time"
+
+	"github.com/rs/zerolog"
+)
+
+var ExcludeFromLogging = regexp.MustCompile(`\/(ready|live|metrics)$`)
+
+type LoggingResponseWriter struct {
+	http.ResponseWriter
+	statusCode int
+	body       *bytes.Buffer
+}
+
+func loggingMiddleware(appCtx context.Context, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if ExcludeFromLogging.Match([]byte(r.URL.Path)) {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		log := zerolog.Ctx(appCtx)
+
+		start := time.Now()
+		lrr := newLoggingResponseWriter(w)
+		next.ServeHTTP(lrr, r)
+
+		log.Debug().
+			Str("path", r.URL.Path).
+			Any("query", r.URL.Query()).
+			Int("statusCode", lrr.statusCode).
+			Str("protocol", r.Proto).
+			Str("remote", r.RemoteAddr).
+			Dur("duration", time.Since(start)).
+			Msg("http request served")
+
+		// Log response body
+		trcLog := log.Trace().
+			Str("path", r.URL.Path).
+			Int("statusCode", lrr.statusCode)
+
+		// Check if it's JSON
+		firstByte, err := lrr.body.ReadByte()
+		if err != nil {
+			trcLog.Err(errors.New("invalid response body")).Send()
+			return
+		}
+		lrr.body.UnreadByte()
+
+		if firstByte == '{' {
+			trcLog = trcLog.RawJSON("response", lrr.body.Bytes())
+		} else {
+			trcLog = trcLog.Bytes("response", lrr.body.Bytes())
+		}
+		trcLog.Msg("response payload")
+	})
+}
+
+// Implement Flush to support the http.Flusher interface
+func (w *LoggingResponseWriter) Flush() {
+	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+func (w *LoggingResponseWriter) WriteHeader(code int) {
+	w.statusCode = code
+	w.ResponseWriter.WriteHeader(code)
+}
+
+func (w *LoggingResponseWriter) Write(b []byte) (int, error) {
+	w.body.Write(b)
+	return w.ResponseWriter.Write(b)
+}
+
+func newLoggingResponseWriter(w http.ResponseWriter) *LoggingResponseWriter {
+	return &LoggingResponseWriter{
+		ResponseWriter: w, statusCode: http.StatusOK, body: bytes.NewBuffer(nil),
+	}
+}