package server import ( "context" "fmt" "github.com/gorilla/websocket" "google.golang.org/protobuf/proto" "regexp" "runtime/debug" "sonet/api/gen/auth" "sonet/pkg/grpc/generic" "sonet/pkg/plugins/cache" "sonet/pkg/protocol" "sonet/pkg/protocol/session" "sonet/pkg/utils/logger" "time" ) type ConnHandler struct { postalServerAddress string grpcFactory *generic.GrpcGenericClientFactory sessionStore session.Store // 当前连接用户,内存缓存 subjectStore cache.MultiLevelCache } func NewConnHandler(postalServerAddress string, grpcFactory *generic.GrpcGenericClientFactory, sessionStore session.Store, subjectStore cache.MultiLevelCache) *ConnHandler { return &ConnHandler{ grpcFactory: grpcFactory, postalServerAddress: postalServerAddress, sessionStore: sessionStore, subjectStore: subjectStore, } } func (c *ConnHandler) handleConn(conn *websocket.Conn) { wsConn := newWsConn(conn) var sessionUid string defer func() { // 捕获其他错误 if r := recover(); r != nil { logger.Error("NetClient recover error: ", r) // 输出堆栈信息 logger.Error("NetClient recover error stack: ", string(debug.Stack())) } }() defer func() { if err := conn.Close(); err != nil { logger.Error("conn close: ", err) } // TODO 连接关闭,mq发送关闭事件 // 删除连接 if sessionUid != "" { c.sessionStore.Delete(sessionUid) err := c.subjectStore.Del(context.Background(), sessionUid) if err != nil { logger.Error("del cluster subject store uid error: ", sessionUid) } logger.Info("subject offline: ", sessionUid) } }() for { if err := conn.SetReadDeadline(time.Time{}); err != nil { return } t, message, err := conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) { logger.Error("unexpected close error: ", err) } else { // 读失败 logger.Errorf("NetClient ReadMessage error: %T, %v", err, err) } return } if t != websocket.BinaryMessage { continue } payload, err := protocol.DecodeSo(message) if err != nil { logger.Errorf("decode message error: len=%d", len(message), err) return } header := payload.Header service, method := header.Svc, header.Target if sessionUid == "" && !(service == "Auth" && method == "Verify") { writeError(wsConn, header.SeqId, session.UnauthorizedRequestError.Error()) continue } // grpc generic call ctx := context.Background() grpcClient, err := c.grpcFactory.GetClient(ctx, service) if err != nil { logger.Error("get grpc generic client error: ", err) writeError(wsConn, header.SeqId, fmt.Sprintf("service %s not avaliable: %s", header.Svc, err.Error())) continue } // put session if sessionUid != "" { ctx = session.PutSubject(ctx, session.NewRpcSubject(sessionUid)) } resp, err := grpcClient.InvokeUnary(ctx, header.Target, payload.Body) if err != nil { writeError(wsConn, header.SeqId, unwrapRpcError(err.Error())) logger.Error("grpc generic call error: ", err) continue } // write response if resp == nil { // proto.Empty header.Target = "" } else { payload.Body, err = resp.Marshal() if err != nil { logger.Error("generic call response marshal error: ", err) writeError(wsConn, header.SeqId, "server error") continue } header.Target = resp.XXX_MessageName() } // auth verify success if service == "Auth" && method == "Verify" { sessionUid, err = c.extractAuthVerify(wsConn, payload.Body) if err != nil { logger.Error("extractAuthVerify error: ", err) writeError(wsConn, header.SeqId, "server error") continue } } header.Type = protocol.TypeResponse resMessage, err := protocol.EncodeSo(payload) if err != nil { logger.Error("grpc generic call error: ", err) continue } err = wsConn.Write(resMessage) if err != nil { logger.Errorf("write 2 %s response error: \n", sessionUid, err) } } } func (c *ConnHandler) extractAuthVerify(conn *wsConn, resp []byte) (sessionUid string, err error) { authSubject := &auth.Subject{} err = proto.Unmarshal(resp, authSubject) if err != nil { logger.Error("grpc generic call error: ", err) return } // 设置连接身份信息 sessionUid = authSubject.Uid gateSubject := &session.GateSubject{ Uid: authSubject.Uid, Online: 1, Time: time.Now().UnixMilli(), Gate: c.postalServerAddress, } // 存储session err = c.subjectStore.Set(context.Background(), sessionUid, gateSubject) // store to redis cluster cache if err != nil { return } // 关闭旧的链接 oldConn, ok := c.sessionStore.Load(sessionUid) if ok { // TODO nats offline / force load subject target cluster call offline if err = oldConn.(*wsConn).conn.Close(); err != nil { logger.Error("old conn close error: ", err) } logger.Infof("close subject old conn: %s\n", sessionUid) } // 存储session到内存 c.sessionStore.Store(sessionUid, conn) logger.Infof("subject online: %s\n", sessionUid) return } var reg = regexp.MustCompile("^rpc error:.*?desc ?= ?(.*)$") // unwrapRpcError 提取 grpc err: fmt.Sprintf("rpc error: code = %s desc = %s", s.Code(), s.Message()) func unwrapRpcError(err string) string { finds := reg.FindStringSubmatch(err) if len(finds) > 1 { return finds[1] } return err } func writeError(conn session.NetConn, seqId int32, errMsg string) { header := &protocol.Header{} payload := &protocol.Payload{Header: header} header.Magic = protocol.Magic header.Type = protocol.TypeError header.Status = 50 header.SeqId = seqId payload.Body = []byte(errMsg) message, err := protocol.EncodeSo(payload) if err != nil { logger.Error("NetClient writeError encode error: ", err) return } err = conn.Write(message) if err != nil { logger.Error("net conn write error msg fail: ", err) return } return }