From 71702002cfdad4cf4a55a92c03ac660903a65744 Mon Sep 17 00:00:00 2001
From: Julien Dessaux
Date: Sun, 13 Oct 2024 16:10:22 +0200
Subject: feat(logger): implement optional body logging

---
 pkg/logger/body_writer.go | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

(limited to 'pkg/logger/body_writer.go')

diff --git a/pkg/logger/body_writer.go b/pkg/logger/body_writer.go
index 60da151..f57a03e 100644
--- a/pkg/logger/body_writer.go
+++ b/pkg/logger/body_writer.go
@@ -2,6 +2,7 @@ package logger
 
 import (
 	"bufio"
+	"bytes"
 	"errors"
 	"net"
 	"net/http"
@@ -9,11 +10,20 @@ import (
 
 type bodyWriter struct {
 	http.ResponseWriter
-	status int
+	body    *bytes.Buffer
+	maxSize int
+	bytes   int
+	status  int
 }
 
-func newBodyWriter(writer http.ResponseWriter) *bodyWriter {
+func newBodyWriter(writer http.ResponseWriter, recordBody bool) *bodyWriter {
+	var body *bytes.Buffer
+	if recordBody {
+		body = bytes.NewBufferString("")
+	}
 	return &bodyWriter{
+		body:           body,
+		maxSize:        64000,
 		ResponseWriter: writer,
 		status:         http.StatusNotImplemented,
 	}
@@ -21,6 +31,15 @@ func newBodyWriter(writer http.ResponseWriter) *bodyWriter {
 
 // implements http.ResponseWriter
 func (w *bodyWriter) Write(b []byte) (int, error) {
+	if w.body != nil {
+		if w.body.Len()+len(b) > w.maxSize {
+			w.body.Write(b[:w.maxSize-w.body.Len()])
+		} else {
+			w.body.Write(b)
+		}
+	}
+
+	w.bytes += len(b)
 	return w.ResponseWriter.Write(b)
 }
 
-- 
cgit v1.2.3