From 6742410025d5e99c60045bb314730799f0e1e5ce Mon Sep 17 00:00:00 2001
From: Matthis Holleville <matthish29@gmail.com>
Date: Sun, 16 Apr 2023 23:00:24 +0200
Subject: [PATCH] feat: init logging middleware on server mode

Signed-off-by: Matthis Holleville <matthish29@gmail.com>
---
 pkg/server/log.go    | 77 ++++++++++++++++++++++++++++++++++++++++++++
 pkg/server/server.go | 17 ++++++----
 2 files changed, 87 insertions(+), 7 deletions(-)
 create mode 100644 pkg/server/log.go

diff --git a/pkg/server/log.go b/pkg/server/log.go
new file mode 100644
index 0000000000..7115d964e1
--- /dev/null
+++ b/pkg/server/log.go
@@ -0,0 +1,77 @@
+package server
+
+import (
+	"bytes"
+	"net/http"
+	"time"
+
+	"go.uber.org/zap"
+)
+
+type loggingResponseWriter struct {
+	http.ResponseWriter
+	statusCode int
+	buf        *bytes.Buffer
+}
+
+func NewLoggingResponseWriter(w http.ResponseWriter) *loggingResponseWriter {
+	return &loggingResponseWriter{
+		w,
+		http.StatusOK,
+		&bytes.Buffer{},
+	}
+}
+
+func (lrw *loggingResponseWriter) WriteHeader(code int) {
+	lrw.statusCode = code
+	lrw.ResponseWriter.WriteHeader(code)
+}
+
+func (lrw *loggingResponseWriter) Write(b []byte) (int, error) {
+	return lrw.buf.Write(b)
+}
+
+func (lrw *loggingResponseWriter) Flush() {
+	if f, ok := lrw.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+	lrw.ResponseWriter.Write(lrw.buf.Bytes())
+}
+
+func logRequest(logger *zap.Logger, fields []zap.Field, statusCode int, message string) {
+	if statusCode >= 400 {
+		logger.Error(message, fields...)
+	} else {
+		logger.Info("request completed", fields...)
+	}
+}
+
+func loggingMiddleware(next http.Handler) http.Handler {
+	config := zap.NewProductionConfig()
+	config.DisableCaller = true
+	logger, err := config.Build()
+	if err != nil {
+		panic(err)
+	}
+	defer logger.Sync()
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		lrw := NewLoggingResponseWriter(w)
+		start := time.Now()
+		defer func() {
+			duration := time.Since(start).Milliseconds()
+			fields := []zap.Field{
+				zap.Int64("duration_ms", duration),
+				zap.String("method", r.Method),
+				zap.String("remote_addr", r.RemoteAddr),
+				zap.Int("status_code", lrw.statusCode),
+				zap.String("url", r.URL.Path),
+			}
+			logRequest(logger, fields, lrw.statusCode, lrw.buf.String())
+		}()
+
+		next.ServeHTTP(lrw, r)
+
+		lrw.Flush()
+	})
+}
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 9217aa947b..341a611159 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -51,30 +51,32 @@ func (s *Config) analyzeHandler(w http.ResponseWriter, r *http.Request) {
 	config, err := analysis.NewAnalysis(s.Backend, language, []string{}, namespace, nocache, explain)
 	if err != nil {
 		health.Failure++
-		fmt.Fprintf(w, err.Error())
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
 	}
 
 	err = config.RunAnalysis()
 	if err != nil {
 		color.Red("Error: %v", err)
 		health.Failure++
-		fmt.Fprintf(w, err.Error())
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
 	}
 
 	if explain {
 		err := config.GetAIResults(s.Output, anonymize)
 		if err != nil {
-			color.Red("Error: %v", err)
 			health.Failure++
-			fmt.Fprintf(w, err.Error())
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
 		}
 	}
 
 	out, err := config.PrintOutput(s.Output)
 	if err != nil {
-		color.Red("Error: %v", err)
 		health.Failure++
-		fmt.Fprintf(w, err.Error())
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
 	}
 
 	health.Success++
@@ -82,11 +84,12 @@ func (s *Config) analyzeHandler(w http.ResponseWriter, r *http.Request) {
 }
 
 func (s *Config) Serve() error {
+	handler := loggingMiddleware(http.DefaultServeMux)
 	http.Handle("/metrics", promhttp.Handler())
 	http.HandleFunc("/analyze", s.analyzeHandler)
 	http.HandleFunc("/healthz", s.healthzHandler)
 	color.Green("Starting server on port %s", s.Port)
-	err := http.ListenAndServe(":"+s.Port, nil)
+	err := http.ListenAndServe(":"+s.Port, handler)
 	if err != nil {
 		fmt.Printf("error starting server: %s\n", err)
 		return err