From 2e4be2800f70478f5d70ff1107d68ffa020f2a36 Mon Sep 17 00:00:00 2001
From: laodaming <11058467+laudamine@user.noreply.gitee.com>
Date: Wed, 11 Oct 2023 18:54:14 +0800
Subject: [PATCH] fix

---
 .../internal/handler/datatransferhandler.go   | 15 -----------
 .../internal/logic/datatransferlogic.go       | 27 +++++--------------
 server/websocket/internal/types/types.go      |  4 +++
 server_api/websocket.api                      |  6 ++++-
 4 files changed, 16 insertions(+), 36 deletions(-)
 delete mode 100644 server/websocket/internal/handler/datatransferhandler.go

diff --git a/server/websocket/internal/handler/datatransferhandler.go b/server/websocket/internal/handler/datatransferhandler.go
deleted file mode 100644
index c167b31c..00000000
--- a/server/websocket/internal/handler/datatransferhandler.go
+++ /dev/null
@@ -1,15 +0,0 @@
-package handler
-
-import (
-	"fusenapi/server/websocket/internal/logic"
-	"fusenapi/server/websocket/internal/svc"
-	"net/http"
-)
-
-func DataTransferHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
-		// 创建一个业务逻辑层实例
-		l := logic.NewDataTransferLogic(r.Context(), svcCtx)
-		l.DataTransfer(w, r)
-	}
-}
diff --git a/server/websocket/internal/logic/datatransferlogic.go b/server/websocket/internal/logic/datatransferlogic.go
index beb5b5c8..c15caf01 100644
--- a/server/websocket/internal/logic/datatransferlogic.go
+++ b/server/websocket/internal/logic/datatransferlogic.go
@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"fusenapi/constants"
+	"fusenapi/server/websocket/internal/types"
 	"fusenapi/utils/auth"
 	"fusenapi/utils/basic"
 	"fusenapi/utils/encryption_decryption"
@@ -107,7 +108,7 @@ type wsConnectItem struct {
 }
 
 // 请求建立连接，升级websocket协议
-func (l *DataTransferLogic) DataTransfer(w http.ResponseWriter, r *http.Request) {
+func (l *DataTransferLogic) DataTransfer(req *types.DataTransferReq, w http.ResponseWriter, r *http.Request) {
 	origin := r.Header.Get("Origin")
 	//判断是不是允许的跨域
 	if !openDebug {
@@ -119,27 +120,13 @@ func (l *DataTransferLogic) DataTransfer(w http.ResponseWriter, r *http.Request)
 		}
 	}
 	//把子协议携带的token设置到标准token头信息中
-	secWebsocketProtocol := r.Header.Get("Sec-Websocket-Protocol")
-	oldWid := ""
+	token := r.Header.Get("Sec-Websocket-Protocol")
+	oldWid := req.Wid
 	//有token是正常用户，无则是白板用户，也可以连接
-	if secWebsocketProtocol != "" {
-		s := strings.Split(secWebsocketProtocol, ",")
-		if len(s) != 2 {
-			w.Write([]byte("invalid secWebsocketProtocol param"))
-			return
-		}
-		//有效token
-		token := strings.Trim(s[0], " ")
-		if token != "empty_token" {
-			r.Header.Set("Authorization", "Bearer "+token)
-		}
-		//有效wid
-		wid := strings.Trim(s[1], " ")
-		if wid != "empty_wid" {
-			oldWid = wid
-		}
+	if token != "" {
+		r.Header.Set("Authorization", "Bearer "+token)
 		//设置Sec-Websocket-Protocol
-		upgrader.Subprotocols = []string{secWebsocketProtocol}
+		upgrader.Subprotocols = []string{token}
 	}
 	//判断下是否火狐浏览器(获取浏览器第一条消息返回有收不到的bug需要延迟1秒)
 	userAgent := r.Header.Get("User-Agent")
diff --git a/server/websocket/internal/types/types.go b/server/websocket/internal/types/types.go
index 1308ca96..f5fbe2c6 100644
--- a/server/websocket/internal/types/types.go
+++ b/server/websocket/internal/types/types.go
@@ -5,6 +5,10 @@ import (
 	"fusenapi/utils/basic"
 )
 
+type DataTransferReq struct {
+	Wid string `form:"wid"`
+}
+
 type RenderNotifyReq struct {
 	TaskId  string `json:"task_id"`
 	UserId  int64  `json:"user_id"`
diff --git a/server_api/websocket.api b/server_api/websocket.api
index 2cf9b172..62dfd83b 100644
--- a/server_api/websocket.api
+++ b/server_api/websocket.api
@@ -11,7 +11,7 @@ import "basic.api"
 service websocket {
 	//websocket数据交互
 	@handler DataTransferHandler
-	get /api/websocket/data_transfer(request) returns (response);
+	get /api/websocket/data_transfer(DataTransferReq) returns (response);
 	//云渲染完了通知接口
 	@handler RenderNotifyHandler
 	post /api/websocket/render_notify(RenderNotifyReq) returns (response);
@@ -20,6 +20,10 @@ service websocket {
 	post /api/websocket/common_notify(CommonNotifyReq) returns (response);
 }
 
+//websocket数据交互[
+type DataTransferReq {
+	Wid string `form:"wid"`
+}
 //渲染完了通知接口
 type RenderNotifyReq {
 	TaskId  string `json:"task_id"`