From ebfcbf095a76d208aa58ba966edba2a92a3f5137 Mon Sep 17 00:00:00 2001
From: Abhishek Shroff <shroff@kudosoft.net>
Date: Tue, 5 Mar 2024 12:40:37 +0530
Subject: [PATCH] Allow logging response body

---
 internal/command/serve.go | 41 ++++++++++++++++++++++++++++++++-------
 1 file changed, 34 insertions(+), 7 deletions(-)

diff --git a/internal/command/serve.go b/internal/command/serve.go
index 7253c663..d0f97884 100644
--- a/internal/command/serve.go
+++ b/internal/command/serve.go
@@ -1,14 +1,14 @@
 package command
 
 import (
-	"fmt"
+	"bytes"
 	"time"
 
 	"github.com/fvbock/endless"
 	"github.com/gin-contrib/cors"
 	"github.com/gin-gonic/gin"
 	"github.com/shroff/phylum/server/internal/webdav"
-	log "github.com/sirupsen/logrus"
+	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
@@ -19,17 +19,17 @@ func setupServeCommand() *cobra.Command {
 		Short: "Run the server",
 		Run: func(cmd *cobra.Command, args []string) {
 			config := viper.GetViper()
-			engine := createEngine(config.GetBool("cors_enabled"), config.GetStringSlice("cors_origins"))
+			engine := createEngine(config.GetBool("log_body"), config.GetBool("cors_enabled"), config.GetStringSlice("cors_origins"))
 
 			setupWebdav(engine.Group(config.GetString("webdav_prefix")))
 
 			server := endless.NewServer(config.GetString("listen"), engine)
 			server.BeforeBegin = func(addr string) {
-				log.Info(fmt.Sprintf("Listening on %s\n", addr))
+				logrus.Info("Listening on " + addr)
 			}
 
 			if err := server.ListenAndServe(); err != nil {
-				log.Fatal(err.Error())
+				logrus.Fatal(err.Error())
 			}
 		},
 	}
@@ -38,7 +38,10 @@ func setupServeCommand() *cobra.Command {
 	flags.StringP("listen", "l", ":1234", "Listen Addres")
 	viper.BindPFlag("listen", flags.Lookup("listen"))
 
-	flags.Bool("cors-enabled", false, "CORS enabled")
+	flags.Bool("log-body", false, "Log Response Body (Must be used with --debug)")
+	viper.BindPFlag("log_body", flags.Lookup("log-body"))
+
+	flags.Bool("cors-enabled", false, "Enabled CORS")
 	viper.BindPFlag("cors_enabled", flags.Lookup("cors-enabled"))
 
 	flags.StringSlice("cors-origins", []string{"*"}, "CORS origins")
@@ -65,11 +68,14 @@ func setupWebdav(r *gin.RouterGroup) {
 	r.Handle("PROPPATCH", "/*path", handler.HandleRequest)
 }
 
-func createEngine(corsEnabled bool, corsOrigins []string) *gin.Engine {
+func createEngine(logBody bool, corsEnabled bool, corsOrigins []string) *gin.Engine {
 	if !debug {
 		gin.SetMode(gin.ReleaseMode)
 	}
 	engine := gin.Default()
+	if debug && logBody {
+		engine.Use(logBodyMiddleware)
+	}
 
 	if corsEnabled {
 		engine.Use(cors.New(cors.Config{
@@ -85,3 +91,24 @@ func createEngine(corsEnabled bool, corsOrigins []string) *gin.Engine {
 
 	return engine
 }
+
+type logBodyWriter struct {
+	gin.ResponseWriter
+	body *bytes.Buffer
+}
+
+func (w logBodyWriter) Write(b []byte) (int, error) {
+	w.body.Write(b)
+	return w.ResponseWriter.Write(b)
+}
+
+func logBodyMiddleware(c *gin.Context) {
+	logBodyWriter := &logBodyWriter{
+		ResponseWriter: c.Writer,
+		body:           bytes.NewBuffer(make([]byte, 0, 1024)),
+	}
+	c.Writer = logBodyWriter
+	c.Next()
+	logrus.Trace(c.Request.URL)
+	logrus.Trace(logBodyWriter.body.String())
+}