fusenapi/server/websocket/internal/logic/datatransferlogic.go

package logic

//websocket连接
import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"fusenapi/constants"
	"fusenapi/server/websocket/internal/types"
	"fusenapi/utils/auth"
	"fusenapi/utils/basic"
	"fusenapi/utils/encryption_decryption"
	"fusenapi/utils/websocket_data"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"

	"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 (
	//临时对象缓存池
	buffPool = sync.Pool{
		New: func() interface{} {
			return bytes.Buffer{}
		},
	}
	//升级websocket
	upgrader = websocket.Upgrader{
		//最大可读取大小 1M
		ReadBufferSize: 1024,
		//最大可写大小 1M
		WriteBufferSize: 1024,
		//握手超时时间15s
		HandshakeTimeout: time.Second * 15,
		//允许跨域
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
		//写的缓冲队列
		WriteBufferPool: &buffPool,
		//是否支持压缩
		EnableCompression: true,
	}
	//websocket连接存储
	mapConnPool = sync.Map{}
	//每个websocket连接入口缓冲队列长度默认值
	websocketInChanLen = 2000
	//每个websocket连接出口缓冲队列长度默认值
	websocketOutChanLen = 2000
)

// 每个连接的连接基本属性
type wsConnectItem struct {
	conn                 *websocket.Conn      //websocket的连接(基本属性)
	connExpireTime       int64                //websocket过期时间（跟随连接时候的token过期时间）
	userAgent            string               //用户代理头信息(基本属性，用于重连标识验证因素之一)
	logic                *DataTransferLogic   //logic(基本属性，用于获取上下文，配置或者操作数据库)
	closeChan            chan struct{}        //ws连接关闭chan(基本属性)
	isClose              bool                 //是否已经关闭(基本属性)
	uniqueId             string               //ws连接唯一标识(基本属性)
	inChan               chan []byte          //接受消息缓冲队列(基本属性)
	outChan              chan []byte          //要发送回客户端的消息缓冲队列(基本属性)
	mutex                sync.Mutex           //互斥锁(基本属性)
	debug                *auth.Debug          //是否开启debug
	userId               int64                //用户id(基本属性)
	guestId              int64                //游客id(基本属性)
	extendRenderProperty extendRenderProperty //扩展云渲染属性(扩展属性)
}

// 请求建立连接，升级websocket协议
func (l *DataTransferLogic) DataTransfer(req *types.DataTransferReq, w http.ResponseWriter, r *http.Request) {
	//把子协议携带的token设置到标准token头信息中
	tokens := r.Header.Get("Sec-Websocket-Protocol")
	oldWid := req.Wid
	oldWid = strings.Trim(oldWid, " ")
	//有token是正常用户，无则是白板用户，也可以连接
	if tokens != "" {
		token := ""
		debugToken := ""
		tokenSlice := strings.Split(tokens, ",")
		switch len(tokenSlice) {
		case 1:
			token = strings.Trim(tokenSlice[0], " ")
		case 2:
			token = strings.Trim(tokenSlice[0], " ")
			debugToken = strings.Trim(tokenSlice[1], " ")
		default:
			logx.Error("invalid ws token:", tokens)
			return
		}
		if token != "empty_token" && token != "" {
			r.Header.Set("Authorization", "Bearer "+token)
		}
		if debugToken != "empty_debug_token" && debugToken != "" {
			r.Header.Set("Debug-Token", debugToken)
		}
		//设置Sec-Websocket-Protocol
		upgrader.Subprotocols = strings.Split(tokens, ",")
	}
	//判断下是否火狐浏览器(获取浏览器第一条消息返回有收不到的bug需要延迟1秒)
	userAgent := r.Header.Get("User-Agent")
	//是否火狐浏览器
	isFirefoxBrowser := false
	if strings.Contains(userAgent, "Firefox") {
		isFirefoxBrowser = true
	}
	//升级websocket
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		logx.Error("http upgrade websocket err:", err)
		return
	}
	//支持写压缩
	conn.EnableWriteCompression(true)
	//鉴权不成功后断开
	var (
		userInfo *auth.UserInfo
		isAuth   bool
	)
	isAuth, userInfo = l.checkAuth(r)
	if !isAuth {
		//未授权响应消息
		l.unAuthResponse(conn, isFirefoxBrowser, "unAuth")
		conn.Close()
		return
	}
	//设置连接
	ws, err := l.setConnPool(conn, userInfo, isFirefoxBrowser, userAgent, oldWid)
	if err != nil {
		conn.Close()
		return
	}
	//循环读客户端信息
	go ws.reciveBrowserMessage()
	//消费出口数据并发送浏览器端
	go ws.consumeOutChanData()
	//消费入口数据
	go ws.consumeInChanData()
	//消费渲染缓冲队列
	go ws.consumeRenderImageData()
	//心跳
	ws.heartbeat()
}

// 设置连接
func (l *DataTransferLogic) setConnPool(conn *websocket.Conn, userInfo *auth.UserInfo, isFirefoxBrowser bool, userAgent, oldWid string) (wsConnectItem, error) {
	//生成连接唯一标识（失败重试10次）
	uniqueId, err := l.getUniqueId(userInfo, userAgent, 10)
	if err != nil {
		//发送获取唯一标识失败的消息
		if isFirefoxBrowser {
			time.Sleep(time.Second * 1) //兼容下火狐(直接发回去收不到第一条消息：有待研究)
		}
		l.sendGetUniqueIdErrResponse(conn, userInfo.Debug)
		return wsConnectItem{}, err
	}
	//传入绑定的wid判断是否可重用
	if oldWid != "" {
		for i := 0; i < 1; i++ {
			//解析传入的wid是不是属于自己的用户的
			decryptionWid, err := encryption_decryption.CBCDecrypt(oldWid)
			if err != nil {
				logx.Error("解密wid失败:", err)
				break
			}
			lendecryptionWid := len(decryptionWid)
			//合成client后缀,不是同个后缀的不能复用
			userPart := getUserJoinPart(userInfo.UserId, userInfo.GuestId, userAgent)
			lenUserPart := len(userPart)
			//长度太短
			if lendecryptionWid <= lenUserPart {
				logx.Error("复用的连接标识太短，不符合重用条件")
				break
			}
			//尾部不同不能复用
			if decryptionWid[lendecryptionWid-lenUserPart:] != userPart {
				logx.Error("尾部用户信息不同，不符合重用条件")
				break
			}
			//存在是不能给他申请重新绑定
			if _, ok := mapConnPool.Load(oldWid); ok {
				logx.Error("复用的连接标识已被其他客户端使用，不符合重用条件,用户id:", userInfo.UserId, "  guest_id:", userInfo.GuestId)
				break
			}
			logx.Info("====复用旧的ws连接成功====")
			uniqueId = oldWid
		}
	}
	//默认过期时间
	connExpireTime := time.Now().UTC().Add(time.Second * time.Duration(l.svcCtx.Config.Auth.AccessExpire)).Unix()
	if userInfo.Exp > 0 {
		connExpireTime = userInfo.Exp
	}
	renderCtx, renderCtxCancelFunc := context.WithCancel(l.ctx)
	ws := wsConnectItem{
		conn:           conn,
		connExpireTime: connExpireTime,
		userAgent:      userAgent,
		logic:          l,
		closeChan:      make(chan struct{}, 1),
		isClose:        false,
		uniqueId:       uniqueId,
		inChan:         make(chan []byte, websocketInChanLen),
		outChan:        make(chan []byte, websocketOutChanLen),
		mutex:          sync.Mutex{},
		userId:         userInfo.UserId,
		guestId:        userInfo.GuestId,
		extendRenderProperty: extendRenderProperty{
			renderChan:          make(chan websocket_data.RenderImageReqMsg, renderChanLen),
			renderCtx:           renderCtx,
			renderCtxCancelFunc: renderCtxCancelFunc,
		},
		debug: userInfo.Debug,
	}
	//********强制开启debug 后面删掉
	e := int64(1700359131)
	ws.debug = &auth.Debug{
		Exp:              &e,
		IsCache:          1,
		IsAllTemplateTag: 0,
	}
	//保存连接
	mapConnPool.Store(uniqueId, ws)
	//非白板用户，需要为这个用户建立map索引便于通过用户查询
	createUserConnPoolElement(userInfo.UserId, userInfo.GuestId, uniqueId)
	if isFirefoxBrowser {
		time.Sleep(time.Second * 1) //兼容下火狐(直接发回去收不到第一条消息：有待研究)
	}
	ws.sendToOutChan(ws.respondDataFormat(constants.WEBSOCKET_CONNECT_SUCCESS, websocket_data.ConnectSuccessMsg{Wid: uniqueId}))
	//发送累加统计连接数
	increaseWebsocketConnectCount()
	return ws, nil
}

// 获取唯一id
func (l *DataTransferLogic) getUniqueId(userInfo *auth.UserInfo, userAgent string, retryTimes int) (uniqueId string, err error) {
	if retryTimes < 0 {
		return "", errors.New("failed to get unique id")
	}
	//后面拼接上用户id
	uniqueId = hex.EncodeToString([]byte(uuid.New().String())) + getUserJoinPart(userInfo.UserId, userInfo.GuestId, userAgent)
	//存在则从新获取
	if _, ok := mapConnPool.Load(uniqueId); ok {
		uniqueId, err = l.getUniqueId(userInfo, userAgent, retryTimes-1)
		if err != nil {
			return "", err
		}
	}
	//加密
	uniqueId, err = encryption_decryption.CBCEncrypt(uniqueId)
	if err != nil {
		return "", err
	}
	return uniqueId, nil
}

// 鉴权
func (l *DataTransferLogic) checkAuth(r *http.Request) (isAuth bool, userInfo *auth.UserInfo) {
	// 解析JWT token,并对空用户进行判断
	userInfo, err := basic.ParseJwtToken(r, l.svcCtx)
	if err != nil {
		logx.Error("未授权：", err.Error())
		return false, nil
	}
	if userInfo.UserId > 0 {
		userInfo.GuestId = 0
	}
	//白板用户
	return true, userInfo
}

// 心跳检测
func (w *wsConnectItem) heartbeat() {
	tick := time.Tick(time.Second * 5)
	for {
		select {
		case <-w.closeChan:
			return
		case <-tick:
			//看看token是否过期了
			if w.connExpireTime > 0 && w.connExpireTime < time.Now().UTC().Unix() {
				logx.Info("token过期，关闭连接:", w.uniqueId)
				w.close()
				return
			}
			//查看debug的时间是否过期
			if w.debug != nil && w.debug.Exp != nil && *w.debug.Exp < time.Now().UTC().Unix() {
				w.debug = nil
			}
			if err := w.conn.WriteMessage(websocket.PongMessage, nil); err != nil {
				logx.Error("发送心跳信息异常，关闭连接:", w.uniqueId, err)
				w.close()
				return
			}
		}
	}
}

// 关闭websocket连接
func (w *wsConnectItem) close() {
	w.mutex.Lock()
	defer w.mutex.Unlock()
	logx.Info("###websocket:", w.uniqueId, "  uid:", w.userId, "  gid:", w.guestId, " is closing....")
	//发送关闭信息
	_ = w.conn.WriteMessage(websocket.CloseMessage, nil)
	w.conn.Close()
	mapConnPool.Delete(w.uniqueId)
	if !w.isClose {
		w.isClose = true
		close(w.closeChan)
		//删除用户级索引
		deleteUserConnPoolElement(w.userId, w.guestId, w.uniqueId)
		//减少连接数统计
		decreaseWebsocketConnectCount()
	}
	logx.Info("###websocket:", w.uniqueId, "  uid:", w.userId, "  gid:", w.guestId, " is closed")
}

// 读取出口缓冲队列数据输出返回给浏览器端
func (w *wsConnectItem) consumeOutChanData() {
	defer func() {
		if err := recover(); err != nil {
			logx.Error("consumeOutChanData panic:", err)
		}
	}()
	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
			}
		}
	}
}

// 消费websocket入口数据池中的数据
func (w *wsConnectItem) consumeInChanData() {
	defer func() {
		if err := recover(); err != nil {
			logx.Error("consumeInChanData:", err)
		}
	}()
	for {
		select {
		case <-w.closeChan:
			return
		case data := <-w.inChan:
			//对不同消息类型分发处理
			w.allocationProcessing(data)
		}
	}
}

// 接受浏览器端发来的消息并写入入口缓冲队列
func (w *wsConnectItem) reciveBrowserMessage() {
	defer func() {
		if err := recover(); err != nil {
			logx.Error("acceptBrowserMessage panic:", err)
		}
	}()
	for {
		select {
		case <-w.closeChan: //如果关闭了
			return
		default: //收取消息
			msgType, data, err := w.conn.ReadMessage()
			if err != nil {
				logx.Error("接受信息错误:", err)
				//关闭连接
				w.close()
				return
			}
			switch msgType {
			case websocket.PingMessage, websocket.PongMessage: //心跳消息(过滤不处理)
				continue
			case websocket.BinaryMessage, websocket.TextMessage: //主要消息
				w.sendToInChan(data)
			case websocket.CloseMessage: //客户端主动关闭消息
				w.close()
			}
		}
	}
}

// 把要传递给客户端的数据放入出口缓冲队列
func (w *wsConnectItem) sendToOutChan(data []byte) {
	select {
	case <-w.closeChan:
		return
	case w.outChan <- data:
		return
	}
}

// 发送接受到的消息到入口缓冲队列中
func (w *wsConnectItem) sendToInChan(data []byte) {
	select {
	case <-w.closeChan: //关闭了
		return
	case w.inChan <- data:
		return
	}
}

// 格式化为websocket标准返回格式
func (w *wsConnectItem) respondDataFormat(msgType constants.Websocket, data interface{}) []byte {
	d := websocket_data.DataTransferData{
		T:     msgType,
		D:     data,
		Debug: w.debug != nil,
	}
	b, _ := json.Marshal(d)
	return b
}

// 获取用户拼接部分(复用标识用到)
func getUserJoinPart(userId, guestId int64, userAgent string) string {
	if userId > 0 {
		guestId = 0
	}
	return fmt.Sprintf("|_%d_%d_|_%s_|", userId, guestId, userAgent)
}

// 处理入口缓冲队列中不同类型的数据（分发处理）
func (w *wsConnectItem) allocationProcessing(data []byte) {
	var parseInfo websocket_data.DataTransferData
	if err := json.Unmarshal(data, &parseInfo); err != nil {
		w.incomeDataFormatErrResponse("invalid format of income message:" + string(data))
		logx.Error("invalid format of websocket message:", err)
		return
	}
	d, _ := json.Marshal(parseInfo.D)
	//获取工厂实例
	processor := w.newAllocationProcessor(parseInfo.T)
	if processor == nil {
		//logx.Error("未知消息类型:", string(data))
		return
	}
	//执行工厂方法
	processor.allocationMessage(w, d)
}