fix: missing ssh client closes

The webssh backend forgets to close ssh connects created by Sftp & CheckSSH.
It can make tons of orphaned sshd processes exhausting a target host.

This commit fixes the closing mechanism.

Author: cdfmlr

controller/common.go

@@ -30,7 +30,10 @@ func CheckSSH(c *gin.Context) *ResponseBody {
 		responseBody.Msg = err.Error()
 		return &responseBody
 	}
+
 	err = sshClient.GenerateClient()
+	defer sshClient.Close()
+
 	if err != nil {
 		fmt.Println(err)
 		responseBody.Msg = err.Error()

controller/file.go

@@ -112,7 +112,7 @@ func UploadFile(c *gin.Context) *ResponseBody {
 		responseBody.Msg = err.Error()
 		return &responseBody
 	}
-	defer sshClient.Sftp.Close()
+	defer sshClient.Close()
 	file, header, err := c.Request.FormFile("file")
 	if err != nil {
 		responseBody.Msg = err.Error()
@@ -157,7 +157,7 @@ func DownloadFile(c *gin.Context) *ResponseBody {
 		responseBody.Msg = err.Error()
 		return &responseBody
 	}
-	defer sshClient.Sftp.Close()
+	defer sshClient.Close()
 	if sftpFile, err := sshClient.Download(path); err != nil {
 		fmt.Println(err)
 		responseBody.Msg = err.Error()
@@ -228,7 +228,7 @@ func FileList(c *gin.Context) *ResponseBody {
 		responseBody.Msg = err.Error()
 		return &responseBody
 	}
-	defer sshClient.Sftp.Close()
+	defer sshClient.Close()
 	files, err := sshClient.Sftp.ReadDir(path)
 	if err != nil {
 		if strings.Contains(err.Error(), "exist") {

core/models.go

@@ -5,6 +5,7 @@ import (
 	"github.com/pkg/sftp"
 	"golang.org/x/crypto/ssh"
 	"io"
+	"log"
 	"unicode/utf8"
 )
 
@@ -67,3 +68,31 @@ func NewSSHClient() SSHClient {
 	client.Port = 22
 	return client
 }
+
+// Close all closable fields of SSHClient that is opened:
+//
+//	StdinPipe, Session, Sftp, Client
+func (sclient *SSHClient) Close() {
+	defer func() { // just in case
+		if err := recover(); err != nil {
+			log.Println("SSHClient Close recover from panic: ", err)
+		}
+	}()
+
+	if sclient.StdinPipe != nil {
+		sclient.StdinPipe.Close()
+		sclient.StdinPipe = nil
+	}
+	if sclient.Session != nil {
+		sclient.Session.Close()
+		sclient.Session = nil
+	}
+	if sclient.Sftp != nil {
+		sclient.Sftp.Close()
+		sclient.Sftp = nil
+	}
+	if sclient.Client != nil {
+		sclient.Client.Close()
+		sclient.Client = nil
+	}
+}

core/sshclient.go

@@ -121,6 +121,7 @@ func (sclient *SSHClient) Connect(ws *websocket.Conn, timeout time.Duration, clo
 				err := sclient.Session.WindowChange(rows, cols)
 				if err != nil {
 					log.Println(err)
+					close(stopCh)
 					return
 				}
 				continue
@@ -135,13 +136,8 @@ func (sclient *SSHClient) Connect(ws *websocket.Conn, timeout time.Duration, clo
 
 	defer func() {
 		ws.Close()
-		if sclient.Session != nil {
-			sclient.StdinPipe.Close()
-			sclient.Session.Close()
-			sclient.Client.Close()
-			sclient.Session = nil
-			sclient.Client = nil
-		}
+		sclient.Close()
+
 		if err := recover(); err != nil {
 			log.Println(err)
 		}