package service import ( "go_dreamfactory/cmd/v2/lib/common" "go_dreamfactory/cmd/v2/lib" "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 } func NewWsCli(addr string) lib.Handler { return &WsCli{addr: addr} } 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) BuildReq() lib.RawReq { id := time.Now().UnixNano() // TODO: 读流程表配置,确定请求UserMessage b, err := cli.loginReq() if err != nil { panic(err) } rawReq := lib.RawReq{ID: id, Req: b} return rawReq } func (cli *WsCli) Call(req []byte, timeout time.Duration) ([]byte, error) { dialer := &websocket.Dialer{ HandshakeTimeout: timeout, } conn, _, err := dialer.Dial(cli.addr, nil) if err != nil { logrus.Errorf("websocket conn err:%v", err) return nil, err } go func() { timer := time.NewTimer(2 * time.Second) for { timer.Reset(2 * time.Second) <-timer.C if err := conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil { break } } }() // 向连接写数据 conn.WriteMessage(websocket.BinaryMessage, req) // 读数据 var res []byte for { _, data, err := conn.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 }