From e8123352ae2d000fb175204ec2b9bca490410ae0 Mon Sep 17 00:00:00 2001
From: lwch <lwch748@gmail.com>
Date: Wed, 27 Jul 2022 00:20:14 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=AD=A3websocket=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E4=BC=A0=E8=BE=93=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 code/client/rule/code/workspace.go            | 15 ++++++++++++
 .../rule/code/workspace_remote_response.go    | 24 +++++++++++++------
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/code/client/rule/code/workspace.go b/code/client/rule/code/workspace.go
index dc42cc5..86fb3c0 100644
--- a/code/client/rule/code/workspace.go
+++ b/code/client/rule/code/workspace.go
@@ -179,7 +179,12 @@ func (ws *Workspace) remoteRead() {
 		case network.Msg_code_request:
 			go ws.handleRequest(msg)
 		case network.Msg_code_connect:
+			ws.Lock()
+			ws.onMessage[msg.GetCsconn().GetRequestId()] = make(chan *network.Msg, 1024)
+			ws.Unlock()
 			go ws.handleConnect(msg)
+		case network.Msg_code_data:
+			ws.writeMessage(msg.GetCsdata().GetRequestId(), msg)
 		}
 	}
 }
@@ -207,6 +212,10 @@ func (ws *Workspace) localRead() {
 			ws.writeMessage(msg.GetCsrepHdr().GetRequestId(), msg)
 		case network.Msg_code_response_body:
 			ws.writeMessage(msg.GetCsrepBody().GetRequestId(), msg)
+		case network.Msg_code_connect_response:
+			ws.writeMessage(msg.GetCsconnRep().GetRequestId(), msg)
+		case network.Msg_code_data:
+			ws.writeMessage(msg.GetCsdata().GetRequestId(), msg)
 		}
 	}
 }
@@ -226,6 +235,12 @@ func (ws *Workspace) writeMessage(reqID uint64, msg *network.Msg) {
 	}
 }
 
+func (ws *Workspace) chanResponse(reqID uint64) <-chan *network.Msg {
+	ws.RLock()
+	defer ws.RUnlock()
+	return ws.onMessage[reqID]
+}
+
 func (ws *Workspace) onResponse(reqID uint64) *network.Msg {
 	ws.RLock()
 	ch := ws.onMessage[reqID]
diff --git a/code/client/rule/code/workspace_remote_response.go b/code/client/rule/code/workspace_remote_response.go
index 8a44414..ffe4dcf 100644
--- a/code/client/rule/code/workspace_remote_response.go
+++ b/code/client/rule/code/workspace_remote_response.go
@@ -100,18 +100,26 @@ func (ws *Workspace) handleConnect(msg *network.Msg) {
 func (ws *Workspace) ws2remote(wg *sync.WaitGroup, reqID uint64, conn *websocket.Conn) {
 	defer wg.Done()
 	defer conn.Close()
+	defer ws.closeMessage(reqID)
 	for {
 		t, data, err := conn.ReadMessage()
 		if err != nil {
 			logging.Error("read_message [%s] [%s]: %v", ws.id, ws.name, err)
-			send := ws.remote.SendCodeData(ws.target, ws.id, reqID,
-				false, 0, []byte(err.Error()))
-			ws.sendBytes += send
-			ws.sendPacket++
+			ws.SendData(reqID, false, websocket.TextMessage, []byte(err.Error()))
 			return
 		}
+		ws.SendData(reqID, true, t, data)
+	}
+}
+
+func (ws *Workspace) SendData(reqID uint64, ok bool, t int, body []byte) {
+	for i := 0; i < len(body); i += 32 * 1024 {
+		end := i + 32*1024
+		if end > len(body) {
+			end = len(body)
+		}
 		send := ws.remote.SendCodeData(ws.target, ws.id, reqID,
-			true, t, data)
+			ok, t, body[i:end])
 		ws.sendBytes += send
 		ws.sendPacket++
 	}
@@ -120,10 +128,12 @@ func (ws *Workspace) ws2remote(wg *sync.WaitGroup, reqID uint64, conn *websocket
 func (ws *Workspace) remote2ws(wg *sync.WaitGroup, reqID uint64, conn *websocket.Conn) {
 	defer wg.Done()
 	defer conn.Close()
+	defer ws.closeMessage(reqID)
+	ch := ws.chanResponse(reqID)
 	for {
-		msg := ws.onResponse(reqID)
+		msg := <-ch
 		if msg == nil {
-			continue
+			return
 		}
 		if msg.GetXType() != network.Msg_code_data {
 			logging.Error("got invalid message type [%s] [%s]: %s",