aiops/backend/pkg/wshub/manager.go

// 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
}