From 2cdbc4e782c7ad7f42b1ddd3b7be6a24dea580e7 Mon Sep 17 00:00:00 2001
From: Julien Dessaux
Date: Tue, 1 Oct 2024 08:47:32 +0200
Subject: feat(logger): implement a logger middleware

---
 cmd/tfstated/main.go      |  3 ++-
 cmd/tfstated/routes.go    |  1 +
 pkg/logger/body_writer.go | 46 ++++++++++++++++++++++++++++++++
 pkg/logger/middleware.go  | 68 +++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 117 insertions(+), 1 deletion(-)
 create mode 100644 pkg/logger/body_writer.go
 create mode 100644 pkg/logger/middleware.go

diff --git a/cmd/tfstated/main.go b/cmd/tfstated/main.go
index 61248bc..cf93976 100644
--- a/cmd/tfstated/main.go
+++ b/cmd/tfstated/main.go
@@ -14,6 +14,7 @@ import (
 	"time"
 
 	"git.adyxax.org/adyxax/tfstated/pkg/database"
+	"git.adyxax.org/adyxax/tfstated/pkg/logger"
 )
 
 type Config struct {
@@ -49,7 +50,7 @@ func run(
 
 	httpServer := &http.Server{
 		Addr:    net.JoinHostPort(config.Host, config.Port),
-		Handler: mux,
+		Handler: logger.Middleware(mux),
 	}
 	go func() {
 		log.Printf("listening on %s\n", httpServer.Addr)
diff --git a/cmd/tfstated/routes.go b/cmd/tfstated/routes.go
index 32ab111..853970a 100644
--- a/cmd/tfstated/routes.go
+++ b/cmd/tfstated/routes.go
@@ -11,6 +11,7 @@ func addRoutes(
 	db *database.DB,
 ) {
 	mux.Handle("GET /healthz", handleHealthz())
+
 	mux.Handle("GET /", handleGet(db))
 	mux.Handle("POST /", handlePost(db))
 }
diff --git a/pkg/logger/body_writer.go b/pkg/logger/body_writer.go
new file mode 100644
index 0000000..60da151
--- /dev/null
+++ b/pkg/logger/body_writer.go
@@ -0,0 +1,46 @@
+package logger
+
+import (
+	"bufio"
+	"errors"
+	"net"
+	"net/http"
+)
+
+type bodyWriter struct {
+	http.ResponseWriter
+	status int
+}
+
+func newBodyWriter(writer http.ResponseWriter) *bodyWriter {
+	return &bodyWriter{
+		ResponseWriter: writer,
+		status:         http.StatusNotImplemented,
+	}
+}
+
+// implements http.ResponseWriter
+func (w *bodyWriter) Write(b []byte) (int, error) {
+	return w.ResponseWriter.Write(b)
+}
+
+// implements http.ResponseWriter
+func (w *bodyWriter) WriteHeader(code int) {
+	w.status = code
+	w.ResponseWriter.WriteHeader(code)
+}
+
+// implements http.Flusher
+func (w *bodyWriter) Flush() {
+	if f, ok := w.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+// implements http.Hijacker
+func (w *bodyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if hi, ok := w.ResponseWriter.(http.Hijacker); ok {
+		return hi.Hijack()
+	}
+	return nil, nil, errors.New("Hijack not supported")
+}
diff --git a/pkg/logger/middleware.go b/pkg/logger/middleware.go
new file mode 100644
index 0000000..d0a6a65
--- /dev/null
+++ b/pkg/logger/middleware.go
@@ -0,0 +1,68 @@
+package logger
+
+import (
+	"log/slog"
+	"net/http"
+	"runtime/debug"
+	"strconv"
+	"time"
+)
+
+func Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == "GET" && r.URL.Path == "/healthz" {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		defer func() {
+			if err := recover(); err != nil {
+				w.WriteHeader(http.StatusInternalServerError)
+				slog.Error(
+					"panic",
+					"err", err,
+					"trace", string(debug.Stack()),
+				)
+			}
+		}()
+		start := time.Now()
+		path := r.URL.Path
+		query := r.URL.RawQuery
+
+		bw := newBodyWriter(w)
+
+		next.ServeHTTP(bw, r)
+
+		end := time.Now()
+		requestAttributes := []slog.Attr{
+			slog.Time("time", start.UTC()),
+			slog.String("method", r.Method),
+			slog.String("host", r.Host),
+			slog.String("path", path),
+			slog.String("query", query),
+			slog.String("ip", r.RemoteAddr),
+		}
+		responseAttributes := []slog.Attr{
+			slog.Time("time", end.UTC()),
+			slog.Duration("latency", end.Sub(start)),
+			slog.Int("status", bw.status),
+		}
+		attributes := []slog.Attr{
+			{
+				Key:   "request",
+				Value: slog.GroupValue(requestAttributes...),
+			},
+			{
+				Key:   "response",
+				Value: slog.GroupValue(responseAttributes...),
+			},
+		}
+		level := slog.LevelInfo
+		if bw.status >= http.StatusInternalServerError {
+			level = slog.LevelError
+		} else if bw.status >= http.StatusBadRequest && bw.status < http.StatusInternalServerError {
+			level = slog.LevelWarn
+		}
+		slog.LogAttrs(r.Context(), level, strconv.Itoa(bw.status)+": "+http.StatusText(bw.status), attributes...)
+	})
+}
-- 
cgit v1.2.3