package logic import ( "bytes" "encoding/json" "fmt" "fusenapi/constants" "fusenapi/server/websocket/internal/types" "fusenapi/utils/auth" "fusenapi/utils/id_generator" "github.com/gorilla/websocket" "net/http" "sync" "time" "context" "fusenapi/server/websocket/internal/svc" "github.com/zeromicro/go-zero/core/logx" ) type DataTransferLogic struct { logx.Logger ctx context.Context svcCtx *svc.ServiceContext } func NewDataTransferLogic(ctx context.Context, svcCtx *svc.ServiceContext) *DataTransferLogic { return &DataTransferLogic{ Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx, } } var ( //全局websocketid生成器 websocketIdGenerator = id_generator.NewWebsocketId(1) //临时缓存对象池 buffPool = sync.Pool{ New: func() interface{} { return bytes.Buffer{} }, } //升级websocket upgrade = websocket.Upgrader{ ReadBufferSize: 1024 * 10, //最大可读取大小 10M //握手超时时间15s HandshakeTimeout: time.Second * 15, //允许跨域 CheckOrigin: func(r *http.Request) bool { return true }, WriteBufferPool: &buffPool, EnableCompression: true, } //连接map池 mapConnPool = sync.Map{} ) // 每个连接的连接基本属性 type wsConnectItem struct { conn *websocket.Conn //websocket的连接 closeChan chan struct{} //ws连接关闭chan isClose bool //是否已经关闭 uniqueId uint64 //ws连接唯一标识 inChan chan []byte //接受消息缓冲通道 outChan chan []byte //发送回客户端的消息 mutex sync.Mutex //互斥锁 renderProperty renderProperty //扩展云渲染属性 } func (l *DataTransferLogic) DataTransfer(svcCtx *svc.ServiceContext, w http.ResponseWriter, r *http.Request) { //升级websocket conn, err := upgrade.Upgrade(w, r, nil) if err != nil { logx.Error("http upgrade websocket err:", err) return } defer conn.Close() w.Header().Set("Connection", "Upgrade") rsp := types.DataTransferData{} //鉴权不成功10秒后断开 /*isAuth, _ := l.checkAuth(svcCtx, r) if !isAuth { time.Sleep(time.Second) //兼容下火狐 rsp.T = constants.WEBSOCKET_UNAUTH b, _ := json.Marshal(rsp) //先发一条正常信息 _ = conn.WriteMessage(websocket.TextMessage, b) //发送关闭信息 _ = conn.WriteMessage(websocket.CloseMessage, nil) return }*/ //生成连接唯一标识 uniqueId := websocketIdGenerator.Get() ws := wsConnectItem{ conn: conn, uniqueId: uniqueId, closeChan: make(chan struct{}, 1), inChan: make(chan []byte, 100), outChan: make(chan []byte, 100), renderProperty: renderProperty{ renderImageTask: make(map[string]struct{}), renderImageTaskCtlChan: make(chan renderImageControlChanItem, 100), }, } //保存连接 mapConnPool.Store(uniqueId, ws) defer ws.close() //把连接成功消息发回去 time.Sleep(time.Second) //兼容下火狐 rsp.T = constants.WEBSOCKET_CONNECT_SUCCESS rsp.D = uniqueId b, _ := json.Marshal(rsp) _ = conn.WriteMessage(websocket.TextMessage, b) //循环读客户端信息 go ws.readLoop() //循环把数据发送给客户端 go ws.writeLoop() //推消息到云渲染 go ws.sendLoop() //操作连接中渲染任务的增加/删除 go ws.operationRenderTask() //心跳 ws.heartbeat() } // 鉴权 func (l *DataTransferLogic) checkAuth(svcCtx *svc.ServiceContext, r *http.Request) (isAuth bool, userInfo *auth.UserInfo) { // 解析JWT token,并对空用户进行判断 claims, err := svcCtx.ParseJwtToken(r) // 如果解析JWT token出错,则返回未授权的JSON响应并记录错误消息 if err != nil { return false, nil } if claims != nil { // 从token中获取对应的用户信息 userInfo, err = auth.GetUserInfoFormMapClaims(claims) // 如果获取用户信息出错,则返回未授权的JSON响应并记录错误消息 if err != nil { return false, nil } } else { return false, nil } return true, userInfo } // 心跳 func (w *wsConnectItem) heartbeat() { tick := time.Tick(time.Second * 5) for { select { case <-w.closeChan: return case <-tick: //发送心跳信息 if err := w.conn.WriteMessage(websocket.PongMessage, nil); err != nil { logx.Error("发送心跳信息异常,关闭连接:", w.uniqueId, err) w.close() return } } } } // 关闭连接 func (w *wsConnectItem) close() { w.mutex.Lock() defer w.mutex.Unlock() logx.Info("websocket:", w.uniqueId, " is closing...") //发送关闭信息 _ = w.conn.WriteMessage(websocket.CloseMessage, nil) w.conn.Close() mapConnPool.Delete(w.uniqueId) if !w.isClose { w.isClose = true close(w.closeChan) } logx.Info("websocket:", w.uniqueId, " is closed") } // 读取输出返回给客户端 func (w *wsConnectItem) writeLoop() { for { select { case <-w.closeChan: //如果关闭了 return case data := <-w.outChan: if err := w.conn.WriteMessage(websocket.TextMessage, data); err != nil { logx.Error("websocket write loop err:", err) w.close() return } } } } // 接受客户端发来的消息 func (w *wsConnectItem) readLoop() { for { select { case <-w.closeChan: //如果关闭了 return default: msgType, data, err := w.conn.ReadMessage() if err != nil { logx.Error("接受信息错误:", err) //关闭连接 w.close() return } //ping的消息不处理 if msgType != websocket.PingMessage { //消息传入缓冲通道 w.inChan <- data } } } } // 把收到的消息发往不同的地方处理 func (w *wsConnectItem) sendLoop() { for { select { case <-w.closeChan: return case data := <-w.inChan: w.dealwithReciveData(data) } } } // 把要传递给客户端的数据放入outchan func (w *wsConnectItem) sendToOutChan(data []byte) { select { case <-w.closeChan: return case w.outChan <- data: logx.Info("notify send render result to out chan") } } // 获取需要渲染图片的map key func (w *wsConnectItem) getRenderImageMapKey(productId, templateTagId int64, algorithmVersion string) string { return fmt.Sprintf("%d-%d-%s", productId, templateTagId, algorithmVersion) } // 处理接受到的数据 func (w *wsConnectItem) dealwithReciveData(data []byte) { var parseInfo types.DataTransferData if err := json.Unmarshal(data, &parseInfo); err != nil { logx.Error("invalid format of websocket message") return } d, _ := json.Marshal(parseInfo.D) //分消息类型给到不同逻辑处理,可扩展 switch parseInfo.T { //图片渲染 case constants.WEBSOCKET_RENDER_IMAGE: go w.SendToCloudRender(d) default: } }