feat 初次提交

2026-01-03 14:26:09 +08:00
parent 18eb946934
commit 3ea3a3ac6d
48 changed files with 5420 additions and 1 deletions
--- a/net/socket/websocket/websocket.go
+++ b/net/socket/websocket/websocket.go
@@ -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)
+		}
+	}
+}
--- a/net/socket/websocket/wsconn.go
+++ b/net/socket/websocket/wsconn.go
@@ -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)
+}