From b6641b277651be701f572fc9a1041a075b83a9cf Mon Sep 17 00:00:00 2001
From: ssongliu <73214554+ssongliu@users.noreply.github.com>
Date: Tue, 19 Nov 2024 18:24:05 +0800
Subject: [PATCH] feat: optimize the SSH connection method for the host (#7126)

---
 core/app/api/v2/host.go  |  4 ++--
 core/app/service/host.go |  4 ++--
 core/constant/status.go  |  4 ++--
 core/utils/ssh/ssh.go    | 27 ++++++++++-----------------
 4 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/core/app/api/v2/host.go b/core/app/api/v2/host.go
index e62599811..b4d1874cf 100644
--- a/core/app/api/v2/host.go
+++ b/core/app/api/v2/host.go
@@ -277,12 +277,12 @@ func (b *BaseApi) WsSsh(c *gin.Context) {
 		connInfo.PassPhrase = []byte(host.PassPhrase)
 	}
 
-	client, err := connInfo.NewClient()
+	client, err := ssh.NewClient(connInfo)
 	if wshandleError(wsConn, errors.WithMessage(err, "failed to set up the connection. Please check the host information")) {
 		return
 	}
 	defer client.Close()
-	sws, err := terminal.NewLogicSshWsSession(cols, rows, true, connInfo.Client, wsConn)
+	sws, err := terminal.NewLogicSshWsSession(cols, rows, true, client.Client, wsConn)
 	if wshandleError(wsConn, err) {
 		return
 	}
diff --git a/core/app/service/host.go b/core/app/service/host.go
index a0bf50cb8..47acbc5c9 100644
--- a/core/app/service/host.go
+++ b/core/app/service/host.go
@@ -65,7 +65,7 @@ func (u *HostService) TestByInfo(req dto.HostConnTest) bool {
 	if len(req.PassPhrase) != 0 {
 		connInfo.PassPhrase = []byte(req.PassPhrase)
 	}
-	client, err := connInfo.NewClient()
+	client, err := ssh.NewClient(connInfo)
 	if err != nil {
 		return false
 	}
@@ -114,7 +114,7 @@ func (u *HostService) TestLocalConn(id uint) bool {
 		}
 		connInfo.PassPhrase = []byte(host.PassPhrase)
 	}
-	client, err := connInfo.NewClient()
+	client, err := ssh.NewClient(connInfo)
 	if err != nil {
 		return false
 	}
diff --git a/core/constant/status.go b/core/constant/status.go
index f01b0eaca..166c0483c 100644
--- a/core/constant/status.go
+++ b/core/constant/status.go
@@ -1,8 +1,8 @@
 package constant
 
 const (
-	StatusSuccess = "Success"
-	StatusFailed  = "Failed"
+	StatusSuccess = "success"
+	StatusFailed  = "failed"
 
 	// node
 	StatusWaiting     = "waiting"
diff --git a/core/utils/ssh/ssh.go b/core/utils/ssh/ssh.go
index f487182d4..6f6f12c27 100644
--- a/core/utils/ssh/ssh.go
+++ b/core/utils/ssh/ssh.go
@@ -20,13 +20,13 @@ type ConnInfo struct {
 	PrivateKey  []byte        `json:"privateKey"`
 	PassPhrase  []byte        `json:"passPhrase"`
 	DialTimeOut time.Duration `json:"dialTimeOut"`
+}
 
-	Client     *gossh.Client  `json:"client"`
-	Session    *gossh.Session `json:"session"`
-	LastResult string         `json:"lastResult"`
+type SSHClient struct {
+	Client *gossh.Client `json:"client"`
 }
 
-func (c *ConnInfo) NewClient() (*ConnInfo, error) {
+func NewClient(c ConnInfo) (*SSHClient, error) {
 	if strings.Contains(c.Addr, ":") {
 		c.Addr = fmt.Sprintf("[%s]", c.Addr)
 	}
@@ -55,18 +55,12 @@ func (c *ConnInfo) NewClient() (*ConnInfo, error) {
 	}
 	client, err := gossh.Dial(proto, addr, config)
 	if nil != err {
-		return c, err
+		return nil, err
 	}
-	c.Client = client
-	return c, nil
+	return &SSHClient{Client: client}, nil
 }
 
-func (c *ConnInfo) Run(shell string) (string, error) {
-	if c.Client == nil {
-		if _, err := c.NewClient(); err != nil {
-			return "", err
-		}
-	}
+func (c *SSHClient) Run(shell string) (string, error) {
 	session, err := c.Client.NewSession()
 	if err != nil {
 		return "", err
@@ -74,11 +68,10 @@ func (c *ConnInfo) Run(shell string) (string, error) {
 	defer session.Close()
 	buf, err := session.CombinedOutput(shell)
 
-	c.LastResult = string(buf)
-	return c.LastResult, err
+	return string(buf), err
 }
 
-func (c *ConnInfo) Close() {
+func (c *SSHClient) Close() {
 	_ = c.Client.Close()
 }
 
@@ -88,7 +81,7 @@ type SshConn struct {
 	Session     *gossh.Session
 }
 
-func (c *ConnInfo) NewSshConn(cols, rows int) (*SshConn, error) {
+func (c *SSHClient) NewSshConn(cols, rows int) (*SshConn, error) {
 	sshSession, err := c.Client.NewSession()
 	if err != nil {
 		return nil, err