fusenapi/server/websocket/internal/logic/datatransferlogic.go
laodaming 31a3966e39 fix
2023-08-07 19:13:16 +08:00

290 lines
6.9 KiB
Go

package logic
import (
"bytes"
"encoding/json"
"fusenapi/constants"
"fusenapi/initalize"
"fusenapi/utils/auth"
"fusenapi/utils/id_generator"
"fusenapi/utils/websocket_data"
"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{
//最大可读取大小 10M
ReadBufferSize: 1024 * 10,
//握手超时时间15s
HandshakeTimeout: time.Second * 15,
//允许跨域
CheckOrigin: func(r *http.Request) bool {
return true
},
//写的缓存池
WriteBufferPool: &buffPool,
//是否支持压缩
EnableCompression: true,
}
//websocket连接存储
mapConnPool = sync.Map{}
)
// 每个连接的连接基本属性
type wsConnectItem struct {
conn *websocket.Conn //websocket的连接
rabbitMq *initalize.RabbitMqHandle
closeChan chan struct{} //ws连接关闭chan
isClose bool //是否已经关闭
uniqueId uint64 //ws连接唯一标识
inChan chan []byte //接受消息缓冲通道
outChan chan []byte //发送回客户端的消息
mutex sync.Mutex //互斥锁
userId int64
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()
//鉴权不成功10秒后断开
var (
userInfo *auth.UserInfo
isAuth bool
)
isAuth, userInfo = l.checkAuth(svcCtx, r)
if !isAuth {
time.Sleep(time.Second) //兼容下火狐
rsp := websocket_data.DataTransferData{
T: constants.WEBSOCKET_UNAUTH,
D: nil,
}
b, _ := json.Marshal(rsp)
//先发一条正常信息
_ = conn.WriteMessage(websocket.TextMessage, b)
//发送关闭信息
_ = conn.WriteMessage(websocket.CloseMessage, nil)
return
}
//生成连接唯一标识
uniqueId := websocketIdGenerator.Get()
ws := wsConnectItem{
conn: conn,
rabbitMq: l.svcCtx.RabbitMq,
uniqueId: uniqueId,
closeChan: make(chan struct{}, 1),
inChan: make(chan []byte, 1000),
outChan: make(chan []byte, 1000),
renderProperty: renderProperty{
renderImageTask: make(map[string]string),
renderImageTaskCtlChan: make(chan renderImageControlChanItem, 100),
},
}
if userInfo.UserId > 0 {
ws.userId = userInfo.UserId
} else {
ws.userId = userInfo.GuestId
}
//保存连接
mapConnPool.Store(uniqueId, ws)
defer ws.close()
//把连接成功消息发回去
time.Sleep(time.Second) //兼容下火狐
b := ws.respondDataFormat(constants.WEBSOCKET_CONNECT_SUCCESS, uniqueId)
_ = 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
}
//不是登录用户也不是游客
if !userInfo.IsUser() && !userInfo.IsGuest() {
return false, nil
}
return true, userInfo
}
return false, nil
}
// 心跳
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")
}
}
// 格式化返回数据
func (w *wsConnectItem) respondDataFormat(msgType string, data interface{}) []byte {
d := websocket_data.DataTransferData{
T: msgType,
D: data,
}
b, _ := json.Marshal(d)
return b
}
// 处理接受到的数据
func (w *wsConnectItem) dealwithReciveData(data []byte) {
var parseInfo websocket_data.DataTransferData
if err := json.Unmarshal(data, &parseInfo); err != nil {
logx.Error("invalid format of websocket message")
w.outChan <- w.respondDataFormat(constants.WEBSOCKET_ERR_DATA_FORMAT, "invalid format of websocket message:"+string(data))
return
}
d, _ := json.Marshal(parseInfo.D)
//分消息类型给到不同逻辑处理,可扩展
switch parseInfo.T {
//图片渲染
case constants.WEBSOCKET_RENDER_IMAGE:
w.renderImage(d)
default:
}
}