aiops/backend/controller/audit_hub.go

// 终端审计实时流 Hub — 在线会话内容广播给审计查看者 // Logan.Li
package controller

import (
	"log"
	"net/http"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// AuditHub 终端审计流中心 — 管理所有活跃会话的审计订阅者
type AuditHub struct {
	mu       sync.RWMutex
	watchers map[uint]map[chan []byte]struct{} // sessionID -> set of watcher channels
}

// NewAuditHub 创建审计 Hub
func NewAuditHub() *AuditHub {
	return &AuditHub{
		watchers: make(map[uint]map[chan []byte]struct{}),
	}
}

// Subscribe 订阅某个会话的实时输出
func (h *AuditHub) Subscribe(sessionID uint) chan []byte {
	ch := make(chan []byte, 256)
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.watchers[sessionID] == nil {
		h.watchers[sessionID] = make(map[chan []byte]struct{})
	}
	h.watchers[sessionID][ch] = struct{}{}
	return ch
}

// Unsubscribe 取消订阅
func (h *AuditHub) Unsubscribe(sessionID uint, ch chan []byte) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if watchers, ok := h.watchers[sessionID]; ok {
		delete(watchers, ch)
		close(ch)
		if len(watchers) == 0 {
			delete(h.watchers, sessionID)
		}
	}
}

// Broadcast 向所有订阅者广播终端输出数据（非阻塞）
func (h *AuditHub) Broadcast(sessionID uint, data []byte) {
	h.mu.RLock()
	watchers, ok := h.watchers[sessionID]
	h.mu.RUnlock()
	if !ok || len(watchers) == 0 {
		return
	}
	for ch := range watchers {
		select {
		case ch <- data:
		default:
		}
	}
}

// Cleanup 清理某个会话的所有订阅者
func (h *AuditHub) Cleanup(sessionID uint) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if watchers, ok := h.watchers[sessionID]; ok {
		for ch := range watchers {
			close(ch)
		}
		delete(h.watchers, sessionID)
	}
}

var auditUpgrader = websocket.Upgrader{
	CheckOrigin:     func(r *http.Request) bool { return true },
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// ServeAuditStream 处理 /ws/audit/:sessionId WebSocket — 实时推送终端输出
func (h *AuditHub) ServeAuditStream(c *gin.Context, sessionID uint, historyData string, sessionInfo map[string]interface{}) {
	conn, err := auditUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("[AuditStream] WS 升级失败 session=%d: %v", sessionID, err)
		return
	}
	defer conn.Close()

	// 发送会话元信息
	conn.WriteJSON(map[string]interface{}{
		"type":    "session_info",
		"payload": sessionInfo,
	})

	// 发送历史数据（已录制的输出）
	if historyData != "" {
		conn.WriteJSON(map[string]interface{}{
			"type": "history",
			"data": historyData,
		})
	}

	// 订阅实时流
	ch := h.Subscribe(sessionID)
	defer h.Unsubscribe(sessionID, ch)

	// goroutine 读取关闭信号
	go func() {
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				break
			}
		}
	}()

	// 推送实时数据
	for data := range ch {
		if err := conn.WriteJSON(map[string]interface{}{
			"type": "output",
			"data": string(data),
		}); err != nil {
			break
		}
	}
}