fix: ws接收消息处理改为函数回调易于扩展

2024-10-10 20:39:15 +08:00
parent 4f9f91ea5a
commit 9127865b12
6 changed files with 338 additions and 318 deletions
--- a/src/modules/ws/controller/ws_ssh.go
+++ b/src/modules/ws/controller/ws_ssh.go
@@ -0,0 +1,118 @@
+package controller
+
+import (
+	"encoding/json"
+	"fmt"
+	"time"
+
+	"be.ems/src/framework/i18n"
+	"be.ems/src/framework/logger"
+	"be.ems/src/framework/utils/ctx"
+	"be.ems/src/framework/utils/ssh"
+	"be.ems/src/framework/vo/result"
+	neService "be.ems/src/modules/network_element/service"
+
+	"github.com/gin-gonic/gin"
+)
+
+// SSH 终端
+//
+// GET /ssh?hostId=1&cols=80&rows=40
+func (s *WSController) SSH(c *gin.Context) {
+	language := ctx.AcceptLanguage(c)
+	var query struct {
+		HostId string `form:"hostId" binding:"required"` // 连接主机ID
+		Cols   int    `form:"cols"`                      // 终端单行字符数
+		Rows   int    `form:"rows"`                      // 终端显示行数
+	}
+	if err := c.ShouldBindQuery(&query); err != nil {
+		c.JSON(400, result.CodeMsg(400, i18n.TKey(language, "app.common.err400")))
+		return
+	}
+	if query.Cols < 80 || query.Cols > 400 {
+		query.Cols = 80
+	}
+	if query.Rows < 40 || query.Rows > 1200 {
+		query.Rows = 40
+	}
+
+	// 登录用户信息
+	loginUser, err := ctx.LoginUser(c)
+	if err != nil {
+		c.JSON(401, result.CodeMsg(401, i18n.TKey(language, err.Error())))
+		return
+	}
+
+	neHost := neService.NewNeHostImpl.SelectById(query.HostId)
+	if neHost.HostID != query.HostId || neHost.HostType != "ssh" {
+		// 没有可访问主机信息数据！
+		c.JSON(200, result.ErrMsg(i18n.TKey(language, "neHost.noData")))
+		return
+	}
+
+	// 创建链接SSH客户端
+	var connSSH ssh.ConnSSH
+	neHost.CopyTo(&connSSH)
+	var client *ssh.ConnSSH
+	var clientErr error
+	if neHost.AuthMode == "2" {
+		client, clientErr = connSSH.NewClientByLocalPrivate()
+	} else {
+		client, clientErr = connSSH.NewClient()
+	}
+	if clientErr != nil {
+		// 连接主机失败，请检查连接参数后重试
+		c.JSON(200, result.ErrMsg(i18n.TKey(language, "neHost.errByHostInfo")))
+		return
+	}
+	defer client.Close()
+
+	// 创建SSH客户端会话
+	clientSession, err := client.NewClientSession(query.Cols, query.Rows)
+	if err != nil {
+		// 连接主机失败，请检查连接参数后重试
+		c.JSON(200, result.ErrMsg(i18n.TKey(language, "neHost.errByHostInfo")))
+		return
+	}
+	defer clientSession.Close()
+
+	// 将 HTTP 连接升级为 WebSocket 连接
+	wsConn := s.wsService.UpgraderWs(c.Writer, c.Request)
+	if wsConn == nil {
+		return
+	}
+	defer wsConn.Close()
+
+	wsClient := s.wsService.ClientCreate(loginUser.UserID, nil, wsConn, clientSession)
+	go s.wsService.ClientWriteListen(wsClient)
+	go s.wsService.ClientReadListen(wsClient, s.wsReceiveService.Shell)
+
+	// 实时读取SSH消息直接输出
+	msTicker := time.NewTicker(100 * time.Millisecond)
+	defer msTicker.Stop()
+	for {
+		select {
+		case ms := <-msTicker.C:
+			outputByte := clientSession.Read()
+			if len(outputByte) > 0 {
+				outputStr := string(outputByte)
+				msgByte, _ := json.Marshal(result.Ok(map[string]any{
+					"requestId": fmt.Sprintf("ssh_%s_%d", neHost.HostID, ms.UnixMilli()),
+					"data":      outputStr,
+				}))
+				wsClient.MsgChan <- msgByte
+
+				// 退出ssh登录
+				// if strings.LastIndex(outputStr, "logout\r\n") != -1 {
+				// 	time.Sleep(1 * time.Second)
+				// 	s.wsService.CloseClient(wsClient.ID)
+				// 	return
+				// }
+			}
+		case <-wsClient.StopChan: // 等待停止信号
+			s.wsService.ClientClose(wsClient.ID)
+			logger.Infof("ws Stop Client UID %s", wsClient.BindUid)
+			return
+		}
+	}
+}