lone-services/services/ws/internal/logic/connectLogic.go

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{}
	l.Logger.Error(sess.Status)
	switch sess.Status {
	case pkgws.StatusPending:
		data = nil
	case pkgws.StatusScanned, pkgws.StatusLogined:
		l.Logger.Error("status", sess.Status)
		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:
		l.Logger.Error("use position")
		if len(sess.Data) > 0 && string(sess.Data) != "null" {
			data = sess.Data
			l.Logger.Error("data=", 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)
}