feat 初次提交
This commit is contained in:
48
net/socket/server.go
Normal file
48
net/socket/server.go
Normal file
@@ -0,0 +1,48 @@
|
||||
package socket
|
||||
|
||||
import (
|
||||
"time"
|
||||
)
|
||||
|
||||
type Action int
|
||||
|
||||
const (
|
||||
// None indicates that no action should occur following an event.
|
||||
None Action = iota
|
||||
// Close closes the connection.
|
||||
Close
|
||||
// Shutdown shutdowns the engine.
|
||||
Shutdown
|
||||
)
|
||||
|
||||
type OpCode byte
|
||||
|
||||
const (
|
||||
OpContinuation OpCode = 0x0
|
||||
OpText OpCode = 0x1
|
||||
OpBinary OpCode = 0x2
|
||||
OpClose OpCode = 0x8
|
||||
OpPing OpCode = 0x9
|
||||
OpPong OpCode = 0xa
|
||||
)
|
||||
|
||||
// ISocketServer 由应用层实现
|
||||
type ISocketServer interface {
|
||||
OnOpen(ISocketConn) ([]byte, Action) // 开启连接
|
||||
OnHandShake(ISocketConn, []byte, func(ISocketConn, []byte)) Action // 开始握手
|
||||
OnMessage(ISocketConn, []byte) Action // 收到消息
|
||||
OnPong(ISocketConn)
|
||||
OnClose(ISocketConn, error) Action // 关闭连接
|
||||
OnTick() (time.Duration, Action)
|
||||
}
|
||||
|
||||
// ISocketConn 由网络层实现
|
||||
type ISocketConn interface {
|
||||
GetParam(key string) interface{}
|
||||
SetParam(key string, values interface{})
|
||||
RemoteAddr() string
|
||||
Ping() error // 需要隔一段时间调用一下，建议20s
|
||||
Write(data []byte) error
|
||||
Close() error
|
||||
IsClose() bool
|
||||
}
|
||||
187
net/socket/websocket/websocket.go
Normal file
187
net/socket/websocket/websocket.go
Normal file
@@ -0,0 +1,187 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"common/net/socket"
|
||||
"context"
|
||||
"github.com/panjf2000/gnet/v2"
|
||||
"github.com/panjf2000/gnet/v2/pkg/logging"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// WSServer 实现GNet库接口
|
||||
type WSServer struct {
|
||||
eng gnet.Engine
|
||||
i socket.ISocketServer
|
||||
logger logging.Logger // 日志
|
||||
upgradeTimeout time.Duration // 升级超时时间
|
||||
unUpgradeConn sync.Map
|
||||
}
|
||||
|
||||
func NewWSServer(i socket.ISocketServer, logger logging.Logger, upgradeTimeout time.Duration) *WSServer {
|
||||
if i == nil {
|
||||
return nil
|
||||
}
|
||||
return &WSServer{
|
||||
i: i,
|
||||
logger: logger,
|
||||
upgradeTimeout: upgradeTimeout,
|
||||
unUpgradeConn: sync.Map{},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *WSServer) Run(addr string, multicore bool, numEventLoop int, tcpNoDelay gnet.TCPSocketOpt, readBufferCap, writeBufferCap int, lockOSThread bool, reusePort, ticker bool, logger logging.Logger) error {
|
||||
return gnet.Run(
|
||||
s,
|
||||
addr,
|
||||
gnet.WithMulticore(multicore),
|
||||
gnet.WithNumEventLoop(numEventLoop),
|
||||
gnet.WithTCPNoDelay(tcpNoDelay),
|
||||
gnet.WithReadBufferCap(readBufferCap),
|
||||
gnet.WithWriteBufferCap(writeBufferCap),
|
||||
gnet.WithLockOSThread(lockOSThread),
|
||||
gnet.WithReusePort(reusePort),
|
||||
gnet.WithTicker(ticker),
|
||||
gnet.WithLogger(logger),
|
||||
)
|
||||
}
|
||||
|
||||
func (s *WSServer) Stop() error {
|
||||
return s.eng.Stop(context.Background())
|
||||
}
|
||||
|
||||
func (s *WSServer) OnBoot(eng gnet.Engine) gnet.Action {
|
||||
s.eng = eng
|
||||
return gnet.None
|
||||
}
|
||||
|
||||
func (s *WSServer) OnShutdown(_ gnet.Engine) {
|
||||
}
|
||||
|
||||
func (s *WSServer) OnOpen(c gnet.Conn) ([]byte, gnet.Action) {
|
||||
ws := &WSConn{
|
||||
Conn: c,
|
||||
isUpgrade: false,
|
||||
openTime: time.Now().Unix(),
|
||||
buf: bytes.Buffer{},
|
||||
logger: s.logger,
|
||||
wsMessageBuf: wsMessageBuf{
|
||||
curHeader: nil,
|
||||
cachedBuf: bytes.Buffer{},
|
||||
},
|
||||
param: make(map[string]interface{}),
|
||||
remoteAddr: c.RemoteAddr().String(),
|
||||
}
|
||||
c.SetContext(ws)
|
||||
s.unUpgradeConn.Store(c.RemoteAddr().String(), ws)
|
||||
d, a := s.i.OnOpen(ws)
|
||||
return d, (gnet.Action)(a)
|
||||
}
|
||||
|
||||
// OnClose fires when a connection has been closed.
|
||||
// The parameter err is the last known connection error.
|
||||
func (s *WSServer) OnClose(c gnet.Conn, err error) (action gnet.Action) {
|
||||
s.unUpgradeConn.Delete(c.RemoteAddr().String())
|
||||
ws, ok := c.Context().(*WSConn)
|
||||
if ok {
|
||||
ws.isClose = true
|
||||
ws.logger.Warnf("connection close, err: %v", err)
|
||||
return gnet.Action(s.i.OnClose(ws, err))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// OnTraffic fires when a local socket receives data from the peer.
|
||||
func (s *WSServer) OnTraffic(c gnet.Conn) gnet.Action {
|
||||
ws := c.Context().(*WSConn)
|
||||
if action := ws.readBytesBuf(c); action != gnet.None {
|
||||
return action
|
||||
}
|
||||
|
||||
if !ws.isUpgrade {
|
||||
data, action := ws.upgrade()
|
||||
if len(data) > 0 {
|
||||
s.unUpgradeConn.Delete(c.RemoteAddr().String())
|
||||
action = gnet.Action(s.i.OnHandShake(ws, data, s.OnHandShakeFinish))
|
||||
}
|
||||
return action
|
||||
}
|
||||
|
||||
if msg, err := ws.readWsMessages(); err != nil {
|
||||
ws.logger.Errorf("read ws messages err: %v", err)
|
||||
return gnet.Close
|
||||
} else if msg != nil {
|
||||
for _, m := range msg {
|
||||
if socket.OpCode(m.OpCode) == socket.OpPong {
|
||||
s.i.OnPong(ws)
|
||||
continue
|
||||
}
|
||||
if socket.OpCode(m.OpCode) == socket.OpClose {
|
||||
return gnet.Close
|
||||
}
|
||||
if socket.OpCode(m.OpCode) == socket.OpPing {
|
||||
continue
|
||||
}
|
||||
a := s.i.OnMessage(ws, m.Payload)
|
||||
if gnet.Action(a) != gnet.None {
|
||||
return gnet.Action(a)
|
||||
}
|
||||
}
|
||||
}
|
||||
return gnet.None
|
||||
}
|
||||
|
||||
// OnTick fires immediately after the engine starts and will fire again
|
||||
// following the duration specified by the delay return value.
|
||||
func (s *WSServer) OnTick() (delay time.Duration, action gnet.Action) {
|
||||
now := time.Now().Unix()
|
||||
|
||||
delConn := make([]string, 0)
|
||||
s.unUpgradeConn.Range(func(key, value interface{}) bool {
|
||||
k, ok := key.(string)
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
|
||||
v, ok := value.(*WSConn)
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
if now-v.openTime >= int64(s.upgradeTimeout.Seconds()) {
|
||||
delConn = append(delConn, k)
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
for _, k := range delConn {
|
||||
wsConn, _ := s.unUpgradeConn.LoadAndDelete(k)
|
||||
if wsConn == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
v, ok := wsConn.(*WSConn)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if err := v.Close(); err != nil {
|
||||
v.logger.Errorf("upgrade ws time out close socket error: %v", err)
|
||||
}
|
||||
}
|
||||
d, a := s.i.OnTick()
|
||||
delay = d
|
||||
action = gnet.Action(a)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// OnHandShakeFinish 握手完成
|
||||
func (s *WSServer) OnHandShakeFinish(conn socket.ISocketConn, hsResp []byte) {
|
||||
ws := conn.(*WSConn)
|
||||
if err := ws.Conn.AsyncWrite(hsResp, nil); err != nil {
|
||||
ws.logger.Errorf("OnHandShakeFinish err: %v", err)
|
||||
if err = ws.Close(); err != nil {
|
||||
ws.logger.Errorf("OnHandShakeFinish Close error: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
202
net/socket/websocket/wsconn.go
Normal file
202
net/socket/websocket/wsconn.go
Normal file
@@ -0,0 +1,202 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"errors"
|
||||
"github.com/gobwas/ws"
|
||||
"github.com/gobwas/ws/wsutil"
|
||||
"github.com/panjf2000/gnet/v2"
|
||||
"github.com/panjf2000/gnet/v2/pkg/logging"
|
||||
"io"
|
||||
"net/url"
|
||||
)
|
||||
|
||||
// WSConn 实现ISocketConn接口
|
||||
type WSConn struct {
|
||||
gnet.Conn
|
||||
buf bytes.Buffer
|
||||
logger logging.Logger
|
||||
isUpgrade bool // 是否已经升级
|
||||
upgradeResp []byte // 升级响应
|
||||
isClose bool
|
||||
param map[string]interface{}
|
||||
openTime int64 // 开启连接的时间
|
||||
remoteAddr string // 远程ID地址
|
||||
wsMessageBuf
|
||||
}
|
||||
|
||||
type wsMessageBuf struct {
|
||||
curHeader *ws.Header
|
||||
cachedBuf bytes.Buffer
|
||||
}
|
||||
|
||||
type readWrite struct {
|
||||
io.Reader
|
||||
io.Writer
|
||||
}
|
||||
|
||||
func (w *WSConn) readBytesBuf(c gnet.Conn) gnet.Action {
|
||||
size := c.InboundBuffered()
|
||||
if size <= 0 {
|
||||
return gnet.None
|
||||
}
|
||||
buf := make([]byte, size)
|
||||
|
||||
read, err := c.Read(buf)
|
||||
if err != nil {
|
||||
if w.logger != nil {
|
||||
w.logger.Errorf("ws read bytes buf error", err)
|
||||
}
|
||||
return gnet.Close
|
||||
}
|
||||
|
||||
if read < size {
|
||||
if w.logger != nil {
|
||||
w.logger.Errorf("read bytes len err! size: %d read: %d", size, read)
|
||||
}
|
||||
return gnet.Close
|
||||
}
|
||||
w.buf.Write(buf)
|
||||
return gnet.None
|
||||
}
|
||||
|
||||
func (w *WSConn) upgrade() (data []byte, action gnet.Action) {
|
||||
buf := &w.buf
|
||||
tmpReader := bytes.NewReader(buf.Bytes())
|
||||
oldLen := tmpReader.Len()
|
||||
result := &bytes.Buffer{}
|
||||
tempWriter := bufio.NewWriter(result)
|
||||
|
||||
var err error = nil
|
||||
up := ws.Upgrader{
|
||||
OnRequest: w.OnRequest,
|
||||
}
|
||||
_, err = up.Upgrade(readWrite{tmpReader, tempWriter})
|
||||
skipN := oldLen - tmpReader.Len()
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { //数据不完整
|
||||
return
|
||||
}
|
||||
buf.Next(skipN)
|
||||
if w.logger != nil {
|
||||
w.logger.Errorf("ws upgrade err: %v", err.Error())
|
||||
}
|
||||
|
||||
action = gnet.Close
|
||||
return
|
||||
}
|
||||
buf.Next(skipN)
|
||||
//if w.logger != nil {
|
||||
// w.logger.Infof("ws upgrade success conn upgrade websocket protocol!")
|
||||
//}
|
||||
_ = tempWriter.Flush()
|
||||
data = result.Bytes()
|
||||
w.isUpgrade = true
|
||||
return
|
||||
}
|
||||
|
||||
func (w *WSConn) readWsMessages() (messages []wsutil.Message, err error) {
|
||||
in := &w.buf
|
||||
for {
|
||||
if w.curHeader == nil {
|
||||
if in.Len() < ws.MinHeaderSize { //头长度至少是2
|
||||
return
|
||||
}
|
||||
var head ws.Header
|
||||
//有可能不完整，构建新的 reader 读取 head 读取成功才实际对 in 进行读操作
|
||||
tmpReader := bytes.NewReader(in.Bytes())
|
||||
oldLen := tmpReader.Len()
|
||||
head, err = ws.ReadHeader(tmpReader)
|
||||
skipN := oldLen - tmpReader.Len()
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { //数据不完整
|
||||
return messages, nil
|
||||
}
|
||||
in.Next(skipN)
|
||||
return nil, err
|
||||
}
|
||||
_, err = io.CopyN(&w.cachedBuf, in, int64(skipN))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
w.curHeader = &head
|
||||
}
|
||||
dataLen := (int)(w.curHeader.Length)
|
||||
if dataLen > 0 {
|
||||
if in.Len() >= dataLen {
|
||||
_, err = io.CopyN(&w.cachedBuf, in, int64(dataLen))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
} else { //数据不完整
|
||||
if w.logger != nil {
|
||||
w.logger.Debugf("ws read ws message incomplete data", in.Len(), dataLen)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
if w.curHeader.Fin { //当前 header 已经是一个完整消息
|
||||
messages, err = wsutil.ReadClientMessage(&w.cachedBuf, messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
w.cachedBuf.Reset()
|
||||
}
|
||||
w.curHeader = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (w *WSConn) OnRequest(u []byte) error {
|
||||
parsedURL, err := url.Parse(string(u))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for key, value := range parsedURL.Query() {
|
||||
w.SetParam(key, value[0])
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *WSConn) GetParam(key string) interface{} {
|
||||
return w.param[key]
|
||||
}
|
||||
|
||||
func (w *WSConn) SetParam(key string, values interface{}) {
|
||||
w.param[key] = values
|
||||
}
|
||||
|
||||
func (w *WSConn) RemoteAddr() string {
|
||||
return w.remoteAddr
|
||||
}
|
||||
|
||||
func (w *WSConn) Write(data []byte) error {
|
||||
return w.write(data, ws.OpBinary)
|
||||
}
|
||||
|
||||
func (w *WSConn) Ping() (err error) {
|
||||
return w.write(make([]byte, 0), ws.OpPing)
|
||||
}
|
||||
|
||||
func (w *WSConn) Close() (err error) {
|
||||
defer func(Conn gnet.Conn) {
|
||||
err = Conn.Close()
|
||||
}(w.Conn)
|
||||
return w.write(make([]byte, 0), ws.OpClose)
|
||||
}
|
||||
|
||||
func (w *WSConn) IsClose() bool {
|
||||
return w.isClose
|
||||
}
|
||||
|
||||
func (w *WSConn) write(data []byte, opCode ws.OpCode) error {
|
||||
if w.isClose {
|
||||
return errors.New("connection has close")
|
||||
}
|
||||
buf := bytes.Buffer{}
|
||||
if err := wsutil.WriteServerMessage(&buf, opCode, data); err != nil {
|
||||
return err
|
||||
}
|
||||
return w.Conn.AsyncWrite(buf.Bytes(), nil)
|
||||
}
|
||||
Reference in New Issue
Block a user