feat(ws): A3 WebSocket 统一连接管理 — 心跳/注册表/监控/优雅关闭
- 新增 pkg/wshub/manager.go 统一 WebSocket 连接管理器 (WSManager) - 6 个 WebSocket handler 全部迁移: WebShell/kubectl/PodExec/Audit/Metrics/Pipeline - 心跳机制: ping(30s) + pong handler + read deadline(60s) 自动检测死连接 - 连接注册表: 元数据(类型/用户/远程地址/连接时间/最后活跃) - 每用户连接限制 MaxConnsPerUser=15 防资源耗尽 - 监控端点: GET /api/v1/ws/stats + GET /api/v1/ws/connections - 优雅关闭: CloseAll() 发送 CloseGoingAway 帧并关闭所有连接 - 消除 3 个重复 websocket.Upgrader 变量统一为 WSManager 内部 upgrader - 6/6 A3 测试通过, 17/17 回归测试通过
This commit is contained in:
attacker committed
1 parent
00a98e210d
commit
e4f1018ca1
13 files changed
+508
-58
No files matched your search
@@ -0,0 +1,6 @@
|
||||
package wshub
|
||||
|
||||
import "errors"
|
||||
|
||||
// ErrTooManyConnections is returned when a user exceeds MaxConnsPerUser.
|
||||
var ErrTooManyConnections = errors.New("too many websocket connections for user")
|
||||
@@ -0,0 +1,299 @@
|
||||
// Package wshub provides a unified WebSocket connection manager for the AIOps platform.
|
||||
// It centralizes connection lifecycle management: registration, heartbeat (ping/pong),
|
||||
// read/write deadlines, metrics, and graceful shutdown.
|
||||
//
|
||||
// All WebSocket handlers (WebShell, kubectl, Pod exec, audit stream, metrics push,
|
||||
// pipeline log stream) share a single WSManager instance, enabling:
|
||||
// - Dead-client detection via ping/pong + read deadlines
|
||||
// - Central connection registry with metadata (type, user, remote addr, uptime)
|
||||
// - Real-time monitoring via /api/v1/ws/connections and /api/v1/ws/stats
|
||||
// - Graceful shutdown (CloseAll closes every tracked connection)
|
||||
// - Per-user connection limits (resource exhaustion protection)
|
||||
package wshub
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// Connection types — identifies which feature a WebSocket serves.
|
||||
const (
|
||||
TypeTerminal = "terminal" // WebShell SSH terminal
|
||||
TypeKubectl = "kubectl" // kubectl command terminal
|
||||
TypePodExec = "pod_exec" // K8s pod exec terminal
|
||||
TypeAudit = "audit" // Audit real-time stream (live monitor)
|
||||
TypeMetrics = "metrics" // Dashboard metrics push
|
||||
TypePipeline = "pipeline" // Pipeline run log stream
|
||||
)
|
||||
|
||||
// Config tunes the manager behavior.
|
||||
type Config struct {
|
||||
// PingInterval how often to send ping frames to clients.
|
||||
PingInterval time.Duration
|
||||
// PongTimeout is the max wait for a pong before treating the connection as dead.
|
||||
PongTimeout time.Duration
|
||||
// WriteTimeout is the max time for a single write operation.
|
||||
WriteTimeout time.Duration
|
||||
// MaxConnsPerUser limits how many concurrent WS connections a single user can hold.
|
||||
// 0 = unlimited.
|
||||
MaxConnsPerUser int
|
||||
// ReadBufferSize / WriteBufferSize for the internal upgrader.
|
||||
ReadBufferSize int
|
||||
WriteBufferSize int
|
||||
}
|
||||
|
||||
// DefaultConfig returns production-safe defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
PingInterval: 30 * time.Second,
|
||||
PongTimeout: 60 * time.Second,
|
||||
WriteTimeout: 10 * time.Second,
|
||||
MaxConnsPerUser: 15,
|
||||
ReadBufferSize: 1024,
|
||||
WriteBufferSize: 1024,
|
||||
}
|
||||
}
|
||||
|
||||
// ConnInfo holds metadata for a tracked WebSocket connection.
|
||||
type ConnInfo struct {
|
||||
ID uint64
|
||||
Type string
|
||||
UserID uint
|
||||
Username string
|
||||
RemoteAddr string
|
||||
UserAgent string
|
||||
ConnectTime time.Time
|
||||
LastActivity time.Time
|
||||
conn *websocket.Conn
|
||||
mgr *WSManager
|
||||
}
|
||||
|
||||
// Uptime returns how long the connection has been open.
|
||||
func (ci *ConnInfo) Uptime() time.Duration {
|
||||
return time.Since(ci.ConnectTime)
|
||||
}
|
||||
|
||||
// WSManager is the central WebSocket connection manager.
|
||||
// Create one per application (typically in router/routes.go) and inject it into
|
||||
// every WebSocket handler.
|
||||
type WSManager struct {
|
||||
cfg Config
|
||||
upgrader websocket.Upgrader
|
||||
|
||||
conns sync.Map // id → *ConnInfo
|
||||
next uint64 // atomic counter for connection IDs
|
||||
|
||||
mu sync.Mutex
|
||||
closed bool
|
||||
}
|
||||
|
||||
// New creates a WSManager with the given config.
|
||||
func New(cfg Config) *WSManager {
|
||||
if cfg.PingInterval <= 0 {
|
||||
cfg = DefaultConfig()
|
||||
}
|
||||
return &WSManager{
|
||||
cfg: cfg,
|
||||
upgrader: websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true // JWT auth is enforced by middleware before upgrade
|
||||
},
|
||||
ReadBufferSize: cfg.ReadBufferSize,
|
||||
WriteBufferSize: cfg.WriteBufferSize,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Upgrader returns the underlying gorilla upgrader for handlers that need
|
||||
// low-level access (e.g. custom subprotocols). Most handlers should use Upgrade() instead.
|
||||
func (m *WSManager) Upgrader() *websocket.Upgrader {
|
||||
return &m.upgrader
|
||||
}
|
||||
|
||||
// Upgrade performs the HTTP→WebSocket upgrade, registers the connection,
|
||||
// and returns the raw conn + ConnInfo. The caller MUST defer Unregister(info).
|
||||
//
|
||||
// connType must be one of the Type* constants. userID/username come from JWT context.
|
||||
func (m *WSManager) Upgrade(c *gin.Context, connType string, userID uint, username string) (*websocket.Conn, *ConnInfo, error) {
|
||||
// Enforce per-user connection limit (resource exhaustion protection).
|
||||
if m.cfg.MaxConnsPerUser > 0 && userID > 0 {
|
||||
if m.countByUser(userID) >= m.cfg.MaxConnsPerUser {
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{
|
||||
"code": 429,
|
||||
"message": "WebSocket 连接数超出限制，请关闭其他终端后重试",
|
||||
})
|
||||
return nil, nil, ErrTooManyConnections
|
||||
}
|
||||
}
|
||||
|
||||
conn, err := m.upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
id := atomic.AddUint64(&m.next, 1)
|
||||
info := &ConnInfo{
|
||||
ID: id,
|
||||
Type: connType,
|
||||
UserID: userID,
|
||||
Username: username,
|
||||
RemoteAddr: c.ClientIP(),
|
||||
UserAgent: c.GetHeader("User-Agent"),
|
||||
ConnectTime: time.Now(),
|
||||
LastActivity: time.Now(),
|
||||
conn: conn,
|
||||
mgr: m,
|
||||
}
|
||||
m.conns.Store(id, info)
|
||||
return conn, info, nil
|
||||
}
|
||||
|
||||
// StartHeartbeat sets up ping/pong + read deadline for a connection.
|
||||
// It launches a goroutine that sends pings at PingInterval and resets the read
|
||||
// deadline when a pong is received. If no pong arrives within PongTimeout, the
|
||||
// connection is forcibly closed (the handler's ReadMessage will return an error,
|
||||
// unblocking its read loop).
|
||||
//
|
||||
// MUST be called before the handler enters its read loop.
|
||||
func (m *WSManager) StartHeartbeat(conn *websocket.Conn, info *ConnInfo) {
|
||||
// Pong handler — resets the read deadline + updates lastActivity.
|
||||
conn.SetPongHandler(func(string) error {
|
||||
info.LastActivity = time.Now()
|
||||
_ = conn.SetReadDeadline(time.Now().Add(m.cfg.PongTimeout))
|
||||
return nil
|
||||
})
|
||||
|
||||
// Initial read deadline.
|
||||
_ = conn.SetReadDeadline(time.Now().Add(m.cfg.PongTimeout))
|
||||
|
||||
go func() {
|
||||
ticker := time.NewTicker(m.cfg.PingInterval)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
// Check if connection is still registered (not unregistered by handler).
|
||||
if _, ok := m.conns.Load(info.ID); !ok {
|
||||
return
|
||||
}
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(m.cfg.WriteTimeout))
|
||||
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||
return // write failed → connection is dead, stop pinging
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// SetWriteDeadline helper — sets a write deadline on the connection for a single write.
|
||||
func (m *WSManager) SetWriteDeadline(conn *websocket.Conn) {
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(m.cfg.WriteTimeout))
|
||||
}
|
||||
|
||||
// Unregister removes a connection from the registry and closes the underlying socket.
|
||||
// Safe to call multiple times (idempotent).
|
||||
func (m *WSManager) Unregister(info *ConnInfo) {
|
||||
if info == nil {
|
||||
return
|
||||
}
|
||||
if _, loaded := m.conns.LoadAndDelete(info.ID); loaded {
|
||||
_ = info.conn.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// CloseAll closes every tracked connection. Called during graceful shutdown.
|
||||
func (m *WSManager) CloseAll() {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if m.closed {
|
||||
return
|
||||
}
|
||||
m.closed = true
|
||||
|
||||
m.conns.Range(func(key, val any) bool {
|
||||
m.conns.Delete(key)
|
||||
if info, ok := val.(*ConnInfo); ok && info != nil {
|
||||
// Send a close frame so clients disconnect cleanly.
|
||||
_ = info.conn.WriteMessage(
|
||||
websocket.CloseMessage,
|
||||
websocket.FormatCloseMessage(websocket.CloseGoingAway, "server shutdown"),
|
||||
)
|
||||
_ = info.conn.Close()
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// Stats returns aggregated connection statistics.
|
||||
type Stats struct {
|
||||
Total int `json:"total"`
|
||||
ByType map[string]int `json:"by_type"`
|
||||
ByUser map[string]int `json:"by_user"`
|
||||
ByTypeUser map[string]int `json:"by_type_user"` // "type:user" → count
|
||||
}
|
||||
|
||||
// GetStats returns current aggregate statistics.
|
||||
func (m *WSManager) GetStats() Stats {
|
||||
s := Stats{ByType: make(map[string]int), ByUser: make(map[string]int), ByTypeUser: make(map[string]int)}
|
||||
m.conns.Range(func(_, val any) bool {
|
||||
if info, ok := val.(*ConnInfo); ok {
|
||||
s.Total++
|
||||
s.ByType[info.Type]++
|
||||
key := info.Username
|
||||
if key == "" {
|
||||
key = "anonymous"
|
||||
}
|
||||
s.ByUser[key]++
|
||||
s.ByTypeUser[info.Type+":"+key]++
|
||||
}
|
||||
return true
|
||||
})
|
||||
return s
|
||||
}
|
||||
|
||||
// ConnectionDetail is the JSON representation of a tracked connection (for monitoring).
|
||||
type ConnectionDetail struct {
|
||||
ID uint64 `json:"id"`
|
||||
Type string `json:"type"`
|
||||
UserID uint `json:"user_id"`
|
||||
Username string `json:"username"`
|
||||
RemoteAddr string `json:"remote_addr"`
|
||||
ConnectTime time.Time `json:"connect_time"`
|
||||
LastActivity time.Time `json:"last_activity"`
|
||||
UptimeSec int64 `json:"uptime_sec"`
|
||||
}
|
||||
|
||||
// ListConnections returns all active connections sorted by connect time (newest first).
|
||||
func (m *WSManager) ListConnections() []ConnectionDetail {
|
||||
list := []ConnectionDetail{} // initialized to avoid JSON null
|
||||
m.conns.Range(func(_, val any) bool {
|
||||
if info, ok := val.(*ConnInfo); ok {
|
||||
list = append(list, ConnectionDetail{
|
||||
ID: info.ID,
|
||||
Type: info.Type,
|
||||
UserID: info.UserID,
|
||||
Username: info.Username,
|
||||
RemoteAddr: info.RemoteAddr,
|
||||
ConnectTime: info.ConnectTime,
|
||||
LastActivity: info.LastActivity,
|
||||
UptimeSec: int64(info.Uptime().Seconds()),
|
||||
})
|
||||
}
|
||||
return true
|
||||
})
|
||||
return list
|
||||
}
|
||||
|
||||
// countByUser returns the number of active connections for a given user.
|
||||
func (m *WSManager) countByUser(userID uint) int {
|
||||
count := 0
|
||||
m.conns.Range(func(_, val any) bool {
|
||||
if info, ok := val.(*ConnInfo); ok && info.UserID == userID {
|
||||
count++
|
||||
}
|
||||
return true
|
||||
})
|
||||
return count
|
||||
}
|
||||
Reference in new issue
Block a user