aiChat/internal/controller/chatWs.go

package controller

import (
	"aiChat/internal/model"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/ghttp"
	"github.com/gogf/gf/v2/os/glog"
)

// ChatWs 聊天websocket请求
var ChatWs = func(r *ghttp.Request) {
	session, _ := model.GetSession(r)
	model.SessionCtx.Init(r, &model.Context{
		JSession: session,
	})
	wsChat := model.NewMessage(r.Context())
	// 创建ws链接
	ws, err := r.WebSocket()
	if err != nil {
		glog.Error(wsChat.Ctx, session.AccountId, "建立连接出错", err)
		r.Exit()
	}
	if session.AccountId == 0 {
		_ = ws.WriteJSON(g.Map{
			"error_code": -1,
			"error_msg":  "无用户身份",
		})
		return
	}
	for {
		// 接受客户端信息
		_, msg, err := ws.ReadMessage()
		if err != nil {
			glog.Info(wsChat.Ctx, session.AccountId, "接收消息出错", err)
			_ = ws.Close()
			return
		}
		rv := model.ChatLimit.GetBucket(wsChat.Ctx, session.AccountId).TakeAvailable(1)
		g.Log().Info(wsChat.Ctx, "c----", rv)
		if rv == 0 {
			g.Log().Info(wsChat.Ctx, "请求频繁")
			_ = ws.WriteJSON(g.Map{"error_code": 0, "error_msg": g.Cfg().MustGet(wsChat.Ctx, "limit.exceedMsg").String()})
			return
		}
		// 处理ws
		wsChat.Handle(ws, msg)
	}
}