package logic import ( "context" "errors" "fmt" "net/http" "time" "lone-services/pkg/redis" pkgws "lone-services/pkg/ws" "lone-services/services/ws/internal/svc" "lone-services/services/ws/internal/types" "github.com/gorilla/websocket" jsoniter "github.com/json-iterator/go" "github.com/zeromicro/go-zero/core/logx" ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r *http.Request) bool { return true }, } type ConnectLogic struct { logx.Logger ctx context.Context svcCtx *svc.ServiceContext } func NewConnectLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ConnectLogic { return &ConnectLogic{ Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx, } } func (l *ConnectLogic) Connect(w http.ResponseWriter, r *http.Request, req *types.ConnectReq) error { if req.SessionId == "" { return fmt.Errorf("missing session_id") } sess, err := pkgws.GetSession(l.ctx, req.SessionId) if err != nil { if errors.Is(err, redis.Nil) { return fmt.Errorf("qrcode session expired") } l.Errorf("get session: %v", err) return fmt.Errorf("qrcode session invalid") } conn, err := upgrader.Upgrade(w, r, nil) if err != nil { return err } client := &svc.Client{SessionID: req.SessionId, Conn: conn} l.svcCtx.Hub.Add(client) if err := l.pushSnapshot(client, sess); err != nil { l.Errorf("push snapshot session=%s: %v", req.SessionId, err) l.svcCtx.Hub.Remove(req.SessionId, conn) _ = conn.Close() return nil } go l.readLoop(client) go l.autoClose(client) return nil } func (l *ConnectLogic) pushSnapshot(c *svc.Client, sess *pkgws.Session) error { var data interface{} switch sess.Status { case pkgws.StatusPending: data = nil case pkgws.StatusScanned, pkgws.StatusLogined: if len(sess.Data) > 0 && string(sess.Data) != "null" { var parsed any if err := jsoniter.Unmarshal(sess.Data, &parsed); err != nil { data = jsoniter.RawMessage(sess.Data) } else { data = parsed } } else if sess.Status == pkgws.StatusScanned { data = map[string]int{"ws_type": pkgws.WsTypeScanned} } else { data = nil } case pkgws.StatusUsePosition: if len(sess.Data) > 0 && string(sess.Data) != "null" { data = sess.Data } else { data = nil } default: data = nil } payload, err := jsoniter.Marshal(pkgws.Success(data)) if err != nil { return err } l.svcCtx.Hub.Send(c.SessionID, payload) return nil } func (l *ConnectLogic) readLoop(c *svc.Client) { defer func() { _ = c.Conn.Close() l.svcCtx.Hub.Remove(c.SessionID, c.Conn) }() for { if _, _, err := c.Conn.ReadMessage(); err != nil { return } } } func (l *ConnectLogic) autoClose(c *svc.Client) { timer := time.NewTimer(pkgws.SessionTTL) defer timer.Stop() <-timer.C l.Infof("timeout close session=%s", c.SessionID) _ = c.Conn.Close() l.svcCtx.Hub.Remove(c.SessionID, c.Conn) }