From 0142e04ef55fa2d8754c68efd2a25b5df17d1470 Mon Sep 17 00:00:00 2001
From: laodaming <11058467+laudamine@user.noreply.gitee.com>
Date: Wed, 26 Jul 2023 17:30:35 +0800
Subject: [PATCH] fix

---
 .../internal/logic/datatransferlogic.go       | 57 +++++++++----------
 1 file changed, 26 insertions(+), 31 deletions(-)

diff --git a/server/websocket/internal/logic/datatransferlogic.go b/server/websocket/internal/logic/datatransferlogic.go
index ae4663cb..0d8571f6 100644
--- a/server/websocket/internal/logic/datatransferlogic.go
+++ b/server/websocket/internal/logic/datatransferlogic.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"fusenapi/constants"
 	"fusenapi/server/websocket/internal/types"
+	"fusenapi/utils/auth"
 	"github.com/google/uuid"
 	"github.com/gorilla/websocket"
 	"github.com/zeromicro/go-zero/rest/httpx"
@@ -79,44 +80,17 @@ func (l *DataTransferLogic) DataTransfer(svcCtx *svc.ServiceContext, w http.Resp
 	}
 	defer conn.Close()
 	rsp := types.DataTransferData{}
-	/*isAuth := true
-	// 解析JWT token,并对空用户进行判断
-	claims, err := svcCtx.ParseJwtToken(r)
-	// 如果解析JWT token出错,则返回未授权的JSON响应并记录错误消息
-	if err != nil {
+	//鉴权不成功10秒后断开
+	if !l.checkAuth(svcCtx, r) {
 		rsp.T = constants.WEBSOCKET_UNAUTH
 		rsp.D = "unAuth"
 		b, _ := json.Marshal(rsp)
 		_ = conn.WriteMessage(websocket.TextMessage, b)
-		isAuth = false
-	}
-	if claims != nil {
-		// 从token中获取对应的用户信息
-		_, err = auth.GetUserInfoFormMapClaims(claims)
-		// 如果获取用户信息出错,则返回未授权的JSON响应并记录错误消息
-		if err != nil {
-			rsp.T = constants.WEBSOCKET_UNAUTH
-			rsp.D = "unAuth!!"
-			b, _ := json.Marshal(rsp)
-			_ = conn.WriteMessage(websocket.TextMessage, b)
-			isAuth = false
-		}
-	} else {
-		// 如果claims为nil,则认为用户身份为白板用户
-		rsp.T = constants.WEBSOCKET_UNAUTH
-		rsp.D = "unAuth!!!"
-		b, _ := json.Marshal(rsp)
-		_ = conn.WriteMessage(websocket.TextMessage, b)
-		isAuth = false
-	}
-	//不是授权的连接(10秒后关闭)
-	if !isAuth {
 		select {
-		case <-time.After(time.Second * 10):
-			conn.Close()
+		case <-time.After(time.Second * 10): //10秒后断开
 			return
 		}
-	}*/
+	}
 	//生成连接唯一标识
 	flag := uuid.New().String() + time.Now().Format("20060102150405")
 	ws := wsConnectItem{
@@ -150,6 +124,27 @@ func (l *DataTransferLogic) DataTransfer(svcCtx *svc.ServiceContext, w http.Resp
 	ws.heartbeat()
 }
 
+// 鉴权
+func (l *DataTransferLogic) checkAuth(svcCtx *svc.ServiceContext, r *http.Request) bool {
+	// 解析JWT token,并对空用户进行判断
+	claims, err := svcCtx.ParseJwtToken(r)
+	// 如果解析JWT token出错,则返回未授权的JSON响应并记录错误消息
+	if err != nil {
+		return false
+	}
+	if claims != nil {
+		// 从token中获取对应的用户信息
+		_, err = auth.GetUserInfoFormMapClaims(claims)
+		// 如果获取用户信息出错,则返回未授权的JSON响应并记录错误消息
+		if err != nil {
+			return false
+		}
+	} else {
+		return false
+	}
+	return true
+}
+
 // 心跳
 func (w *wsConnectItem) heartbeat() {
 	rsp := types.DataTransferData{