rmm2/backend/ws/hub.go

package ws

import (
	"encoding/json"
	"log"
	"sync"
	"time"
)

type AgentConnection struct {
	AgentID    string
	Connection *Connection
	LastSeen   time.Time
}

type Hub struct {
	agents    map[string]*AgentConnection
	agentsMux sync.RWMutex
	
	register   chan *AgentConnection
	unregister chan *AgentConnection
	broadcast  chan []byte
	
	tunnelManager *TunnelManager
}

func NewHub() *Hub {
	hub := &Hub{
		agents:     make(map[string]*AgentConnection),
		register:   make(chan *AgentConnection),
		unregister: make(chan *AgentConnection),
		broadcast:  make(chan []byte, 256),
	}
	
	hub.tunnelManager = NewTunnelManager(hub)
	return hub
}

func (h *Hub) Run() {
	log.Println("WebSocket Hub startet...")
	
	// Cleanup-Timer für inaktive Verbindungen
	cleanupTicker := time.NewTicker(5 * time.Minute)
	defer cleanupTicker.Stop()
	
	for {
		select {
		case agent := <-h.register:
			h.registerAgent(agent)
			
		case agent := <-h.unregister:
			h.unregisterAgent(agent)
			
		case message := <-h.broadcast:
			h.broadcastMessage(message)
			
		case <-cleanupTicker.C:
			h.cleanupInactiveAgents()
		}
	}
}

func (h *Hub) registerAgent(agent *AgentConnection) {
	h.agentsMux.Lock()
	defer h.agentsMux.Unlock()
	
	// Bestehende Verbindung schließen falls vorhanden
	if existing, ok := h.agents[agent.AgentID]; ok {
		log.Printf("Agent %s: Ersetzt bestehende Verbindung", agent.AgentID)
		existing.Connection.Close()
	}
	
	h.agents[agent.AgentID] = agent
	agent.LastSeen = time.Now()
	
	log.Printf("Agent %s: WebSocket-Verbindung registriert", agent.AgentID)
}

func (h *Hub) unregisterAgent(agent *AgentConnection) {
	h.agentsMux.Lock()
	defer h.agentsMux.Unlock()
	
	if existing, ok := h.agents[agent.AgentID]; ok && existing.Connection == agent.Connection {
		delete(h.agents, agent.AgentID)
		existing.Connection.Close()
		log.Printf("Agent %s: WebSocket-Verbindung entfernt", agent.AgentID)
		
		// Alle Tunnel für diesen Agent schließen
		h.tunnelManager.CloseAgentTunnels(agent.AgentID)
	}
}

func (h *Hub) broadcastMessage(message []byte) {
	h.agentsMux.RLock()
	defer h.agentsMux.RUnlock()
	
	for agentID, agent := range h.agents {
		select {
		case agent.Connection.Send <- message:
		default:
			// Agent kann keine Nachrichten mehr empfangen
			log.Printf("Agent %s: Send-Channel blockiert, trenne Verbindung", agentID)
			close(agent.Connection.Send)
			delete(h.agents, agentID)
		}
	}
}

func (h *Hub) SendToAgent(agentID string, message []byte) error {
	h.agentsMux.RLock()
	agent, exists := h.agents[agentID]
	h.agentsMux.RUnlock()
	
	if !exists {
		return ErrAgentNotConnected
	}
	
	agent.LastSeen = time.Now()
	
	select {
	case agent.Connection.Send <- message:
		return nil
	case <-time.After(5 * time.Second):
		return ErrAgentSendBlocked
	}
}

func (h *Hub) SendBinaryToAgent(agentID string, data []byte) error {
	h.agentsMux.RLock()
	agent, exists := h.agents[agentID]
	h.agentsMux.RUnlock()
	
	if !exists {
		return ErrAgentNotConnected
	}
	
	agent.LastSeen = time.Now()
	
	select {
	case agent.Connection.BinarySend <- data:
		return nil
	case <-time.After(5 * time.Second):
		return ErrAgentSendBlocked
	}
}

func (h *Hub) IsAgentConnected(agentID string) bool {
	h.agentsMux.RLock()
	defer h.agentsMux.RUnlock()
	
	_, exists := h.agents[agentID]
	return exists
}

func (h *Hub) GetConnectedAgents() []string {
	h.agentsMux.RLock()
	defer h.agentsMux.RUnlock()
	
	agents := make([]string, 0, len(h.agents))
	for agentID := range h.agents {
		agents = append(agents, agentID)
	}
	
	return agents
}

func (h *Hub) GetAgentStats() map[string]interface{} {
	h.agentsMux.RLock()
	defer h.agentsMux.RUnlock()
	
	stats := map[string]interface{}{
		"connected_agents": len(h.agents),
		"active_tunnels":   h.tunnelManager.GetActiveTunnelCount(),
	}
	
	return stats
}

func (h *Hub) cleanupInactiveAgents() {
	h.agentsMux.Lock()
	defer h.agentsMux.Unlock()
	
	timeout := 10 * time.Minute
	now := time.Now()
	
	for agentID, agent := range h.agents {
		if now.Sub(agent.LastSeen) > timeout {
			log.Printf("Agent %s: Timeout, trenne Verbindung", agentID)
			delete(h.agents, agentID)
			agent.Connection.Close()
			h.tunnelManager.CloseAgentTunnels(agentID)
		}
	}
}

// GetTunnelManager gibt den TunnelManager zurück (für API-Handler)
func (h *Hub) GetTunnelManager() *TunnelManager {
	return h.tunnelManager
}

// RegisterAgent registriert eine neue Agent-Verbindung
func (h *Hub) RegisterAgent(agentID string, conn *Connection) {
	agent := &AgentConnection{
		AgentID:    agentID,
		Connection: conn,
		LastSeen:   time.Now(),
	}
	
	h.register <- agent
}

// UnregisterAgent entfernt eine Agent-Verbindung
func (h *Hub) UnregisterAgent(agentID string, conn *Connection) {
	agent := &AgentConnection{
		AgentID:    agentID,
		Connection: conn,
	}
	
	h.unregister <- agent
}

// ProcessAgentMessage verarbeitet eingehende Nachrichten von Agents
func (h *Hub) ProcessAgentMessage(agentID string, messageType int, data []byte) error {
	h.agentsMux.RLock()
	agent, exists := h.agents[agentID]
	h.agentsMux.RUnlock()
	
	if exists {
		agent.LastSeen = time.Now()
	}
	
	// Binary Messages sind normalerweise Tunnel-Daten
	if messageType == 2 { // websocket.BinaryMessage
		return h.tunnelManager.ProcessBinaryMessage(agentID, data)
	}
	
	// Text Messages sind JSON-Responses/Events vom Agent
	var msg struct {
		Type   string                 `json:"type"`
		ID     string                 `json:"id"`
		Status string                 `json:"status"`
		Data   map[string]interface{} `json:"data"`
	}
	if err := json.Unmarshal(data, &msg); err != nil {
		return nil // Ignorieren wenn nicht parsebar
	}

	// Tunnel-Connect Responses verarbeiten
	if msg.Type == "response" && msg.Status == "ok" && msg.Data != nil {
		if sessionID, ok := msg.Data["session_id"].(string); ok {
			h.tunnelManager.ConfirmSession(sessionID)
		}
	}

	return nil
}