From f20a9dd46b686e6d8360655ce58e0f31302ba3dd Mon Sep 17 00:00:00 2001
From: Seokho Son <shsongist@gmail.com>
Date: Tue, 11 Jun 2024 20:32:22 +0900
Subject: [PATCH] Display ssh std to streaming pipe

---
 .../server/middlewares/custom-middleware.go   |  6 +--
 src/core/common/client.go                     |  4 +-
 src/core/mcis/remoteCommand.go                | 53 ++++++++++++++-----
 3 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/src/api/rest/server/middlewares/custom-middleware.go b/src/api/rest/server/middlewares/custom-middleware.go
index 047198cdb..d0a5e50f7 100644
--- a/src/api/rest/server/middlewares/custom-middleware.go
+++ b/src/api/rest/server/middlewares/custom-middleware.go
@@ -130,11 +130,11 @@ func ResponseBodyDump() echo.MiddlewareFunc {
 
 			// Get the request ID
 			reqID := c.Get("RequestID").(string)
-			log.Trace().Msgf("(BodyDump middleware) Request ID: %s", reqID)
+			//log.Trace().Msgf("(BodyDump middleware) Request ID: %s", reqID)
 
 			// Get the content type
 			contentType := c.Response().Header().Get(echo.HeaderContentType)
-			log.Trace().Msgf("contentType: %s", contentType)
+			//log.Trace().Msgf("contentType: %s", contentType)
 
 			// log.Debug().Msgf("Request body: %s", string(reqBody))
 			// log.Debug().Msgf("Response body: %s", string(resBody))
@@ -143,7 +143,7 @@ func ResponseBodyDump() echo.MiddlewareFunc {
 			if contentType == echo.MIMEApplicationJSONCharsetUTF8 || contentType == echo.MIMEApplicationJSON {
 				// Load or check the request by ID
 				if v, ok := common.RequestMap.Load(reqID); ok {
-					log.Trace().Msg("OK, common.RequestMap.Load(reqID)")
+					//log.Trace().Msg("OK, common.RequestMap.Load(reqID)")
 					details := v.(common.RequestDetails)
 					details.EndTime = time.Now()
 
diff --git a/src/core/common/client.go b/src/core/common/client.go
index 8161a6b81..57969b2ec 100644
--- a/src/core/common/client.go
+++ b/src/core/common/client.go
@@ -129,7 +129,7 @@ func ExecuteHttpRequest[B any, T any](
 
 				return nil
 			} else {
-				log.Trace().Msg("Cache item expired!")
+				//log.Trace().Msg("Cache item expired!")
 				clientCache.Delete(requestKey)
 			}
 		}
@@ -223,7 +223,7 @@ func ExecuteHttpRequest[B any, T any](
 			log.Trace().Msg("Fesult is nil, not caching")
 		} else {
 			clientCache.Store(requestKey, CacheItem[T]{Response: *result, ExpiresAt: time.Now().Add(cacheDuration)})
-			log.Trace().Msg("Cached successfully!")
+			//log.Trace().Msg("Cached successfully!")
 		}
 	}
 
diff --git a/src/core/mcis/remoteCommand.go b/src/core/mcis/remoteCommand.go
index 77e3b3beb..9dabd365f 100644
--- a/src/core/mcis/remoteCommand.go
+++ b/src/core/mcis/remoteCommand.go
@@ -18,7 +18,9 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net"
+	"os"
 	"sync"
 	"time"
 
@@ -510,13 +512,6 @@ func runSSH(bastionInfo sshInfo, targetInfo sshInfo, cmds []string) (map[int]str
 	client := ssh.NewClient(ncc, chans, reqs)
 	defer client.Close()
 
-	// Create a new SSH session
-	session, err := client.NewSession()
-	if err != nil {
-		return stdoutMap, stderrMap, err
-	}
-	defer session.Close()
-
 	// Run the commands
 	for i, cmd := range cmds {
 		// Create a new SSH session for each command
@@ -524,18 +519,48 @@ func runSSH(bastionInfo sshInfo, targetInfo sshInfo, cmds []string) (map[int]str
 		if err != nil {
 			return stdoutMap, stderrMap, err
 		}
-		defer session.Close()
+		defer session.Close() // Ensure session is closed
+
+		// Get pipes for stdout and stderr
+		stdoutPipe, err := session.StdoutPipe()
+		if err != nil {
+			return stdoutMap, stderrMap, err
+		}
+
+		stderrPipe, err := session.StderrPipe()
+		if err != nil {
+			return stdoutMap, stderrMap, err
+		}
 
-		// Capture the output
+		// Start the command
+		if err := session.Start(cmd); err != nil {
+			return stdoutMap, stderrMap, err
+		}
+
+		// Read stdout and stderr
 		var stdoutBuf, stderrBuf bytes.Buffer
-		session.Stdout = &stdoutBuf
-		session.Stderr = &stderrBuf
+		stdoutDone := make(chan struct{})
+		stderrDone := make(chan struct{})
+
+		go func() {
+			io.Copy(io.MultiWriter(os.Stdout, &stdoutBuf), stdoutPipe)
+			close(stdoutDone)
+		}()
+
+		go func() {
+			io.Copy(io.MultiWriter(os.Stderr, &stderrBuf), stderrPipe)
+			close(stderrDone)
+		}()
+
+		// Wait for the command to finish
+		err = session.Wait()
+		<-stdoutDone
+		<-stderrDone
 
-		// Run the command
-		err = session.Run(cmd)
 		if err != nil {
 			stderrMap[i] = fmt.Sprintf("(%s)\nStderr: %s", err, stderrBuf.String())
-			break // Stop if the command fails
+			stdoutMap[i] = stdoutBuf.String()
+			break
 		}
 
 		stdoutMap[i] = stdoutBuf.String()