broad-announce/cmd/ingestd/wstail.go

// WebSocket live-tail endpoint for ingestd (M5). The
// server-sent companion to /v1/ingest/ws: the operator opens
// this socket and receives a stream of accepted alert events.
//
// Endpoint:
//
//	GET /v1/tail/ws[?company_id=acme-001&token=***]    HTTP Upgrade → WebSocket
//
// Auth: a static token. The token is supplied as
//   - Authorization: Bearer <token>
//   - X-BA-Tail-Token: <token>
//   - ?token=<token> query param  (for wscat)
//
// The static token is set via BA_INGESTD_TAIL_TOKEN. M11 will
// swap this for a JWT signed by admind.
//
// Protocol (server → client only):
//
//   {"alert_id":"...","company_id":"...","source_id":"...",
//    "severity":"...","title":"...","received_at":"...",
//    "transport":"http|mqtt|ws","dedupe_count":N}
//
// One frame per accepted alert. The full alert payload is NOT
// included (could be 256 KB); the operator gets the metadata
// needed to filter, then drills in via the per-alert GET
// endpoint (M9) for the body.
//
// Backpressure: a slow client has events dropped (counter
// ticks); the producer never blocks. The client may keep the
// connection open and catch up; we don't disconnect on drops.

package main

import (
	"log/slog"
	"net/http"
	"strings"
	"time"

	"git3.techno-world.net/lrosales/broad-announce/internal/observability"
	"git3.techno-world.net/lrosales/broad-announce/internal/tailhub"
)

// wsTailDeps is the WS tail handler's dependency set.
type wsTailDeps struct {
	// Token is the static auth token (env: BA_INGESTD_TAIL_TOKEN).
	// Empty disables the endpoint (the route is still wired but
	// every request gets 503). This keeps the safety property
	// "if the env var is unset, the endpoint is a no-op".
	Token string

	// Hub is the in-process tail hub. Subscriptions are
	// created on each upgrade and released on close.
	Hub *tailhub.Hub

	// Metrics is the shared observability bundle. We update
	// tail_subscribers (gauge) and tail_dropped_total on
	// disconnect / slow consumer.
	Metrics *observability.IngestdMetrics

	// Logger is the tail's per-handler logger. Set by main.
	Logger *slog.Logger
}

// handleTail is the /v1/tail/ws upgrade handler. It runs in
// its own goroutine per connection (the gorilla default).
func (d *wsTailDeps) handleTail(w http.ResponseWriter, r *http.Request) {
	if d.Token == "" {
		http.Error(w, "tail endpoint disabled (BA_INGESTD_TAIL_TOKEN unset)", http.StatusServiceUnavailable)
		return
	}
	if !d.checkToken(r) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	if d.Hub == nil {
		http.Error(w, "tail hub not configured", http.StatusServiceUnavailable)
		return
	}

	// Company filter (optional)
	companyID := strings.TrimSpace(r.URL.Query().Get("company_id"))

	// Subscribe BEFORE the upgrade so events that arrive in
	// the window between dial-completion and the per-goroutine
	// stream-loop startup are not lost. The hub's buffered
	// channel (64) absorbs the burst; if the client is too
	// slow, the drops counter ticks.
	sub := d.Hub.Subscribe(tailhub.Filter{CompanyID: companyID})
	// Defer Unsubscribe BEFORE the deferred metric update so
	// the gauge observes the post-Unsubscribe subscriber count
	// (defers run LIFO).
	defer func() {
		sub.Unsubscribe()
		if d.Metrics != nil {
			d.Metrics.TailSubscribers.Set(float64(d.Hub.Stats().Subscribers))
		}
	}()
	if d.Metrics != nil {
		d.Metrics.TailSubscribers.Set(float64(d.Hub.Stats().Subscribers))
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}

	// Stream loop. We hold the conn write-locked while sending
	// each frame; gorilla's NextWriter handles that. We also
	// poll for client-initiated close (close frame or read
	// error) to break the loop promptly.
	conn.SetReadLimit(512) // we don't read frames from the
	// client, but gorilla requires a non-zero read limit.
	// 512 bytes is enough for an empty close frame.

	stop := make(chan struct{})
	go func() {
		for {
			if _, _, err := conn.NextReader(); err != nil {
				close(stop)
				return
			}
		}
	}()

	for {
		select {
		case <-stop:
			return
		case ev, ok := <-sub.C:
			if !ok {
				return
			}
			_ = conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := conn.WriteJSON(ev); err != nil {
				// Could be a slow client; the drop is already
				// counted in sub.Drops.
				if d.Metrics != nil {
					d.Metrics.TailDropped.WithLabelValues("write_error").Inc()
				}
				return
			}
			// Per-subscription drops that happened in the hub
			// (channel full). Tick our tail_dropped metric so
			// the operator sees the backpressure in Grafana.
			//
			// We use Add(Drops.Load()) and then reset Drops to
			// 0 so each tick reflects the events that have
			// happened since the last successful write. The
			// counter still increments monotonically (Add is
			// monotonic), which is what dashboards want.
			if d.Metrics != nil && sub.Drops.Load() > 0 {
				dropped := sub.Drops.Load()
				sub.Drops.Store(0)
				d.Metrics.TailDropped.WithLabelValues("slow_consumer").Add(float64(dropped))
			}
		}
	}
}

// checkToken validates the per-request token. The order is:
//   1. Authorization: Bearer <token>
//   2. X-BA-Tail-Token: <token>
//   3. ?token=<token>
func (d *wsTailDeps) checkToken(r *http.Request) bool {
	if h := r.Header.Get("Authorization"); h != "" {
		const p = "Bearer "
		if strings.HasPrefix(h, p) && strings.TrimPrefix(h, p) == d.Token {
			return true
		}
	}
	if h := r.Header.Get("X-BA-Tail-Token"); h != "" && h == d.Token {
		return true
	}
	if q := r.URL.Query().Get("token"); q != "" && q == d.Token {
		return true
	}
	return false
}