perf: 优化ws模块协程资源消耗

2024-08-07 15:27:36 +08:00
parent 712a9fee0b
commit 0f98508169
7 changed files with 376 additions and 145 deletions
--- a/src/modules/ws/service/ws_receive.impl.go
+++ b/src/modules/ws/service/ws_receive.impl.go
@@ -20,12 +20,22 @@ var NewWSReceiveImpl = &WSReceiveImpl{}
 // WSReceiveImpl WebSocket消息接收处理 服务层处理
 type WSReceiveImpl struct{}

-// AsyncReceive 接收业务异步处理
-func (s *WSReceiveImpl) AsyncReceive(client *model.WSClient, reqMsg model.WSRequest) {
+// Commont 接收通用业务处理
+func (s *WSReceiveImpl) close(client *model.WSClient) {
+	// 主动关闭
+	resultByte, _ := json.Marshal(result.OkMsg("user initiated closure"))
+	client.MsgChan <- resultByte
+	// 等待1s后关闭连接
+	time.Sleep(1 * time.Second)
+	NewWSImpl.ClientClose(client.ID)
+}
+
+// Commont 接收通用业务处理
+func (s *WSReceiveImpl) Commont(client *model.WSClient, reqMsg model.WSRequest) {
 	// 必传requestId确认消息
 	if reqMsg.RequestID == "" {
 		msg := "message requestId is required"
-		logger.Infof("ws AsyncReceive UID %s err: %s", client.BindUid, msg)
+		logger.Infof("ws Commont UID %s err: %s", client.BindUid, msg)
 		msgByte, _ := json.Marshal(result.ErrMsg(msg))
 		client.MsgChan <- msgByte
 		return
@@ -36,34 +46,8 @@ func (s *WSReceiveImpl) AsyncReceive(client *model.WSClient, reqMsg model.WSRequ

 	switch reqMsg.Type {
 	case "close":
-		// 主动关闭
-		resultByte, _ := json.Marshal(result.OkMsg("user initiated closure"))
-		client.MsgChan <- resultByte
-		// 等待1s后关闭连接
-		time.Sleep(1 * time.Second)
-		client.StopChan <- struct{}{}
-	case "ssh":
-		// SSH会话消息接收直接写入会话
-		command := reqMsg.Data.(string)
-		sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
-		_, err = sshClientSession.Write(command)
-	case "ssh_resize":
-		// SSH会话窗口重置
-		msgByte, _ := json.Marshal(reqMsg.Data)
-		var data struct {
-			Cols int `json:"cols"`
-			Rows int `json:"rows"`
-		}
-		err = json.Unmarshal(msgByte, &data)
-		if err == nil {
-			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
-			err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
-		}
-	case "telnet":
-		// Telnet会话消息接收直接写入会话
-		command := reqMsg.Data.(string)
-		telnetClientSession := client.ChildConn.(*telnet.TelnetClientSession)
-		_, err = telnetClientSession.Write(command)
+		s.close(client)
+		return
 	case "ps":
 		resByte, err = processor.GetProcessData(reqMsg.RequestID, reqMsg.Data)
 	case "net":
@@ -85,7 +69,146 @@ func (s *WSReceiveImpl) AsyncReceive(client *model.WSClient, reqMsg model.WSRequ
 	}

 	if err != nil {
-		logger.Warnf("ws AsyncReceive UID %s err: %s", client.BindUid, err.Error())
+		logger.Warnf("ws Commont UID %s err: %s", client.BindUid, err.Error())
+		msgByte, _ := json.Marshal(result.ErrMsg(err.Error()))
+		client.MsgChan <- msgByte
+		return
+	}
+	if len(resByte) > 0 {
+		client.MsgChan <- resByte
+	}
+}
+
+// Shell 接收终端交互业务处理
+func (s *WSReceiveImpl) Shell(client *model.WSClient, reqMsg model.WSRequest) {
+	// 必传requestId确认消息
+	if reqMsg.RequestID == "" {
+		msg := "message requestId is required"
+		logger.Infof("ws Shell UID %s err: %s", client.BindUid, msg)
+		msgByte, _ := json.Marshal(result.ErrMsg(msg))
+		client.MsgChan <- msgByte
+		return
+	}
+
+	var resByte []byte
+	var err error
+
+	switch reqMsg.Type {
+	case "close":
+		s.close(client)
+		return
+	case "ssh":
+		// SSH会话消息接收写入会话
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type ssh not connected")
+		} else {
+			command := reqMsg.Data.(string)
+			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+			_, err = sshClientSession.Write(command)
+		}
+	case "ssh_resize":
+		// SSH会话窗口重置
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type ssh not connected")
+		} else {
+			msgByte, _ := json.Marshal(reqMsg.Data)
+			var data struct {
+				Cols int `json:"cols"`
+				Rows int `json:"rows"`
+			}
+			err = json.Unmarshal(msgByte, &data)
+			if err == nil {
+				sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+				err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
+			}
+		}
+	case "telnet":
+		// Telnet会话消息接收写入会话
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type telnet not connected")
+		} else {
+			command := reqMsg.Data.(string)
+			telnetClientSession := client.ChildConn.(*telnet.TelnetClientSession)
+			_, err = telnetClientSession.Write(command)
+		}
+	default:
+		err = fmt.Errorf("message type not supported")
+	}
+
+	if err != nil {
+		logger.Warnf("ws Shell UID %s err: %s", client.BindUid, err.Error())
+		msgByte, _ := json.Marshal(result.ErrMsg(err.Error()))
+		client.MsgChan <- msgByte
+		if err == io.EOF {
+			// 等待1s后关闭连接
+			time.Sleep(1 * time.Second)
+			client.StopChan <- struct{}{}
+		}
+		return
+	}
+	if len(resByte) > 0 {
+		client.MsgChan <- resByte
+	}
+}
+
+// ShellView 接收查看文件终端交互业务处理
+func (s *WSReceiveImpl) ShellView(client *model.WSClient, reqMsg model.WSRequest) {
+	// 必传requestId确认消息
+	if reqMsg.RequestID == "" {
+		msg := "message requestId is required"
+		logger.Infof("ws Shell UID %s err: %s", client.BindUid, msg)
+		msgByte, _ := json.Marshal(result.ErrMsg(msg))
+		client.MsgChan <- msgByte
+		return
+	}
+
+	var resByte []byte
+	var err error
+
+	switch reqMsg.Type {
+	case "close":
+		s.close(client)
+		return
+	case "ssh":
+		// SSH会话消息接收写入会话
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type ssh not connected")
+		} else {
+			command := reqMsg.Data.(string)
+			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+			_, err = sshClientSession.Write(command)
+		}
+	case "ssh_resize":
+		// SSH会话窗口重置
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type ssh not connected")
+		} else {
+			msgByte, _ := json.Marshal(reqMsg.Data)
+			var data struct {
+				Cols int `json:"cols"`
+				Rows int `json:"rows"`
+			}
+			err = json.Unmarshal(msgByte, &data)
+			if err == nil {
+				sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+				err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
+			}
+		}
+	case "telnet":
+		// Telnet会话消息接收写入会话
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type telnet not connected")
+		} else {
+			command := reqMsg.Data.(string)
+			telnetClientSession := client.ChildConn.(*telnet.TelnetClientSession)
+			_, err = telnetClientSession.Write(command)
+		}
+	default:
+		err = fmt.Errorf("message type not supported")
+	}
+
+	if err != nil {
+		logger.Warnf("ws Shell UID %s err: %s", client.BindUid, err.Error())
 		msgByte, _ := json.Marshal(result.ErrMsg(err.Error()))
 		client.MsgChan <- msgByte
 		if err == io.EOF {