From 77a89f62cd11b596760e90e8c8c688b114ea4abe Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Thu, 19 Aug 2021 10:21:54 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=B8=20=E4=BC=98=E5=8C=96=20Terminal=20?=
 =?UTF-8?q?=E8=BF=9E=E6=8E=A5=E5=81=A5=E5=A3=AE=E6=80=A7=E3=80=81=E8=BE=93?=
 =?UTF-8?q?=E5=85=A5=E7=BC=93=E5=AD=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/dashboard/controller/common_page.go   | 78 +++++++++++++++++------
 resource/template/dashboard/terminal.html |  5 ++
 2 files changed, 65 insertions(+), 18 deletions(-)

diff --git a/cmd/dashboard/controller/common_page.go b/cmd/dashboard/controller/common_page.go
index 51418bb..6c2893e 100644
--- a/cmd/dashboard/controller/common_page.go
+++ b/cmd/dashboard/controller/common_page.go
@@ -161,8 +161,6 @@ func (cp *commonPage) ws(c *gin.Context) {
 }
 
 func (cp *commonPage) terminal(c *gin.Context) {
-	log.Println("terminal connected", c.Request.URL)
-	defer log.Println("terminal disconnected", c.Request.URL)
 	terminalID := c.Param("id")
 	cp.terminalsLock.Lock()
 	if terminalID == "" || cp.terminals[terminalID] == nil {
@@ -181,6 +179,7 @@ func (cp *commonPage) terminal(c *gin.Context) {
 	cp.terminalsLock.Unlock()
 
 	defer func() {
+		// 清理 context
 		cp.terminalsLock.Lock()
 		defer cp.terminalsLock.Unlock()
 		delete(cp.terminals, terminalID)
@@ -272,8 +271,17 @@ func (cp *commonPage) terminal(c *gin.Context) {
 	}
 	defer conn.Close()
 
+	log.Println("terminal connected", isAgent, c.Request.URL)
+	defer log.Println("terminal disconnected", isAgent, c.Request.URL)
+
 	if isAgent {
 		terminal.agentConn = conn
+		defer func() {
+			// Agent断开链接时断开用户连接
+			if terminal.userConn != nil {
+				terminal.userConn.Close()
+			}
+		}()
 	} else {
 		terminal.userConn = conn
 		defer func() {
@@ -284,23 +292,57 @@ func (cp *commonPage) terminal(c *gin.Context) {
 		}()
 	}
 
+	deadlineCh := make(chan interface{})
+	go func() {
+		connectDeadline := time.NewTimer(time.Second * 15)
+		<-connectDeadline.C
+		close(deadlineCh)
+	}()
+
+	dataCh := make(chan []byte)
+	go func() {
+		for {
+			msgType, data, err := conn.ReadMessage()
+			if err != nil {
+				return
+			}
+			// 将文本消息转换为命令输入
+			if msgType == websocket.TextMessage {
+				data = append([]byte{0}, data...)
+			}
+
+			dataCh <- data
+		}
+	}()
+
+	var dataBuffer [][]byte
+	var distConn *websocket.Conn
+
 	for {
-		msgType, data, err := conn.ReadMessage()
-		if err != nil {
-			return
-		}
-		// 将文本消息转换为命令输入
-		if msgType == websocket.TextMessage {
-			data = append([]byte{0}, data...)
-		}
-		// 传递给对方
-		if isAgent {
-			err = terminal.userConn.WriteMessage(websocket.BinaryMessage, data)
-		} else {
-			err = terminal.agentConn.WriteMessage(websocket.BinaryMessage, data)
-		}
-		if err != nil {
-			return
+		select {
+		case <-deadlineCh:
+			if distConn == nil {
+				return
+			}
+		case data := <-dataCh:
+			dataBuffer = append(dataBuffer, data)
+			if distConn == nil {
+				// 传递给对方
+				if isAgent {
+					distConn = terminal.userConn
+				} else {
+					distConn = terminal.agentConn
+				}
+			}
+			if distConn != nil {
+				for i := 0; i < len(dataBuffer); i++ {
+					err = distConn.WriteMessage(websocket.BinaryMessage, dataBuffer[i])
+					if err != nil {
+						return
+					}
+				}
+				dataBuffer = dataBuffer[:0]
+			}
 		}
 	}
 }
diff --git a/resource/template/dashboard/terminal.html b/resource/template/dashboard/terminal.html
index d1825e6..896b988 100644
--- a/resource/template/dashboard/terminal.html
+++ b/resource/template/dashboard/terminal.html
@@ -48,6 +48,11 @@
             onResize()
         }
 
+        socket.onclose = () => {
+            alert('Terminal 连接超时或会话已结束')
+            window.close()
+        }
+
         socket.onerror = () => {
             alert('Terminal 连接失败，请检查 /terminal/* 的 WebSocket 反代情况')
         }