broad-announce/internal/wsclient/client.go

// Package wsclient is a thin wrapper around gorilla/websocket
// for the M5 WebSocket ingest path. Used by:
//
//   - loadgen/cmd/ws (the loadgen-ws publisher)
//   - scripts/m5_smoke.sh test programs (the failure-path
//     binaries)
//
// The wrapper hides the URL/dial, the optional Origin, and the
// read/write deadlines. It does NOT hide the JSON envelope
// shape — callers marshal their own alert bodies and parse the
// server's ack frames.
package wsclient

import (
	"errors"
	"fmt"
	"net/url"
	"time"

	"github.com/gorilla/websocket"
)

// Config is the dial parameters.
type Config struct {
	// URL is the WebSocket endpoint, e.g. ws://localhost:8800/v1/ingest/ws
	URL string
	// APIKey is the source API key (the same
	// `acme-001:prom-prod:s3cret-acme` triple the HTTP and MQTT
	// loadgens use). The first frame sent to the server is
	// `{"api_key": "..."}`. The server replies with either
	// `{"ready": true}` or `{"error": "..."}`. The Connect
	// helper stashes that first reply on the client so callers
	// can distinguish "auth accepted" from "auth rejected"
	// without sending a test alert.
	APIKey       string
	Origin       string
	DialTimeout  time.Duration
	WriteTimeout time.Duration
	ReadTimeout  time.Duration
}

// Client is the live WS connection. The auth reply is kept so
// callers can surface "tail: unauthorized" instead of just
// "tail: dial ok".
type Client struct {
	cfg       Config
	conn      *websocket.Conn
	authReply []byte
}

// Connect dials the WS endpoint, sends the auth frame, and
// returns a ready-to-use client. The auth frame is the only
// non-alert frame; subsequent SendAlert calls are pure alert
// frames.
func Connect(cfg Config) (*Client, error) {
	if cfg.URL == "" {
		return nil, errors.New("wsclient: empty URL")
	}
	if cfg.APIKey == "" {
		return nil, errors.New("wsclient: empty APIKey")
	}
	if cfg.DialTimeout == 0 {
		cfg.DialTimeout = 10 * time.Second
	}
	if cfg.WriteTimeout == 0 {
		cfg.WriteTimeout = 30 * time.Second
	}
	if cfg.ReadTimeout == 0 {
		cfg.ReadTimeout = 30 * time.Second
	}
	if _, err := url.Parse(cfg.URL); err != nil {
		return nil, fmt.Errorf("wsclient: bad URL %q: %w", cfg.URL, err)
	}

	dialer := *websocket.DefaultDialer
	dialer.HandshakeTimeout = cfg.DialTimeout

	headers := map[string][]string{}
	if cfg.Origin != "" {
		headers["Origin"] = []string{cfg.Origin}
	}
	conn, _, err := dialer.Dial(cfg.URL, headers)
	if err != nil {
		return nil, fmt.Errorf("wsclient: dial: %w", err)
	}

	// Auth frame
	authFrame := []byte(fmt.Sprintf(`{"api_key":%q}`, cfg.APIKey))
	_ = conn.SetWriteDeadline(time.Now().Add(cfg.WriteTimeout))
	if err := conn.WriteMessage(websocket.TextMessage, authFrame); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("wsclient: write auth: %w", err)
	}

	// Wait for the server's auth reply
	_ = conn.SetReadDeadline(time.Now().Add(cfg.ReadTimeout))
	_, msg, err := conn.ReadMessage()
	if err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("wsclient: read auth reply: %w", err)
	}
	if len(msg) == 0 {
		_ = conn.Close()
		return nil, errors.New("wsclient: empty auth reply")
	}
	return &Client{cfg: cfg, conn: conn, authReply: msg}, nil
}

// AuthReply returns the server's first frame. The caller can
// parse it ({"ready": true} or {"error": "..."}) to surface a
// clear error.
func (c *Client) AuthReply() []byte { return c.authReply }

// SendAlert writes one alert frame and reads one ack frame.
// Both operations use the per-frame timeouts from Config. The
// ack is the server's reply for THIS alert; if the server is
// publishing a tail event mid-pipeline it will NOT interleave
// here (the tail is a separate connection).
func (c *Client) SendAlert(body []byte) ([]byte, error) {
	_ = c.conn.SetWriteDeadline(time.Now().Add(c.cfg.WriteTimeout))
	if err := c.conn.WriteMessage(websocket.TextMessage, body); err != nil {
		return nil, fmt.Errorf("wsclient: write alert: %w", err)
	}
	_ = c.conn.SetReadDeadline(time.Now().Add(c.cfg.ReadTimeout))
	_, ack, err := c.conn.ReadMessage()
	if err != nil {
		return nil, fmt.Errorf("wsclient: read ack: %w", err)
	}
	return ack, nil
}

// Close sends a graceful close frame and then closes the
// underlying TCP connection. Safe to call multiple times.
//
// We send CloseMessage with code 1000 (normal closure) and an
// empty payload. The server's ReadMessage then returns a
// CloseError with code 1000, which the WS ingest path
// classifies as "closed_clean" instead of
// "closed_protocol_error".
func (c *Client) Close() error {
	if c.conn == nil {
		return nil
	}
	deadline := time.Now().Add(2 * time.Second)
	_ = c.conn.WriteControl(
		websocket.CloseMessage,
		websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
		deadline,
	)
	err := c.conn.Close()
	c.conn = nil
	return err
}

// Conn exposes the underlying gorilla connection. Used by the
// tail client (which doesn't have a request/response model —
// it just reads frames). Wrap with a custom reader if you
// don't want callers reaching in.
func (c *Client) Conn() *websocket.Conn { return c.conn }