rmm2/backend/ws/handler.go

package ws

import (
	"errors"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

var (
	ErrAgentNotConnected = errors.New("Agent nicht verbunden")
	ErrAgentSendBlocked  = errors.New("Agent Send-Channel blockiert")
)

type Connection struct {
	AgentID    string
	Conn       *websocket.Conn
	Send       chan []byte
	BinarySend chan []byte
	Hub        *Hub
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		// Hier könnte Origin-Checking implementiert werden
		return true
	},
}

func NewWebSocketHandler(hub *Hub) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		handleWebSocket(hub, w, r)
	}
}

func handleWebSocket(hub *Hub, w http.ResponseWriter, r *http.Request) {
	// Query-Parameter extrahieren
	agentID := r.URL.Query().Get("agent_id")
	apiKey := r.URL.Query().Get("api_key")
	
	if agentID == "" || apiKey == "" {
		http.Error(w, "agent_id und api_key erforderlich", http.StatusBadRequest)
		return
	}
	
	// API-Key validieren (sollte bereits durch Middleware geschehen)
	// Hier könnten zusätzliche Agent-spezifische Prüfungen stattfinden
	
	// WebSocket-Upgrade
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket-Upgrade fehlgeschlagen: %v", err)
		return
	}
	
	// Connection-Objekt erstellen
	connection := &Connection{
		AgentID:    agentID,
		Conn:       conn,
		Send:       make(chan []byte, 256),
		BinarySend: make(chan []byte, 256),
		Hub:        hub,
	}
	
	// Verbindung bei Hub registrieren
	hub.RegisterAgent(agentID, connection)
	
	// Goroutines für Read/Write starten
	go connection.writePump()
	go connection.readPump()
}

func (c *Connection) readPump() {
	defer func() {
		c.Hub.UnregisterAgent(c.AgentID, c)
		c.Conn.Close()
	}()
	
	// Read-Timeouts setzen
	c.Conn.SetReadLimit(512 * 1024) // 512KB max message size
	c.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.Conn.SetPongHandler(func(string) error {
		c.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})
	
	for {
		messageType, message, err := c.Conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("Agent %s: WebSocket-Fehler: %v", c.AgentID, err)
			}
			break
		}
		
		// Message an Hub weiterleiten
		if err := c.Hub.ProcessAgentMessage(c.AgentID, messageType, message); err != nil {
			log.Printf("Agent %s: Message-Verarbeitung fehlgeschlagen: %v", c.AgentID, err)
		}
	}
}

func (c *Connection) writePump() {
	ticker := time.NewTicker(54 * time.Second) // Ping alle 54s
	defer func() {
		ticker.Stop()
		c.Conn.Close()
	}()
	
	for {
		select {
		case message, ok := <-c.Send:
			c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				c.Conn.WriteMessage(websocket.CloseMessage, nil)
				return
			}
			
			if err := c.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
				log.Printf("Agent %s: Write-Fehler: %v", c.AgentID, err)
				return
			}
			
		case binaryMessage, ok := <-c.BinarySend:
			c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				c.Conn.WriteMessage(websocket.CloseMessage, nil)
				return
			}
			
			if err := c.Conn.WriteMessage(websocket.BinaryMessage, binaryMessage); err != nil {
				log.Printf("Agent %s: Binary-Write-Fehler: %v", c.AgentID, err)
				return
			}
			
		case <-ticker.C:
			c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				log.Printf("Agent %s: Ping-Fehler: %v", c.AgentID, err)
				return
			}
		}
	}
}

func (c *Connection) Close() {
	close(c.Send)
	close(c.BinarySend)
}