rmm2/agent-windows/ws/client.go

package ws

import (
	"crypto/tls"
	"fmt"
	"log"
	"net/url"
	"time"

	"github.com/gorilla/websocket"
)

type Client struct {
	backendURL   string
	agentID      string
	apiKey       string
	insecure     bool
	conn         *websocket.Conn
	handler      *Handler
	reconnectWait time.Duration
}

func NewClient(backendURL, agentID, apiKey string, insecure bool, handler *Handler) *Client {
	return &Client{
		backendURL:   backendURL,
		agentID:      agentID,
		apiKey:       apiKey,
		insecure:     insecure,
		handler:      handler,
		reconnectWait: 5 * time.Second,
	}
}

func (c *Client) Run(stop <-chan struct{}) {
	for {
		select {
		case <-stop:
			return
		default:
		}

		if err := c.connect(stop); err != nil {
			log.Printf("WS-Verbindung fehlgeschlagen: %v — erneuter Versuch in %v", err, c.reconnectWait)
			select {
			case <-stop:
				return
			case <-time.After(c.reconnectWait):
			}
			if c.reconnectWait < 60*time.Second {
				c.reconnectWait *= 2
			}
		} else {
			c.reconnectWait = 5 * time.Second
		}
	}
}

func (c *Client) connect(stop <-chan struct{}) error {
	u, err := url.Parse(c.backendURL)
	if err != nil {
		return err
	}
	switch u.Scheme {
	case "https":
		u.Scheme = "wss"
	case "http":
		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 := websocket.DefaultDialer
	if c.insecure {
		dialer = &websocket.Dialer{
			TLSClientConfig:  &tls.Config{InsecureSkipVerify: true},
			HandshakeTimeout: 10 * time.Second,
		}
	}

	conn, _, err := dialer.Dial(u.String(), nil)
	if err != nil {
		return fmt.Errorf("Dial fehlgeschlagen: %w", err)
	}
	defer conn.Close()
	c.conn = conn
	log.Printf("WS verbunden mit %s", c.backendURL)

	conn.SetReadDeadline(time.Now().Add(70 * time.Second))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(70 * time.Second))
		return nil
	})

	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			_, msg, err := conn.ReadMessage()
			if err != nil {
				if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
					log.Printf("WS Read-Fehler: %v", err)
				}
				return
			}
			if resp := c.handler.Handle(msg); resp != nil {
				conn.WriteMessage(websocket.TextMessage, resp)
			}
		}
	}()

	select {
	case <-stop:
		conn.WriteMessage(websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		return nil
	case <-done:
		return fmt.Errorf("Verbindung getrennt")
	}
}

func (c *Client) SendMessage(data []byte) error {
	if c.conn == nil {
		return fmt.Errorf("nicht verbunden")
	}
	return c.conn.WriteMessage(websocket.TextMessage, data)
}