From 01cf8f7eee38e717db3b3518b1ecebccb3a70da7 Mon Sep 17 00:00:00 2001
From: ioito <qu_xuan@icloud.com>
Date: Tue, 25 Jul 2023 14:24:55 +0800
Subject: [PATCH] fix(webconsole): msg write

---
 cmd/climc/shell/compute/webconsole.go | 12 ++++++++--
 pkg/webconsole/server/ssh_server.go   | 33 +++++++++++----------------
 pkg/webconsole/session/ssh_session.go |  6 ++---
 3 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/cmd/climc/shell/compute/webconsole.go b/cmd/climc/shell/compute/webconsole.go
index 0ba026fbd0..a87e381922 100644
--- a/cmd/climc/shell/compute/webconsole.go
+++ b/cmd/climc/shell/compute/webconsole.go
@@ -52,15 +52,23 @@ func init() {
 			return err
 		}
 		connectParams := resp.GetConnectParams()
-		if protocol, err := resp.GetConnectProtocol(); err != nil {
+		protocol, err := resp.GetConnectProtocol()
+		if err != nil {
 			return err
-		} else if !utils.IsInStringArray(protocol, []string{
+		}
+		if !utils.IsInStringArray(protocol, []string{
 			command.PROTOCOL_TTY, webconsole_api.VNC,
 			webconsole_api.SPICE, webconsole_api.WMKS, webconsole_api.WS,
 		}) {
 			fmt.Println(connectParams)
 			return nil
 		}
+		if protocol == webconsole_api.WS {
+			u, err = url.Parse(fmt.Sprintf("%s/ws", opt.WebconsoleUrl))
+			if err != nil {
+				return err
+			}
+		}
 
 		newQuery := url.Values{}
 		newQuery.Set("data", base64.StdEncoding.EncodeToString([]byte(connectParams)))
diff --git a/pkg/webconsole/server/ssh_server.go b/pkg/webconsole/server/ssh_server.go
index 5db1daecc8..88b8ddde50 100644
--- a/pkg/webconsole/server/ssh_server.go
+++ b/pkg/webconsole/server/ssh_server.go
@@ -20,8 +20,8 @@ import (
 	"io"
 	"net/http"
 	"time"
-	"unicode/utf8"
 
+	"github.com/anacrolix/sync"
 	"github.com/gorilla/websocket"
 	"github.com/pkg/errors"
 	"golang.org/x/crypto/ssh"
@@ -62,28 +62,21 @@ func NewSshServer(s *session.SSession) (*WebsocketServer, error) {
 }
 
 type WebSocketBufferWriter struct {
-	s  *session.SSession
-	ws *websocket.Conn
+	s    *session.SSession
+	ws   *websocket.Conn
+	lock sync.Mutex
 }
 
 func (w *WebSocketBufferWriter) Write(p []byte) (int, error) {
-	if !utf8.Valid(p) {
-		bufStr := string(p)
-		buf := make([]rune, 0, len(bufStr))
-		for _, r := range bufStr {
-			if r == utf8.RuneError {
-				buf = append(buf, []rune("@")...)
-			} else {
-				buf = append(buf, r)
-			}
-		}
-		p = []byte(string(buf))
-	}
-	if w.s != nil {
-		go w.s.GetRecorder().Write("", string(p))
-	}
+	w.lock.Lock()
+	defer w.lock.Unlock()
+
+	go w.s.GetRecorder().Write("", string(p))
 	err := w.ws.WriteMessage(websocket.BinaryMessage, p)
-	return len(p), err
+	if err != nil {
+		return 0, err
+	}
+	return len(p), nil
 }
 
 func (s *WebsocketServer) initWs(w http.ResponseWriter, r *http.Request) error {
@@ -148,7 +141,7 @@ func (s *WebsocketServer) initWs(w http.ResponseWriter, r *http.Request) error {
 		ssh.TTY_OP_OSPEED: 14400,
 	}
 
-	err = s.session.RequestPty("xterm", 120, 32, modes)
+	err = s.session.RequestPty("xterm-256color", 120, 32, modes)
 	if err != nil {
 		return errors.Wrapf(err, "request pty xterm")
 	}
diff --git a/pkg/webconsole/session/ssh_session.go b/pkg/webconsole/session/ssh_session.go
index 1f544c8144..9c4628c15f 100644
--- a/pkg/webconsole/session/ssh_session.go
+++ b/pkg/webconsole/session/ssh_session.go
@@ -59,9 +59,6 @@ func NewSshSession(ctx context.Context, us *mcclient.ClientSession, name, ip str
 		Username: username,
 		Password: password,
 	}
-	if len(ret.name) == 0 {
-		ret.name = ret.Username
-	}
 	if port <= 0 {
 		ret.Port = 22
 	}
@@ -85,6 +82,9 @@ func (s *SSshSession) GetProtocol() string {
 }
 
 func (s *SSshSession) GetRecordObject() *recorder.Object {
+	if len(s.name) == 0 {
+		s.name = s.Username
+	}
 	return recorder.NewObject(s.id, s.name, "server", s.Username, jsonutils.Marshal(map[string]interface{}{"ip": s.Host, "port": s.Port}))
 }