From 580bd5cc9316d04ef15e123f17dbe4eee70a250a Mon Sep 17 00:00:00 2001
From: Gordon <46924906+FGadvancer@users.noreply.github.com>
Date: Tue, 16 Apr 2024 20:07:22 +0800
Subject: [PATCH] refactor: websocket update info.

---
 internal/msggateway/constant.go    |  2 +-
 internal/msggateway/context.go     |  4 ++--
 internal/msggateway/http_error.go  |  6 +++++-
 internal/msggateway/long_conn.go   | 18 ++++++++++++++----
 internal/msggateway/n_ws_server.go | 14 ++++++++++++--
 5 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/internal/msggateway/constant.go b/internal/msggateway/constant.go
index 477fcfde6..64664ac0a 100644
--- a/internal/msggateway/constant.go
+++ b/internal/msggateway/constant.go
@@ -26,7 +26,7 @@ const (
 	Compression             = "compression"
 	GzipCompressionProtocol = "gzip"
 	BackgroundStatus        = "isBackground"
-	ErrResp                 = "errResp"
+	SendResponse            = "isMsgResp"
 )
 
 const (
diff --git a/internal/msggateway/context.go b/internal/msggateway/context.go
index 87ed25d75..6c80ece1b 100644
--- a/internal/msggateway/context.go
+++ b/internal/msggateway/context.go
@@ -149,8 +149,8 @@ func (c *UserConnContext) GetCompression() bool {
 	return false
 }
 
-func (c *UserConnContext) ShouldSendError() bool {
-	errResp, exists := c.Query(ErrResp)
+func (c *UserConnContext) ShouldSendResp() bool {
+	errResp, exists := c.Query(SendResponse)
 	if exists {
 		b, err := strconv.ParseBool(errResp)
 		if err != nil {
diff --git a/internal/msggateway/http_error.go b/internal/msggateway/http_error.go
index 4732680b2..8d9d03522 100644
--- a/internal/msggateway/http_error.go
+++ b/internal/msggateway/http_error.go
@@ -14,8 +14,12 @@
 
 package msggateway
 
-import "github.com/openimsdk/tools/apiresp"
+import (
+	"github.com/openimsdk/tools/apiresp"
+	"github.com/openimsdk/tools/log"
+)
 
 func httpError(ctx *UserConnContext, err error) {
+	log.ZWarn(ctx, "ws connection error", err)
 	apiresp.HttpError(ctx.RespWriter, err)
 }
diff --git a/internal/msggateway/long_conn.go b/internal/msggateway/long_conn.go
index 92bc189b1..7d5bef4c3 100644
--- a/internal/msggateway/long_conn.go
+++ b/internal/msggateway/long_conn.go
@@ -145,7 +145,7 @@ func (d *GWebSocket) SetPingHandler(handler PingPongHandler) {
 	d.conn.SetPingHandler(handler)
 }
 
-func (d *GWebSocket) RespErrInfo(err error, w http.ResponseWriter, r *http.Request) error {
+func (d *GWebSocket) RespondWithError(err error, w http.ResponseWriter, r *http.Request) error {
 	if err := d.GenerateLongConn(w, r); err != nil {
 		return err
 	}
@@ -163,6 +163,16 @@ func (d *GWebSocket) RespErrInfo(err error, w http.ResponseWriter, r *http.Reque
 	return nil
 }
 
-// func (d *GWebSocket) CheckSendConnDiffNow() bool {
-//	return d.conn == d.sendConn
-//}
+func (d *GWebSocket) RespondWithSuccess() error {
+	data, err := json.Marshal(apiresp.ParseError(nil))
+	if err != nil {
+		_ = d.Close()
+		return errs.WrapMsg(err, "json marshal failed")
+	}
+
+	if err := d.WriteMessage(MessageText, data); err != nil {
+		_ = d.Close()
+		return errs.WrapMsg(err, "WriteMessage failed")
+	}
+	return nil
+}
diff --git a/internal/msggateway/n_ws_server.go b/internal/msggateway/n_ws_server.go
index 21213fca8..7af6c2dea 100644
--- a/internal/msggateway/n_ws_server.go
+++ b/internal/msggateway/n_ws_server.go
@@ -391,11 +391,11 @@ func (ws *WsServer) wsHandler(w http.ResponseWriter, r *http.Request) {
 	resp, err := ws.authClient.ParseToken(connContext, connContext.GetToken())
 	if err != nil {
 		// If there's an error parsing the Token, decide whether to send the error message via WebSocket based on the context flag
-		shouldSendError := connContext.ShouldSendError()
+		shouldSendError := connContext.ShouldSendResp()
 		if shouldSendError {
 			// Create a WebSocket connection object and attempt to send the error message via WebSocket
 			wsLongConn := newGWebSocket(WebSocket, ws.handshakeTimeout, ws.writeBufferSize)
-			if err := wsLongConn.RespErrInfo(err, w, r); err == nil {
+			if err := wsLongConn.RespondWithError(err, w, r); err == nil {
 				// If the error message is successfully sent via WebSocket, stop processing
 				return
 			}
@@ -419,6 +419,16 @@ func (ws *WsServer) wsHandler(w http.ResponseWriter, r *http.Request) {
 		// If creating the long connection fails, return an error via HTTP and stop processing
 		httpError(connContext, err)
 		return
+	} else {
+		// Check if a normal response should be sent via WebSocket
+		shouldSendSuccessResp := connContext.ShouldSendResp()
+		if shouldSendSuccessResp {
+			// Attempt to send a success message through WebSocket
+			if err := wsLongConn.RespondWithSuccess(); err == nil {
+				// If the success message is successfully sent, end further processing
+				return
+			}
+		}
 	}
 
 	// Retrieve a client object from the client pool, reset its state, and associate it with the current WebSocket long connection