tangmingyou
3 years ago
6 changed files with 392 additions and 7 deletions
@ -0,0 +1,161 @@
|
||||
package gnet_server |
||||
|
||||
import ( |
||||
"bytes" |
||||
"fmt" |
||||
"github.com/gobwas/ws" |
||||
"github.com/gobwas/ws/wsutil" |
||||
"github.com/panjf2000/gnet/v2" |
||||
"github.com/panjf2000/gnet/v2/pkg/logging" |
||||
"io" |
||||
) |
||||
|
||||
type wsCodec struct { |
||||
upgraded bool // 链接是否升级
|
||||
buf bytes.Buffer // 从实际socket中读取到的数据缓存
|
||||
wsMsgBuf wsMessageBuf // ws 消息缓存
|
||||
} |
||||
|
||||
type wsMessageBuf struct { |
||||
firstHeader *ws.Header |
||||
curHeader *ws.Header |
||||
cachedBuf bytes.Buffer |
||||
} |
||||
|
||||
type readWrite struct { |
||||
io.Reader |
||||
io.Writer |
||||
} |
||||
|
||||
func (w *wsCodec) upgrade(c gnet.Conn) (ok bool, action gnet.Action) { |
||||
if w.upgraded { |
||||
ok = true |
||||
return |
||||
} |
||||
buf := &w.buf |
||||
tmpReader := bytes.NewReader(buf.Bytes()) |
||||
oldLen := tmpReader.Len() |
||||
logging.Infof("do Upgrade") |
||||
|
||||
hs, err := ws.Upgrade(readWrite{tmpReader, c}) |
||||
skipN := oldLen - tmpReader.Len() |
||||
if err != nil { |
||||
if err == io.EOF || err == io.ErrUnexpectedEOF { //数据不完整
|
||||
return |
||||
} |
||||
buf.Next(skipN) |
||||
logging.Infof("conn[%v] [err=%v]", c.RemoteAddr().String(), err.Error()) |
||||
action = gnet.Close |
||||
return |
||||
} |
||||
buf.Next(skipN) |
||||
logging.Infof("conn[%v] upgrade websocket protocol! Handshake: %v", c.RemoteAddr().String(), hs) |
||||
if err != nil { |
||||
logging.Infof("conn[%v] [err=%v]", c.RemoteAddr().String(), err.Error()) |
||||
action = gnet.Close |
||||
return |
||||
} |
||||
ok = true |
||||
w.upgraded = true |
||||
return |
||||
} |
||||
func (w *wsCodec) readBufferBytes(c gnet.Conn) gnet.Action { |
||||
size := c.InboundBuffered() |
||||
buf := make([]byte, size, size) |
||||
read, err := c.Read(buf) |
||||
if err != nil { |
||||
logging.Infof("read err! %w", err) |
||||
return gnet.Close |
||||
} |
||||
if read < size { |
||||
logging.Infof("read bytes len err! size: %d read: %d", size, read) |
||||
return gnet.Close |
||||
} |
||||
w.buf.Write(buf) |
||||
return gnet.None |
||||
} |
||||
func (w *wsCodec) Decode(c gnet.Conn) (outs []wsutil.Message, err error) { |
||||
fmt.Println("do Decode") |
||||
messages, err := w.readWsMessages() |
||||
if err != nil { |
||||
logging.Infof("Error reading message! %v", err) |
||||
return nil, err |
||||
} |
||||
if messages == nil || len(messages) <= 0 { //没有读到完整数据 不处理
|
||||
return |
||||
} |
||||
for _, message := range messages { |
||||
if message.OpCode.IsControl() { |
||||
err = wsutil.HandleClientControlMessage(c, message) |
||||
if err != nil { |
||||
return |
||||
} |
||||
continue |
||||
} |
||||
if message.OpCode == ws.OpText || message.OpCode == ws.OpBinary { |
||||
outs = append(outs, message) |
||||
} |
||||
} |
||||
return |
||||
} |
||||
|
||||
func (w *wsCodec) readWsMessages() (messages []wsutil.Message, err error) { |
||||
msgBuf := &w.wsMsgBuf |
||||
in := &w.buf |
||||
for { |
||||
if msgBuf.curHeader == nil { |
||||
if in.Len() < ws.MinHeaderSize { //头长度至少是2
|
||||
return |
||||
} |
||||
var head ws.Header |
||||
if in.Len() >= ws.MaxHeaderSize { |
||||
head, err = ws.ReadHeader(in) |
||||
if err != nil { |
||||
return messages, err |
||||
} |
||||
} else { //有可能不完整，构建新的 reader 读取 head 读取成功才实际对 in 进行读操作
|
||||
tmpReader := bytes.NewReader(in.Bytes()) |
||||
oldLen := tmpReader.Len() |
||||
head, err = ws.ReadHeader(tmpReader) |
||||
skipN := oldLen - tmpReader.Len() |
||||
if err != nil { |
||||
if err == io.EOF || err == io.ErrUnexpectedEOF { //数据不完整
|
||||
return messages, nil |
||||
} |
||||
in.Next(skipN) |
||||
return nil, err |
||||
} |
||||
in.Next(skipN) |
||||
} |
||||
|
||||
msgBuf.curHeader = &head |
||||
err = ws.WriteHeader(&msgBuf.cachedBuf, head) |
||||
if err != nil { |
||||
return nil, err |
||||
} |
||||
} |
||||
dataLen := (int)(msgBuf.curHeader.Length) |
||||
if dataLen > 0 { |
||||
if in.Len() >= dataLen { |
||||
_, err = io.CopyN(&msgBuf.cachedBuf, in, int64(dataLen)) |
||||
if err != nil { |
||||
return |
||||
} |
||||
} else { //数据不完整
|
||||
fmt.Println(in.Len(), dataLen) |
||||
logging.Infof("incomplete data") |
||||
return |
||||
} |
||||
} |
||||
if msgBuf.curHeader.Fin { //当前 header 已经是一个完整消息
|
||||
messages, err = wsutil.ReadClientMessage(&msgBuf.cachedBuf, messages) |
||||
if err != nil { |
||||
return nil, err |
||||
} |
||||
msgBuf.cachedBuf.Reset() |
||||
} else { |
||||
logging.Infof("The data is split into multiple frames") |
||||
} |
||||
msgBuf.curHeader = nil |
||||
} |
||||
} |
||||
@ -0,0 +1,23 @@
|
||||
package gnet_server |
||||
|
||||
import ( |
||||
"github.com/gobwas/ws" |
||||
"github.com/gobwas/ws/wsutil" |
||||
"github.com/panjf2000/gnet/v2" |
||||
) |
||||
|
||||
type gnetConn struct { |
||||
conn gnet.Conn |
||||
uid string |
||||
wsCodec *wsCodec |
||||
} |
||||
|
||||
func newGnetConn(conn gnet.Conn) *gnetConn { |
||||
return &gnetConn{ |
||||
conn: conn, |
||||
} |
||||
} |
||||
|
||||
func (c *gnetConn) Write(msg []byte) error { |
||||
return wsutil.WriteServerMessage(c.conn, ws.OpBinary, msg) |
||||
} |
||||
@ -0,0 +1,184 @@
|
||||
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))) |
||||
} |
||||
Loading…
Reference in new issue