From 7bde4c572b7fc2b5827ec041116bb6336072591f Mon Sep 17 00:00:00 2001
From: Pavel Okhlopkov <36456348+ldmonster@users.noreply.github.com>
Date: Wed, 30 Oct 2024 18:34:17 +0300
Subject: [PATCH] Fix/double logging std (#677)

Signed-off-by: Pavel Okhlopkov <pavel.okhlopkov@flant.com>
Co-authored-by: Pavel Okhlopkov <pavel.okhlopkov@flant.com>
---
 pkg/executor/executor.go      | 66 +++++++++++++++++++++++++++++++----
 pkg/executor/executor_test.go | 16 ++++++++-
 2 files changed, 75 insertions(+), 7 deletions(-)

diff --git a/pkg/executor/executor.go b/pkg/executor/executor.go
index 95d5ddb3..e62b7b81 100644
--- a/pkg/executor/executor.go
+++ b/pkg/executor/executor.go
@@ -1,6 +1,7 @@
 package executor
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/json"
 	"fmt"
@@ -77,17 +78,18 @@ type proxyJSONLogger struct {
 	logProxyHookJSON bool
 }
 
-func (pj *proxyJSONLogger) Write(p []byte) (n int, err error) {
-	pj.buf = append(pj.buf, p...)
-
+func (pj *proxyJSONLogger) Write(p []byte) (int, error) {
 	if !pj.logProxyHookJSON {
-		pj.Entry.Log(log.InfoLevel, strings.TrimSpace(string(pj.buf)))
+		pj.writerScanner(p)
 
 		return len(p), nil
 	}
 
+	// join all parts of json
+	pj.buf = append(pj.buf, p...)
+
 	var line interface{}
-	err = json.Unmarshal(pj.buf, &line)
+	err := json.Unmarshal(pj.buf, &line)
 	if err != nil {
 		if err.Error() == "unexpected end of JSON input" {
 			return len(p), nil
@@ -107,15 +109,67 @@ func (pj *proxyJSONLogger) Write(p []byte) (n int, err error) {
 		logMap[k] = v
 	}
 
-	logLine, _ := json.Marshal(logMap)
+	logLineRaw, _ := json.Marshal(logMap)
+
+	logLine := string(logLineRaw)
 
 	logEntry := pj.WithField(app.ProxyJsonLogKey, true)
 
+	if len(logLine) > 10000 {
+		logLine = fmt.Sprintf("%s:truncated", string(logLine[:10000]))
+
+		truncatedLog, _ := json.Marshal(map[string]string{
+			"truncated": logLine,
+		})
+
+		logEntry.Log(log.FatalLevel, string(truncatedLog))
+	}
+
 	logEntry.Log(log.FatalLevel, string(logLine))
 
 	return len(p), nil
 }
 
+func (pj *proxyJSONLogger) writerScanner(p []byte) {
+	scanner := bufio.NewScanner(bytes.NewReader(p))
+
+	// Set the buffer size to the maximum token size to avoid buffer overflows
+	scanner.Buffer(make([]byte, bufio.MaxScanTokenSize), bufio.MaxScanTokenSize)
+
+	// Define a split function to split the input into chunks of up to 64KB
+	chunkSize := bufio.MaxScanTokenSize // 64KB
+	splitFunc := func(data []byte, atEOF bool) (int, []byte, error) {
+		if len(data) >= chunkSize {
+			return chunkSize, data[:chunkSize], nil
+		}
+
+		return bufio.ScanLines(data, atEOF)
+	}
+
+	// Use the custom split function to split the input
+	scanner.Split(splitFunc)
+
+	// Scan the input and write it to the logger using the specified print function
+	for scanner.Scan() {
+		// prevent empty logging
+		str := strings.TrimSpace(scanner.Text())
+		if str == "" {
+			continue
+		}
+
+		if len(str) > 10000 {
+			str = fmt.Sprintf("%s:truncated", str[:10000])
+		}
+
+		pj.Entry.Info(str)
+	}
+
+	// If there was an error while scanning the input, log an error
+	if err := scanner.Err(); err != nil {
+		pj.Entry.Errorf("Error while reading from Writer: %s", err)
+	}
+}
+
 func Output(cmd *exec.Cmd) (output []byte, err error) {
 	// TODO context: hook name, hook phase, hook binding
 	// TODO observability
diff --git a/pkg/executor/executor_test.go b/pkg/executor/executor_test.go
index d3815555..a3cae4d3 100644
--- a/pkg/executor/executor_test.go
+++ b/pkg/executor/executor_test.go
@@ -55,7 +55,7 @@ func TestRunAndLogLines(t *testing.T) {
 		cmd := exec.Command("cat", f.Name())
 		_, err = RunAndLogLines(cmd, map[string]string{"a": "b"})
 		assert.NoError(t, err)
-		assert.Contains(t, buf.String(), `\",\"output\":\"stdout\"}" a=b output=stdout proxyJsonLog=true`)
+		assert.Contains(t, buf.String(), `:truncated\"}" a=b output=stdout proxyJsonLog=true`)
 
 		buf.Reset()
 	})
@@ -82,6 +82,20 @@ func TestRunAndLogLines(t *testing.T) {
 
 		buf.Reset()
 	})
+
+	t.Run("multiline non json", func(t *testing.T) {
+		app.LogProxyHookJSON = false
+		cmd := exec.Command("echo", `
+a b
+c d
+`)
+		_, err := RunAndLogLines(cmd, map[string]string{"foor": "baar"})
+		assert.NoError(t, err)
+		assert.Contains(t, buf.String(), `level=info msg="a b" foor=baar output=stdout`)
+		assert.Contains(t, buf.String(), `level=info msg="c d" foor=baar output=stdout`)
+
+		buf.Reset()
+	})
 }
 
 var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")