|
|
package server
|
|
|
|
|
|
import (
|
|
|
"context"
|
|
|
"fmt"
|
|
|
"net"
|
|
|
"net/http"
|
|
|
"strings"
|
|
|
"sync"
|
|
|
"time"
|
|
|
|
|
|
"git.noahlan.cn/noahlan/nnet/v2/internal/connection"
|
|
|
"git.noahlan.cn/noahlan/nnet/v2/pkg/config"
|
|
|
pkgerrors "git.noahlan.cn/noahlan/nnet/v2/pkg/errors"
|
|
|
protocolpkg "git.noahlan.cn/noahlan/nnet/v2/pkg/protocol"
|
|
|
unpackerpkg "git.noahlan.cn/noahlan/nnet/v2/pkg/unpacker"
|
|
|
"github.com/gorilla/websocket"
|
|
|
)
|
|
|
|
|
|
// WebSocketServer is a real WebSocket transport server. It reuses the same
|
|
|
// router, codec, protocol and connection manager stack as the TCP server.
|
|
|
type WebSocketServer struct {
|
|
|
*Server
|
|
|
|
|
|
addr string
|
|
|
isWSS bool
|
|
|
httpServer *http.Server
|
|
|
listener net.Listener
|
|
|
upgrader websocket.Upgrader
|
|
|
ctx context.Context
|
|
|
cancel context.CancelFunc
|
|
|
mu sync.RWMutex
|
|
|
started bool
|
|
|
msgHandler *messageHandler
|
|
|
handlerTimeout time.Duration
|
|
|
}
|
|
|
|
|
|
// NewWebSocketServer 创建WebSocket服务器（ws或wss）
|
|
|
func NewWebSocketServer(cfg *config.Config) (*WebSocketServer, error) {
|
|
|
if cfg == nil {
|
|
|
cfg = config.DefaultConfig()
|
|
|
cfg.Addr = "ws://:6995"
|
|
|
}
|
|
|
|
|
|
addr, isWSS := normalizeWebSocketAddr(cfg.Addr)
|
|
|
if isWSS && !cfg.TLSEnabled {
|
|
|
cfg.TLSEnabled = true
|
|
|
}
|
|
|
|
|
|
gnetSrv, err := newGnetServer(cfg, ProtocolTCP)
|
|
|
if err != nil {
|
|
|
return nil, err
|
|
|
}
|
|
|
base := newServerWithGnet(gnetSrv, cfg)
|
|
|
configHelper := newServerConfigHelper(cfg)
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
|
|
return &WebSocketServer{
|
|
|
Server: base,
|
|
|
addr: addr,
|
|
|
isWSS: isWSS,
|
|
|
upgrader: websocket.Upgrader{
|
|
|
CheckOrigin: func(r *http.Request) bool { return true },
|
|
|
},
|
|
|
ctx: ctx,
|
|
|
cancel: cancel,
|
|
|
msgHandler: gnetSrv.eventHandler.msgHandler,
|
|
|
handlerTimeout: configHelper.HandlerTimeout(),
|
|
|
}, nil
|
|
|
}
|
|
|
|
|
|
// NewWSSServer 创建WSS服务器（WebSocket Secure）
|
|
|
func NewWSSServer(cfg *config.Config) (*WebSocketServer, error) {
|
|
|
if cfg == nil {
|
|
|
cfg = config.DefaultConfig()
|
|
|
}
|
|
|
cfg.TLSEnabled = true
|
|
|
|
|
|
addr := cfg.Addr
|
|
|
if strings.HasPrefix(addr, "ws://") {
|
|
|
cfg.Addr = "wss://" + strings.TrimPrefix(addr, "ws://")
|
|
|
} else if !strings.HasPrefix(addr, "wss://") {
|
|
|
cfg.Addr = "wss://" + addr
|
|
|
}
|
|
|
|
|
|
return NewWebSocketServer(cfg)
|
|
|
}
|
|
|
|
|
|
func normalizeWebSocketAddr(addr string) (string, bool) {
|
|
|
switch {
|
|
|
case strings.HasPrefix(addr, "wss://"):
|
|
|
addr = strings.TrimPrefix(addr, "wss://")
|
|
|
return defaultListenAddr(addr), true
|
|
|
case strings.HasPrefix(addr, "ws://"):
|
|
|
addr = strings.TrimPrefix(addr, "ws://")
|
|
|
return defaultListenAddr(addr), false
|
|
|
default:
|
|
|
return defaultListenAddr(addr), false
|
|
|
}
|
|
|
}
|
|
|
|
|
|
func defaultListenAddr(addr string) string {
|
|
|
if addr == "" || addr == ":" {
|
|
|
return ":6995"
|
|
|
}
|
|
|
return addr
|
|
|
}
|
|
|
|
|
|
// Start 启动WebSocket服务器。与TCP服务器保持一致：启动监听后返回。
|
|
|
func (s *WebSocketServer) Start() error {
|
|
|
s.mu.Lock()
|
|
|
if s.started {
|
|
|
s.mu.Unlock()
|
|
|
return pkgerrors.ErrServerAlreadyStarted
|
|
|
}
|
|
|
|
|
|
s.ctx, s.cancel = context.WithCancel(context.Background())
|
|
|
mux := http.NewServeMux()
|
|
|
mux.HandleFunc("/", s.handleWebSocket)
|
|
|
httpServer := &http.Server{Handler: mux}
|
|
|
|
|
|
listener, err := net.Listen("tcp", s.addr)
|
|
|
if err != nil {
|
|
|
s.mu.Unlock()
|
|
|
return fmt.Errorf("failed to listen on %s: %w", s.addr, err)
|
|
|
}
|
|
|
|
|
|
s.listener = listener
|
|
|
s.httpServer = httpServer
|
|
|
s.started = true
|
|
|
s.gnetServer.mu.Lock()
|
|
|
s.gnetServer.started = true
|
|
|
s.gnetServer.mu.Unlock()
|
|
|
s.mu.Unlock()
|
|
|
|
|
|
for _, hook := range s.serverLifecycleHooks {
|
|
|
if err := hook.OnInit(); err != nil {
|
|
|
_ = listener.Close()
|
|
|
return pkgerrors.New("failed to execute OnInit hook").WithCause(err)
|
|
|
}
|
|
|
}
|
|
|
for _, hook := range s.serverLifecycleHooks {
|
|
|
if err := hook.OnStart(); err != nil {
|
|
|
s.gnetServer.logger.Error("OnStart hook error: %v", err)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
s.gnetServer.logger.Info("WebSocket server started on %s", listener.Addr().String())
|
|
|
go s.serve(listener, httpServer)
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
func (s *WebSocketServer) serve(listener net.Listener, httpServer *http.Server) {
|
|
|
var err error
|
|
|
if s.isWSS {
|
|
|
if s.config.TLS == nil || s.config.TLS.CertFile == "" || s.config.TLS.KeyFile == "" {
|
|
|
err = pkgerrors.New("wss requires TLS cert and key files")
|
|
|
} else {
|
|
|
err = httpServer.ServeTLS(listener, s.config.TLS.CertFile, s.config.TLS.KeyFile)
|
|
|
}
|
|
|
} else {
|
|
|
err = httpServer.Serve(listener)
|
|
|
}
|
|
|
if err != nil && err != http.ErrServerClosed {
|
|
|
s.gnetServer.logger.Error("WebSocket server error: %v", err)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
// Stop 停止WebSocket服务器。
|
|
|
func (s *WebSocketServer) Stop() error {
|
|
|
s.mu.Lock()
|
|
|
if !s.started {
|
|
|
s.mu.Unlock()
|
|
|
return pkgerrors.ErrServerNotStarted
|
|
|
}
|
|
|
httpServer := s.httpServer
|
|
|
shutdownTimeout := s.config.ShutdownTimeout
|
|
|
if shutdownTimeout <= 0 {
|
|
|
shutdownTimeout = 30 * time.Second
|
|
|
}
|
|
|
s.started = false
|
|
|
s.gnetServer.mu.Lock()
|
|
|
s.gnetServer.started = false
|
|
|
s.gnetServer.mu.Unlock()
|
|
|
s.cancel()
|
|
|
s.mu.Unlock()
|
|
|
|
|
|
s.gnetServer.logger.Info("Stopping WebSocket server...")
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
|
|
|
defer cancel()
|
|
|
|
|
|
var stopErr error
|
|
|
if httpServer != nil {
|
|
|
stopErr = httpServer.Shutdown(ctx)
|
|
|
}
|
|
|
|
|
|
for _, conn := range s.connManager.GetAll() {
|
|
|
_ = conn.Close()
|
|
|
}
|
|
|
|
|
|
ticker := time.NewTicker(100 * time.Millisecond)
|
|
|
defer ticker.Stop()
|
|
|
timeout := time.NewTimer(shutdownTimeout)
|
|
|
defer timeout.Stop()
|
|
|
|
|
|
waitConnections:
|
|
|
for {
|
|
|
if s.connManager.Count() == 0 {
|
|
|
break
|
|
|
}
|
|
|
select {
|
|
|
case <-ticker.C:
|
|
|
case <-timeout.C:
|
|
|
s.forceCloseAllConnections()
|
|
|
break waitConnections
|
|
|
}
|
|
|
if s.connManager.Count() == 0 {
|
|
|
break
|
|
|
}
|
|
|
}
|
|
|
|
|
|
for _, hook := range s.serverLifecycleHooks {
|
|
|
if err := hook.OnStop(); err != nil {
|
|
|
s.gnetServer.logger.Error("OnStop hook error: %v", err)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
s.gnetServer.logger.Info("WebSocket server stopped")
|
|
|
return stopErr
|
|
|
}
|
|
|
|
|
|
// Started 检查服务器是否已启动。
|
|
|
func (s *WebSocketServer) Started() bool {
|
|
|
s.mu.RLock()
|
|
|
defer s.mu.RUnlock()
|
|
|
return s.started
|
|
|
}
|
|
|
|
|
|
func (s *WebSocketServer) handleWebSocket(w http.ResponseWriter, r *http.Request) {
|
|
|
wsConn, err := s.upgrader.Upgrade(w, r, nil)
|
|
|
if err != nil {
|
|
|
s.gnetServer.logger.Debug("WebSocket upgrade failed: %v", err)
|
|
|
return
|
|
|
}
|
|
|
|
|
|
conn := connection.NewWebSocketConnection("", wsConn)
|
|
|
connID := conn.ID()
|
|
|
if err := s.connManager.Add(conn); err != nil {
|
|
|
s.gnetServer.logger.Error("Failed to add WebSocket connection: %v", err)
|
|
|
_ = wsConn.Close()
|
|
|
return
|
|
|
}
|
|
|
|
|
|
for _, hook := range s.connLifecycleHooks {
|
|
|
if err := hook.OnOpen(connID, conn.RemoteAddr()); err != nil {
|
|
|
s.gnetServer.logger.Error("OnOpen hook error: %v", err)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
protocol := s.resolveApplicationProtocol()
|
|
|
unpacker := newServerProtocolUnpacker(protocol)
|
|
|
ctxConn := toContextConnection(conn)
|
|
|
|
|
|
defer func() {
|
|
|
for _, hook := range s.connLifecycleHooks {
|
|
|
if err := hook.OnClose(connID, nil); err != nil {
|
|
|
s.gnetServer.logger.Error("OnClose hook error: %v", err)
|
|
|
}
|
|
|
}
|
|
|
_ = s.connManager.Remove(connID)
|
|
|
}()
|
|
|
|
|
|
for {
|
|
|
select {
|
|
|
case <-s.ctx.Done():
|
|
|
return
|
|
|
default:
|
|
|
}
|
|
|
|
|
|
_, data, err := wsConn.ReadMessage()
|
|
|
if err != nil {
|
|
|
return
|
|
|
}
|
|
|
conn.UpdateActive()
|
|
|
|
|
|
messages, err := splitWebSocketMessages(data, unpacker)
|
|
|
if err != nil {
|
|
|
s.gnetServer.logger.Debug("WebSocket unpack failed: %v", err)
|
|
|
return
|
|
|
}
|
|
|
if len(messages) == 0 {
|
|
|
continue
|
|
|
}
|
|
|
|
|
|
for _, message := range messages {
|
|
|
s.msgHandler.handleMessageWithContext(
|
|
|
s.ctx,
|
|
|
ctxConn,
|
|
|
message,
|
|
|
protocol,
|
|
|
s.codecRegistry,
|
|
|
s.handlerTimeout,
|
|
|
)
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
func (s *WebSocketServer) resolveApplicationProtocol() protocolpkg.Protocol {
|
|
|
configHelper := newServerConfigHelper(s.config)
|
|
|
if !configHelper.IsProtocolEncodeEnabled() {
|
|
|
return nil
|
|
|
}
|
|
|
protocol, err := s.protocolManager.Get(s.config.ApplicationProtocol, "")
|
|
|
if err != nil {
|
|
|
s.gnetServer.logger.Debug("Failed to resolve WebSocket application protocol: %v", err)
|
|
|
return nil
|
|
|
}
|
|
|
return protocol
|
|
|
}
|
|
|
|
|
|
type serverUnpackingProtocol interface {
|
|
|
Unpacker() unpackerpkg.Unpacker
|
|
|
}
|
|
|
|
|
|
func newServerProtocolUnpacker(protocol protocolpkg.Protocol) unpackerpkg.Unpacker {
|
|
|
if protocol == nil {
|
|
|
return nil
|
|
|
}
|
|
|
if provider, ok := protocol.(serverUnpackingProtocol); ok {
|
|
|
return provider.Unpacker()
|
|
|
}
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
func splitWebSocketMessages(data []byte, unpacker unpackerpkg.Unpacker) ([][]byte, error) {
|
|
|
if unpacker == nil {
|
|
|
msg := make([]byte, len(data))
|
|
|
copy(msg, data)
|
|
|
return [][]byte{msg}, nil
|
|
|
}
|
|
|
messages, _, _, err := unpacker.Unpack(data)
|
|
|
return messages, err
|
|
|
}
|