From 5ce14480494c4a9c97ad0cd8ba287207d42cecbe Mon Sep 17 00:00:00 2001
From: Krishnan Parthasarathi <krisis@users.noreply.github.com>
Date: Thu, 1 Sep 2022 12:51:04 -0700
Subject: [PATCH] audit: account for response headers separately (#15610)

---
 internal/logger/audit.go               | 16 +++++++++++++---
 internal/logger/message/audit/entry.go |  1 +
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/internal/logger/audit.go b/internal/logger/audit.go
index 9d0c3adfe..f0a8b0c30 100644
--- a/internal/logger/audit.go
+++ b/internal/logger/audit.go
@@ -45,6 +45,8 @@ type ResponseWriter struct {
 	StartTime       time.Time
 	// number of bytes written
 	bytesWritten int
+	// number of bytes of response headers written
+	headerBytesWritten int
 	// Internal recording buffer
 	headers bytes.Buffer
 	body    bytes.Buffer
@@ -86,10 +88,10 @@ func (lrw *ResponseWriter) Write(p []byte) (int, error) {
 // Write the headers into the given buffer
 func (lrw *ResponseWriter) writeHeaders(w io.Writer, statusCode int, headers http.Header) {
 	n, _ := fmt.Fprintf(w, "%d %s\n", statusCode, http.StatusText(statusCode))
-	lrw.bytesWritten += n
+	lrw.headerBytesWritten += n
 	for k, v := range headers {
 		n, _ := fmt.Fprintf(w, "%s: %s\n", k, v[0])
-		lrw.bytesWritten += n
+		lrw.headerBytesWritten += n
 	}
 }
 
@@ -122,11 +124,16 @@ func (lrw *ResponseWriter) Flush() {
 	lrw.ResponseWriter.(http.Flusher).Flush()
 }
 
-// Size - reutrns the number of bytes written
+// Size - returns  the number of bytes written
 func (lrw *ResponseWriter) Size() int {
 	return lrw.bytesWritten
 }
 
+// HeaderSize - returns the number of bytes of response headers written
+func (lrw *ResponseWriter) HeaderSize() int {
+	return lrw.headerBytesWritten
+}
+
 const contextAuditKey = contextKeyType("audit-entry")
 
 // SetAuditEntry sets Audit info in the context.
@@ -187,6 +194,7 @@ func AuditLog(ctx context.Context, w http.ResponseWriter, r *http.Request, reqCl
 			timeToResponse  time.Duration
 			timeToFirstByte time.Duration
 			outputBytes     int64 = -1 // -1: unknown output bytes
+			headerBytes     int64
 		)
 
 		var st *ResponseWriter
@@ -204,6 +212,7 @@ func AuditLog(ctx context.Context, w http.ResponseWriter, r *http.Request, reqCl
 			timeToResponse = time.Now().UTC().Sub(st.StartTime)
 			timeToFirstByte = st.TimeToFirstByte
 			outputBytes = int64(st.Size())
+			headerBytes = int64(st.HeaderSize())
 		}
 
 		entry.API.Name = reqInfo.API
@@ -220,6 +229,7 @@ func AuditLog(ctx context.Context, w http.ResponseWriter, r *http.Request, reqCl
 		entry.API.StatusCode = statusCode
 		entry.API.InputBytes = r.ContentLength
 		entry.API.OutputBytes = outputBytes
+		entry.API.HeaderBytes = headerBytes
 		entry.API.TimeToResponse = strconv.FormatInt(timeToResponse.Nanoseconds(), 10) + "ns"
 		entry.Tags = reqInfo.GetTagsMap()
 		// ttfb will be recorded only for GET requests, Ignore such cases where ttfb will be empty.
diff --git a/internal/logger/message/audit/entry.go b/internal/logger/message/audit/entry.go
index ab76018b0..b88cb0003 100644
--- a/internal/logger/message/audit/entry.go
+++ b/internal/logger/message/audit/entry.go
@@ -53,6 +53,7 @@ type Entry struct {
 		StatusCode      int             `json:"statusCode,omitempty"`
 		InputBytes      int64           `json:"rx"`
 		OutputBytes     int64           `json:"tx"`
+		HeaderBytes     int64           `json:"txHeaders,omitempty"`
 		TimeToFirstByte string          `json:"timeToFirstByte,omitempty"`
 		TimeToResponse  string          `json:"timeToResponse,omitempty"`
 	} `json:"api"`