package ws import ( "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) agent.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 default: 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 default: 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) agent.Connection.Close() delete(h.agents, agentID) 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-Commands/Responses/Events return nil }