sonet/internal/gateway_ws/gnet_server/websocket.go

package gnet_server

import (
	"context"
	"flag"
	"fmt"
	"github.com/gobwas/ws"
	"log"
	"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 GnetWsServer struct {
	gnet.BuiltinEventEngine

	addr      string
	multicore bool
	eng       gnet.Engine
	connected int64
	requests  chan *protocol.Payload
}

func NewGnetWsServer() *GnetWsServer {
	return &GnetWsServer{
		requests: make(chan *protocol.Payload, 1024*16),
	}
}

func (wss *GnetWsServer) Init() {

}

// dispatch 多个 goroutine 进行请求 dispatch
func (wss *GnetWsServer) dispatch(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case payload := <-wss.requests:
			fmt.Println(payload.Header)
		}
	}
}

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
		select {
		case wss.requests <- payload:
		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 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
}

func main() {
	var port int
	var multicore bool

	// Example command: go run main.go --port 8080 --multicore=true
	flag.IntVar(&port, "port", 9080, "server port")
	flag.BoolVar(&multicore, "multicore", true, "multicore")
	flag.Parse()

	wss := &GnetWsServer{addr: fmt.Sprintf("tcp://127.0.0.1:%d", port), multicore: multicore}

	// Start serving!
	log.Println("server exits:", gnet.Run(wss, wss.addr, gnet.WithMulticore(multicore), gnet.WithReusePort(true), gnet.WithTicker(true)))
}