From 543f524d286de0e1904c8a8956ca3aea6e59a254 Mon Sep 17 00:00:00 2001
From: matoval <mat24c@gmail.com>
Date: Wed, 18 Sep 2024 13:06:46 -0700
Subject: [PATCH] Add debug log for work unit payloads

---
 pkg/controlsvc/controlsvc.go | 32 ++++++++++++++++++++++++++++++--
 pkg/workceptor/command.go    | 28 +++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/pkg/controlsvc/controlsvc.go b/pkg/controlsvc/controlsvc.go
index f483cacde..a6647a5af 100644
--- a/pkg/controlsvc/controlsvc.go
+++ b/pkg/controlsvc/controlsvc.go
@@ -4,6 +4,7 @@
 package controlsvc
 
 import (
+	"bufio"
 	"context"
 	"crypto/tls"
 	"encoding/json"
@@ -122,8 +123,35 @@ func (s *SockControl) ReadFromConn(message string, out io.Writer, io Copier) err
 	if err := s.WriteMessage(message); err != nil {
 		return err
 	}
-	if _, err := io.Copy(out, s.conn); err != nil {
-		return err
+	isPayloadDebug := os.Getenv("RECEPTOR_PAYLOAD_DEBUG")
+	if isPayloadDebug != "" {
+		var data string
+		reader := bufio.NewReader(s.conn)
+
+		for {
+			var connectType string
+			if s.conn.LocalAddr().Network() == "unix" {
+				connectType = "unix socket"
+			} else {
+				connectType = "network connection"
+			}
+			response, err := reader.ReadString('\n')
+			if err != nil {
+				if err.Error() != "EOF" {
+					MainInstance.nc.GetLogger().Error("Error reading from %v: %v \n", connectType, err)
+				}
+				break
+			}
+			data += response
+			MainInstance.nc.GetLogger().Debug("Response from %v: %v", connectType, response)
+		}
+		if _, err := out.Write([]byte(data)); err != nil {
+			return err
+		}
+	} else {
+		if _, err := io.Copy(out, s.conn); err != nil {
+			return err
+		}
 	}
 
 	return nil
diff --git a/pkg/workceptor/command.go b/pkg/workceptor/command.go
index 2ff793ce5..be4e1b08b 100644
--- a/pkg/workceptor/command.go
+++ b/pkg/workceptor/command.go
@@ -4,9 +4,11 @@
 package workceptor
 
 import (
+	"bufio"
 	"context"
 	"flag"
 	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"os/signal"
@@ -112,7 +114,31 @@ func commandRunner(command string, params string, unitdir string) error {
 	if err != nil {
 		return err
 	}
-	cmd.Stdin = stdin
+	isPayloadDebug := os.Getenv("RECEPTOR_PAYLOAD_DEBUG")
+	if isPayloadDebug != "" {
+		var data string
+		splitUnitDir := strings.Split(unitdir, "/")
+		workUnitID := splitUnitDir[len(splitUnitDir) - 1]
+		reader := bufio.NewReader(stdin)
+		stdinStream, err := cmd.StdinPipe()
+		if err != nil {
+			return err
+		}
+		for {
+			response, err := reader.ReadString('\n')
+			if err != nil {
+				if err.Error() != "EOF" {
+					MainInstance.nc.GetLogger().Error("Error reading work unit %v stdin: %v\n", workUnitID, err)
+				}
+				break
+			}
+			data += response
+			MainInstance.nc.GetLogger().Debug("Work unit %v stdin: %v", workUnitID, response)
+		}
+		io.WriteString(stdinStream, data)
+	} else {
+		cmd.Stdin = stdin
+	}
 	stdout, err := os.OpenFile(path.Join(unitdir, "stdout"), os.O_CREATE+os.O_WRONLY+os.O_SYNC, 0o600)
 	if err != nil {
 		return err