sonet/internal/gateway_ws/gnet_server/websocket.go

package gnet_server

import (
	"context"
	"fmt"
	"github.com/gobwas/ws"
	"google.golang.org/protobuf/proto"
	"regexp"
	"sonet/api/gen/auth"
	"sonet/api/gen/postal"
	"sonet/pkg/grpc/generic"
	"sonet/pkg/protocol"
	"sonet/pkg/protocol/session"
	"sonet/pkg/utils/logger"
	"sync/atomic"
	"time"

	"github.com/panjf2000/gnet/v2"
	"github.com/panjf2000/gnet/v2/pkg/logging"
)

type gnetRequest struct {
	conn    *gnetConn
	payload *protocol.Payload
}

type GnetWsServer struct {
	gnet.BuiltinEventEngine

	addr      string
	multicore bool
	eng       gnet.Engine
	connected int64

	grpcFactory  *generic.GrpcGenericClientFactory
	sessionStore session.Store // <string, *session.NetSubject> 当前连接用户,内存缓存
	requests     chan *gnetRequest
}

func NewGnetWsServer(
	grpcFactory *generic.GrpcGenericClientFactory,
	sessionStore session.Store) *GnetWsServer {
	return &GnetWsServer{
		grpcFactory:  grpcFactory,
		sessionStore: sessionStore,
	}
}

func (wss *GnetWsServer) Init(ctx context.Context) {
	wss.requests = make(chan *gnetRequest, 1024*16)

	// todo 针对不同服务的协程池
	for i := 0; i < 32; i++ {
		go wss.dispatch(ctx)
	}
}

// dispatch 多个 goroutine 进行请求 dispatch
func (wss *GnetWsServer) dispatch(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case request := <-wss.requests:
			func() {
				conn, payload := request.conn, request.payload
				header := payload.Header
				service, method := header.Svc, header.Target

				// grpc generic call
				ctx := context.Background()
				ctx2, cancel2 := context.WithTimeout(ctx, time.Second*3)
				defer cancel2()
				grpcClient, err := wss.grpcFactory.GetClient(ctx2, service)
				if err != nil {
					logger.Error("get grpc generic client error: ", err)
					writeError(conn, header.SeqId, fmt.Sprintf("service %s not avaliable: %s", header.Svc, err.Error()))
					return
				}

				// put grpc request session
				if conn.uid != "" {
					ctx = session.PutSubject(ctx, session.NewRpcSubject(conn.uid))
				}
				ctx3, cancel3 := context.WithTimeout(ctx, time.Second*10)
				defer cancel3()
				resp, err := grpcClient.InvokeUnary(ctx3, header.Target, payload.Body)
				if err != nil {
					writeError(conn, header.SeqId, unwrapRpcError(err.Error()))
					logger.Error("grpc generic call error: ", err)
					return
				}

				// 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(conn, header.SeqId, "server error")
						return
					}
					header.Target = resp.XXX_MessageName()
				}

				// auth verify success
				if service == "Auth" && method == "Verify" {
					var channel *session.Channel
					conn.uid, channel, err = wss.extractAuthVerify(conn, payload.Body)
					if err != nil {
						logger.Error("extractAuthVerify error: ", err)
						writeError(conn, header.SeqId, "server error")
						return
					}
					// channel
					go wss.processWriting(channel)
				}

				header.Type = protocol.TypeResponse
				resMessage, err := protocol.EncodeSo(payload)
				if err != nil {
					logger.Error("grpc generic call error: ", err)
					return
				}
				err = conn.Write(resMessage)
				if err != nil {
					logger.Error("gws write response error: ", err)
					return
				}
			}()
		}
	}
}

func encodeDeliverMessage(msg *postal.Message) (bytes []byte, err error) {
	header := &protocol.Header{
		Magic:         protocol.Magic,
		Type:          protocol.TypeNotice,
		UrlType:       1,
		SerializeType: 1,
		Svc:           msg.Svc,
		Target:        msg.Msg,
	}
	payload := &protocol.Payload{Header: header, Body: msg.Body}
	bytes, err = protocol.EncodeSo(payload)
	if err != nil {
		logger.Error("encode deliver message error: ", err)
	}
	return
}

// todo close
func (c *GnetWsServer) processWriting(channel *session.Channel) {
	for {
		msg := channel.Ready()
		bytes, err := encodeDeliverMessage(msg)
		if err != nil {
			logger.Error("encode dispatch message error:", err)
			continue
		}
		err = channel.Conn.Write(bytes)
		if err != nil {
			logger.Error("write dispatch message error:", err)
		}
	}
}

func (wss *GnetWsServer) OnBoot(eng gnet.Engine) gnet.Action {
	wss.eng = eng
	logging.Infof("echo server with multi-core=%t is listening on %s", wss.multicore, wss.addr)
	return gnet.None
}

func (wss *GnetWsServer) OnOpen(c gnet.Conn) ([]byte, gnet.Action) {
	// wsCodec, uid, gnetConn
	conn := newGnetConn(c)
	conn.wsCodec = new(wsCodec)
	c.SetContext(conn)

	// c.SetContext(new(wsCodec))
	atomic.AddInt64(&wss.connected, 1)
	return nil, gnet.None
}

func (wss *GnetWsServer) OnClose(c gnet.Conn, err error) (action gnet.Action) {
	if err != nil {
		logging.Warnf("error occurred on connection=%s, %v\n", c.RemoteAddr().String(), err)
	}
	atomic.AddInt64(&wss.connected, -1)
	logging.Infof("conn[%v] disconnected", c.RemoteAddr().String())
	return gnet.None
}

func (wss *GnetWsServer) OnTraffic(c gnet.Conn) (action gnet.Action) {
	conn := c.Context().(*gnetConn)
	if conn.wsCodec.readBufferBytes(c) == gnet.Close {
		return gnet.Close
	}
	ok, action := conn.wsCodec.upgrade(c)
	if !ok {
		return
	}

	if conn.wsCodec.buf.Len() <= 0 {
		return gnet.None
	}
	messages, err := conn.wsCodec.Decode(c)
	if err != nil {
		return gnet.Close
	}
	if messages == nil {
		return
	}
	authorized := conn.uid != ""

	for _, message := range messages {
		if message.OpCode != ws.OpBinary {
			logger.Info("receive message type: ", message.OpCode)
			continue
		}

		// decode payload
		payload, err := protocol.DecodeSo(message.Payload)
		if err != nil {
			logger.Errorf("decode message error: len=%d", len(message.Payload), err)
			return gnet.Close
		}

		header := payload.Header
		service, method := header.Svc, header.Target
		// authorization
		if !authorized {
			if !(service == "Auth" && method == "Verify") {
				writeError(conn, header.SeqId, session.UnauthorizedRequestError.Error())
				return gnet.Close
			}
		}

		// payload -> processing channel
		request := &gnetRequest{conn: conn, payload: payload}
		select {
		case wss.requests <- request:
		default:
			writeError(conn, header.SeqId, "server busy")
		}
		//msgLen := len(message.Payload)
		//if msgLen > 128 {
		//	logging.Infof("conn[%v] receive [op=%v] [msg=%v..., len=%d]", c.RemoteAddr().String(), message.OpCode, string(message.Payload[:128]), len(message.Payload))
		//} else {
		//	logging.Infof("conn[%v] receive [op=%v] [msg=%v, len=%d]", c.RemoteAddr().String(), message.OpCode, string(message.Payload), len(message.Payload))
		//}
		//// This is the echo server
		//err = wsutil.WriteServerMessage(c, message.OpCode, message.Payload)
		//if err != nil {
		//	logging.Infof("conn[%v] [err=%v]", c.RemoteAddr().String(), err.Error())
		//	return gnet.Close
		//}
	}
	return gnet.None
}

func (wss *GnetWsServer) OnTick() (delay time.Duration, action gnet.Action) {
	logging.Infof("[connected-count=%v]", atomic.LoadInt64(&wss.connected))
	return 10 * time.Second, gnet.None
}

func (c *GnetWsServer) extractAuthVerify(conn *gnetConn, resp []byte) (sessionUid string, channel *session.Channel, err error) {
	authSubject := &auth.Subject{}
	err = proto.Unmarshal(resp, authSubject)
	if err != nil {
		logger.Error("grpc generic call error: ", err)
		return
	}
	// 设置连接身份信息
	sessionUid = authSubject.Uid

	// 关闭旧的链接
	oldChannel, ok := c.sessionStore.Load(sessionUid)
	if ok {
		if err = oldChannel.Conn.(*gnetConn).conn.Close(); err != nil {
			logger.Error("close gnet conn error: ", err)
		}
		logger.Infof("close subject old conn: %s\n", sessionUid)
	}
	// 存储session到内存
	channel = session.NewChannel(sessionUid, conn)
	c.sessionStore.Store(sessionUid, channel)
	logger.Infof("subject online: %s\n", sessionUid)
	return
}

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
}

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
}