rmm2/agent-bsd/ws/client.go

package ws

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"log"
	"net/url"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

type Client struct {
	backendURL string
	agentID    string
	apiKey     string
	insecure   bool
	
	conn       *websocket.Conn
	connMux    sync.RWMutex
	writeMux   sync.Mutex
	
	reconnectInterval time.Duration
	pingInterval      time.Duration
	
	ctx        context.Context
	cancel     context.CancelFunc
	
	messageHandler *Handler
	tunnelManager  *TunnelManager
	ptyManager     *PTYManager
}

type Message struct {
	Type    string                 `json:"type"`
	ID      string                 `json:"id,omitempty"`
	Command string                 `json:"command,omitempty"`
	Event   string                 `json:"event,omitempty"`
	Status  string                 `json:"status,omitempty"`
	Params  map[string]interface{} `json:"params,omitempty"`
	Data    interface{}            `json:"data,omitempty"`
	Error   string                 `json:"error,omitempty"`
}

func NewClient(backendURL, agentID, apiKey string, insecure bool) *Client {
	ctx, cancel := context.WithCancel(context.Background())
	
	client := &Client{
		backendURL:        backendURL,
		agentID:          agentID,
		apiKey:           apiKey,
		insecure:         insecure,
		reconnectInterval: 5 * time.Second,
		pingInterval:     30 * time.Second,
		ctx:              ctx,
		cancel:           cancel,
	}
	
	client.messageHandler = NewHandler(client)
	client.tunnelManager = NewTunnelManager(client)
	client.ptyManager = NewPTYManager(client)
	
	return client
}

func (c *Client) Start() error {
	log.Println("WebSocket Client startet...")
	
	go c.reconnectLoop()
	
	return nil
}

func (c *Client) Stop() {
	log.Println("WebSocket Client stoppt...")
	c.cancel()
	
	c.connMux.Lock()
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}
	c.connMux.Unlock()
}

func (c *Client) reconnectLoop() {
	backoff := c.reconnectInterval
	maxBackoff := 5 * time.Minute
	
	for {
		select {
		case <-c.ctx.Done():
			return
		default:
		}
		
		if err := c.connect(); err != nil {
			log.Printf("WebSocket-Verbindung fehlgeschlagen: %v", err)
			
			select {
			case <-time.After(backoff):
				backoff = time.Duration(float64(backoff) * 1.5)
				if backoff > maxBackoff {
					backoff = maxBackoff
				}
			case <-c.ctx.Done():
				return
			}
			continue
		}
		
		connStart := time.Now()
		log.Println("WebSocket-Verbindung hergestellt")
		
		// Ping/Pong und Message-Handling starten
		c.runConnection()
		
		connDuration := time.Since(connStart)
		log.Printf("WebSocket-Verbindung getrennt (Dauer: %v)", connDuration)
		
		// Alte Verbindung explizit schliessen
		c.connMux.Lock()
		if c.conn != nil {
			c.conn.Close()
			c.conn = nil
		}
		c.connMux.Unlock()
		
		// Minimum-Delay zwischen Reconnects (verhindert Tight-Loop)
		if connDuration < 5*time.Second {
			backoff = time.Duration(float64(backoff) * 1.5)
			if backoff > maxBackoff {
				backoff = maxBackoff
			}
			log.Printf("Verbindung zu kurz, warte %v vor Reconnect", backoff)
			select {
			case <-time.After(backoff):
			case <-c.ctx.Done():
				return
			}
		} else {
			backoff = c.reconnectInterval
		}
	}
}

func (c *Client) connect() error {
	// WebSocket-URL aufbauen
	u, err := url.Parse(c.backendURL)
	if err != nil {
		return fmt.Errorf("Backend-URL parsen: %w", err)
	}
	
	if u.Scheme == "https" {
		u.Scheme = "wss"
	} else {
		u.Scheme = "ws"
	}
	
	u.Path = "/api/v1/agent/ws"
	q := u.Query()
	q.Set("agent_id", c.agentID)
	q.Set("api_key", c.apiKey)
	u.RawQuery = q.Encode()
	
	// Dialer mit TLS-Config
	dialer := websocket.Dialer{
		HandshakeTimeout: 10 * time.Second,
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: c.insecure,
		},
	}
	
	conn, _, err := dialer.Dial(u.String(), nil)
	if err != nil {
		return fmt.Errorf("WebSocket-Dial: %w", err)
	}
	
	c.connMux.Lock()
	c.conn = conn
	c.connMux.Unlock()
	
	return nil
}

func (c *Client) runConnection() {
	// Ping-Ticker starten
	pingTicker := time.NewTicker(c.pingInterval)
	defer pingTicker.Stop()
	
	// Pong-Handler setzen
	c.connMux.RLock()
	conn := c.conn
	c.connMux.RUnlock()
	
	if conn == nil {
		return
	}
	
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(c.pingInterval * 2))
		return nil
	})
	conn.SetReadDeadline(time.Now().Add(c.pingInterval * 2))
	
	// Goroutine für eingehende Nachrichten
	msgChan := make(chan Message, 10)
	errChan := make(chan error, 1)
	
	go c.readMessages(conn, msgChan, errChan)
	
	for {
		select {
		case <-c.ctx.Done():
			return
			
		case <-pingTicker.C:
			c.writeMux.Lock()
			c.connMux.RLock()
			conn := c.conn
			c.connMux.RUnlock()
			if conn != nil {
				if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
					c.writeMux.Unlock()
					log.Printf("Ping senden fehlgeschlagen: %v", err)
					return
				}
			}
			c.writeMux.Unlock()
			
		case msg := <-msgChan:
			go c.handleMessage(msg)
			
		case err := <-errChan:
			log.Printf("WebSocket-Lesefehler: %v", err)
			return
		}
	}
}

func (c *Client) readMessages(conn *websocket.Conn, msgChan chan<- Message, errChan chan<- error) {
	defer close(msgChan)
	
	for {
		messageType, data, err := conn.ReadMessage()
		if err != nil {
			errChan <- err
			return
		}
		
		if messageType == websocket.BinaryMessage {
			// PTY-Messages: session_id beginnt mit "pty-"
			if len(data) > 1 {
				idLen := int(data[0])
				if 1+idLen <= len(data) {
					sessionID := string(data[1 : 1+idLen])
					payload := data[1+idLen:]
					if len(sessionID) >= 4 && sessionID[:4] == "pty-" {
						if err := c.ptyManager.WriteToPTY(sessionID, payload); err != nil {
							log.Printf("PTY-Write fehlgeschlagen: %v", err)
						}
						continue
					}
				}
			}
			// Tunnel-Daten verarbeiten
			if err := c.tunnelManager.ProcessBinaryMessage(data); err != nil {
				log.Printf("Binary-Message verarbeiten fehlgeschlagen: %v", err)
			}
			continue
		}
		
		// JSON-Message parsen
		var msg Message
		if err := json.Unmarshal(data, &msg); err != nil {
			log.Printf("JSON-Message parsen fehlgeschlagen: %v", err)
			continue
		}
		
		select {
		case msgChan <- msg:
		case <-c.ctx.Done():
			return
		}
	}
}

func (c *Client) handleMessage(msg Message) {
	switch msg.Type {
	case "command":
		c.messageHandler.HandleCommand(msg)
	case "event":
		// Events vom Backend verarbeiten (falls nötig)
		log.Printf("Event erhalten: %s", msg.Event)
	default:
		log.Printf("Unbekannter Nachrichtentyp: %s", msg.Type)
	}
}

func (c *Client) SendMessage(msg Message) error {
	c.connMux.RLock()
	conn := c.conn
	c.connMux.RUnlock()
	
	if conn == nil {
		return fmt.Errorf("keine WebSocket-Verbindung")
	}
	
	c.writeMux.Lock()
	err := conn.WriteJSON(msg)
	c.writeMux.Unlock()
	return err
}

func (c *Client) SendBinaryData(data []byte) error {
	c.connMux.RLock()
	conn := c.conn
	c.connMux.RUnlock()
	
	if conn == nil {
		return fmt.Errorf("keine WebSocket-Verbindung")
	}
	
	c.writeMux.Lock()
	err := conn.WriteMessage(websocket.BinaryMessage, data)
	c.writeMux.Unlock()
	return err
}

func (c *Client) IsConnected() bool {
	c.connMux.RLock()
	connected := c.conn != nil
	c.connMux.RUnlock()
	
	return connected
}