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