go_dreamfactory/cmd/v2/service/wsCli.go

package service

import (
	"go_dreamfactory/cmd/v2/lib"
	"go_dreamfactory/cmd/v2/lib/common"
	"go_dreamfactory/comm"
	"go_dreamfactory/pb"
	"strings"
	"time"

	"github.com/Pallinder/go-randomdata"
	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
	"google.golang.org/protobuf/proto"
)

type WsCli struct {
	addr    string
	ws      *websocket.Conn
	reqData []byte //请求数据
}

func NewWsCli(addr string, timeout time.Duration) lib.Handler {
	cli := &WsCli{addr: addr}
	cli.connect(timeout)
	return cli
}

func (cli *WsCli) connect(timeout time.Duration) error {
	if cli.ws == nil {
		dialer := &websocket.Dialer{
			HandshakeTimeout: timeout,
		}
		conn, _, err := dialer.Dial(cli.addr, nil)
		if err != nil {
			logrus.Errorf("websocket conn err:%v", err)
			return err
		}
		cli.ws = conn
	}

	// ping
	go func() {
		timer := time.NewTimer(2 * time.Second)
		for {
			timer.Reset(2 * time.Second)
			<-timer.C
			if err := cli.ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				break
			}
		}
	}()
	return nil
}

func (cli *WsCli) loginReq() ([]byte, error) {
	head := &pb.UserMessage{MainType: "user", SubType: "login"}
	sid := "dfz"
	account := randomdata.SillyName()
	head.Sec = common.BuildSecStr(sid, account)
	if comm.ProtoMarshal(&pb.UserLoginReq{
		Sid:     sid,
		Account: account,
	}, head) {
		data, err := proto.Marshal(head)
		if err != nil {
			return nil, err
		}

		return data, nil
	}
	return nil, nil
}

// 检查登录相应
func (cli *WsCli) checkLoginResp(data []byte) bool {
	msg := &pb.UserMessage{}
	if err := proto.Unmarshal(data, msg); err != nil {
		logrus.Error("结果解析失败")
		return false
	}
	if msg.MainType == "user" && msg.SubType == "login" {
		rsp := &pb.UserLoginResp{}
		if !comm.ProtoUnmarshal(msg, rsp) {
			logrus.Error("unmarshal err")
			return false
		}

		if rsp.Data != nil {
			if rsp.Data.Uid != "" {
				logrus.WithField("uid", rsp.Data.Uid).Debug("登录响应")
				return true
			}
		}
		return true
	}
	return false
}

// 检查推送（包括错误推送）
func (cli *WsCli) checkPush(data []byte) bool {
	msg := &pb.UserMessage{}
	if err := proto.Unmarshal(data, msg); err != nil {
		logrus.Error("结果解析失败")
		return false
	}
	methodStr := msg.Data.TypeUrl
	methodName := common.SubStr(methodStr, 20, len(methodStr))

	if strings.HasSuffix(methodName, "Push") {
		if methodName == "NotifyErrorNotifyPush" {
			logrus.WithField("methodName", methodName).Debug("收到错误码")
		} else {
			logrus.WithField("methodName", methodName).Debug("收到推送")
		}
		return true
	}
	return false
}

func (cli *WsCli) SetReq(data []byte) {
	cli.reqData = data
}

func (cli *WsCli) BuildReq() lib.RawReq {
	id := time.Now().UnixNano()
	rawReq := lib.RawReq{ID: id, Req: cli.reqData}
	return rawReq
}

func (cli *WsCli) Call(req []byte) ([]byte, error) {

	// 向连接写数据
	cli.ws.WriteMessage(websocket.BinaryMessage, req)

	// 读数据
	var res []byte
	for {
		_, data, err := cli.ws.ReadMessage()
		if err != nil {
			logrus.Errorf("readMessage err:%v", err)
			break
		}

		if cli.checkLoginResp(data) {
			return data, nil
		} else {
			if !cli.checkPush(data) {
				logrus.Debug("登录失败")
			}
		}
	}

	return res, nil
}

func (cli *WsCli) Check(req lib.RawReq, resp lib.RawResp) *lib.CallResult {
	var result lib.CallResult
	result.Id = resp.ID
	result.Req = req
	result.Resp = resp
	//TODO 解析结果
	msg := &pb.UserMessage{}
	if err := proto.Unmarshal(resp.Resp, msg); err != nil {
		logrus.Error("结果解析失败")
		return &result
	}
	logrus.WithFields(logrus.Fields{"msg": msg}).Debug("检查结果")
	return &result
}