uopi/internal/server/ws.go

package server

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/coder/websocket"

	"github.com/uopi/uopi/internal/access"
	"github.com/uopi/uopi/internal/audit"
	"github.com/uopi/uopi/internal/broker"
	"github.com/uopi/uopi/internal/datasource"
	"github.com/uopi/uopi/internal/metrics"
)

// ── Incoming message types ────────────────────────────────────────────────────

type inMsg struct {
	Type string `json:"type"`

	// subscribe / unsubscribe
	Signals []sigRef `json:"signals,omitempty"`

	// write
	DS    string          `json:"ds,omitempty"`
	Name  string          `json:"name,omitempty"`
	Value json.RawMessage `json:"value,omitempty"`

	// dialogResponse — id of the control-logic dialog being answered.
	ID string `json:"id,omitempty"`

	// history
	Start     time.Time `json:"start"`
	End       time.Time `json:"end"`
	MaxPoints int       `json:"maxPoints,omitempty"`
}

type sigRef struct {
	DS   string `json:"ds"`
	Name string `json:"name"`
}

// ── Outgoing message types ────────────────────────────────────────────────────

type outMsg struct {
	Type string `json:"type"`

	// update
	DS       string `json:"ds,omitempty"`
	Name     string `json:"name,omitempty"`
	TS       string `json:"ts,omitempty"`
	Value    any    `json:"value,omitempty"`
	Quality  string `json:"quality,omitempty"`
	Severity int    `json:"severity,omitempty"` // raw EPICS alarm severity (0=NO_ALARM)
	Status   int    `json:"status,omitempty"`   // raw EPICS alarm status

	// meta
	Meta *metaPayload `json:"meta,omitempty"`

	// history
	Points []histPoint `json:"points,omitempty"`

	// error
	Code    string `json:"code,omitempty"`
	Message string `json:"message,omitempty"`
}

type metaPayload struct {
	Type        string            `json:"type"`
	Unit        string            `json:"unit,omitempty"`
	Description string            `json:"description,omitempty"`
	DisplayLow  float64           `json:"displayLow"`
	DisplayHigh float64           `json:"displayHigh"`
	DriveLow    float64           `json:"driveLow,omitempty"`
	DriveHigh   float64           `json:"driveHigh,omitempty"`
	EnumStrings []string          `json:"enumStrings,omitempty"`
	Writable    bool              `json:"writable"`
	Properties  map[string]string `json:"properties,omitempty"`
	Tags        []string          `json:"tags,omitempty"`
}

type histPoint struct {
	TS    string `json:"ts"`
	Value any    `json:"value"`
}

// ── Handler ───────────────────────────────────────────────────────────────────

type wsHandler struct {
	broker *broker.Broker
	log    *slog.Logger
	// userHeader is the HTTP header from which the per-session end-user identity
	// is read (set by a trusted auth proxy). Empty disables per-user identity.
	userHeader string
	// policy enforces the global access level (blacklist) on signal writes.
	policy *access.Policy
	// audit records signal writes (never nil; audit.Nop when disabled).
	audit audit.Recorder
	// dialogs fans control-logic dialogs to clients and routes responses.
	dialogs *DialogHub
}

func (h *wsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Capture the end-user identity from the trusted proxy header (if enabled)
	// before the connection is upgraded, while the original headers are present.
	var user string
	if h.userHeader != "" {
		user = strings.TrimSpace(r.Header.Get(h.userHeader))
	}
	// Resolve to the configured default_user when the header is absent so the
	// session carries a real identity (used for audit + EPICS write attribution).
	if h.policy != nil {
		user = h.policy.ResolveUser(user)
	}
	clientIP := clientIP(r)

	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
		// Permissive origin check for LAN / SSH-tunnel use; tighten when auth lands.
		InsecureSkipVerify: true,
	})
	if err != nil {
		h.log.Error("ws accept failed", "err", err)
		return
	}
	metrics.IncWsConns()
	defer metrics.DecWsConns()

	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()

	rec := h.audit
	if rec == nil {
		rec = audit.Nop()
	}

	c := &wsClient{
		conn:     conn,
		broker:   h.broker,
		user:     user,
		ip:       clientIP,
		policy:   h.policy,
		audit:    rec,
		dialogs:  h.dialogs,
		outCh:    make(chan []byte, 512),
		updateCh: make(chan broker.Update, 1024),
		subs:     make(map[broker.SignalRef]func()),
		log:      h.log,
	}

	// Register for control-logic dialog pushes for this client's identity.
	if h.dialogs != nil {
		h.dialogs.add(c)
		defer h.dialogs.remove(c)
	}

	var wg sync.WaitGroup
	wg.Add(3)
	go func() { defer wg.Done(); c.readLoop(ctx, cancel) }()
	go func() { defer wg.Done(); c.dispatchLoop(ctx) }()
	go func() { defer wg.Done(); c.writeLoop(ctx, cancel) }()
	wg.Wait()

	// Cancel all subscriptions on disconnect.
	c.subsMu.Lock()
	for _, cancelSub := range c.subs {
		cancelSub()
	}
	c.subsMu.Unlock()

	_ = conn.Close(websocket.StatusNormalClosure, "")
}

// ── wsClient ──────────────────────────────────────────────────────────────────

type wsClient struct {
	conn   *websocket.Conn
	broker *broker.Broker
	log    *slog.Logger
	user   string         // end-user identity from the trusted proxy header ("" if none)
	ip     string         // client address, for audit attribution
	policy  *access.Policy // global access-level enforcement
	audit   audit.Recorder // signal-write audit recorder (never nil)
	dialogs *DialogHub     // control-logic dialog fan-out (nil if disabled)

	outCh    chan []byte        // serialised outgoing messages
	updateCh chan broker.Update // raw updates from the broker

	subsMu sync.Mutex
	subs   map[broker.SignalRef]func() // ref → cancel
}

// readLoop reads incoming WebSocket messages and dispatches them.
// It cancels ctx (and therefore the other goroutines) on any error.
func (c *wsClient) readLoop(ctx context.Context, cancel context.CancelFunc) {
	defer cancel()
	for {
		_, data, err := c.conn.Read(ctx)
		if err != nil {
			if !errors.Is(err, context.Canceled) {
				c.log.Debug("ws read", "err", err)
			}
			return
		}
		c.handleMessage(ctx, data)
	}
}

// dispatchLoop converts broker updates to JSON and enqueues them for writing.
func (c *wsClient) dispatchLoop(ctx context.Context) {
	for {
		select {
		case u := <-c.updateCh:
			if u.Value.MetaUpdate {
				// Datasource refreshed metadata (e.g. after channel connected).
				// Re-send the meta message so the client gets accurate type/unit info.
				c.sendMeta(ctx, u.Ref)
				continue
			}
			msg := outMsg{
				Type:     "update",
				DS:       u.Ref.DS,
				Name:     u.Ref.Name,
				TS:       u.Value.Timestamp.UTC().Format(time.RFC3339Nano),
				Value:    u.Value.Data,
				Quality:  u.Value.Quality.String(),
				Severity: u.Value.Severity,
				Status:   u.Value.Status,
			}
			b, err := json.Marshal(msg)
			if err != nil {
				continue
			}
			select {
			case c.outCh <- b:
			case <-ctx.Done():
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

// writeLoop sends queued messages over the WebSocket connection.
// It cancels ctx on any write error so that readLoop and dispatchLoop exit
// promptly rather than blocking forever waiting for outCh to drain.
func (c *wsClient) writeLoop(ctx context.Context, cancel context.CancelFunc) {
	defer cancel()
	for {
		select {
		case msg := <-c.outCh:
			if err := c.conn.Write(ctx, websocket.MessageText, msg); err != nil {
				return
			}
			metrics.IncMsgOut()
		case <-ctx.Done():
			return
		}
	}
}

// handleMessage parses and routes a single incoming client message.
func (c *wsClient) handleMessage(ctx context.Context, data []byte) {
	metrics.IncMsgIn()
	var msg inMsg
	if err := json.Unmarshal(data, &msg); err != nil {
		c.sendError(ctx, "PARSE_ERROR", "invalid JSON: "+err.Error())
		return
	}

	switch msg.Type {
	case "subscribe":
		c.handleSubscribe(ctx, msg.Signals)
	case "unsubscribe":
		c.handleUnsubscribe(msg.Signals)
	case "write":
		c.handleWrite(ctx, msg)
	case "history":
		c.handleHistory(ctx, msg)
	case "dialogResponse":
		c.handleDialogResponse(ctx, msg)
	default:
		c.sendError(ctx, "UNKNOWN_TYPE", "unknown message type: "+msg.Type)
	}
}

func (c *wsClient) handleSubscribe(ctx context.Context, refs []sigRef) {
	for _, r := range refs {
		ref := broker.SignalRef{DS: r.DS, Name: r.Name}

		c.subsMu.Lock()
		_, already := c.subs[ref]
		c.subsMu.Unlock()
		if already {
			continue
		}

		cancelSub, err := c.broker.Subscribe(ref, c.updateCh)
		if err != nil {
			c.sendError(ctx, "SUBSCRIBE_ERROR", err.Error())
			continue
		}

		c.subsMu.Lock()
		c.subs[ref] = cancelSub
		c.subsMu.Unlock()

		// Send metadata in a goroutine: for EPICS signals GetMetadata may block
		// for several seconds waiting for the CA channel to connect and for
		// ca_pend_io to complete. Running it inline would stall readLoop and
		// prevent any further messages from this client from being processed.
		go c.sendMeta(ctx, ref)
	}
}

func (c *wsClient) handleUnsubscribe(refs []sigRef) {
	for _, r := range refs {
		ref := broker.SignalRef{DS: r.DS, Name: r.Name}
		c.subsMu.Lock()
		if cancel, ok := c.subs[ref]; ok {
			cancel()
			delete(c.subs, ref)
		}
		c.subsMu.Unlock()
	}
}

func (c *wsClient) handleWrite(ctx context.Context, msg inMsg) {
	// Enforce the user's global access level: blacklisted read-only / no-access
	// users may not write signals, regardless of the channel's own writability.
	if c.policy != nil && c.policy.Level(c.policy.ResolveUser(c.user)) < access.LevelWrite {
		c.log.Warn("write: access denied", "ds", msg.DS, "name", msg.Name, "user", c.user)
		c.sendError(ctx, "ACCESS_DENIED", "you do not have write access")
		return
	}

	// Server variables are read-any but write only for control-logic editors, so a
	// panel cannot mutate sequence state unless its user owns control-logic access.
	if msg.DS == "srv" && c.policy != nil && !c.policy.CanEditLogic(c.policy.ResolveUser(c.user)) {
		c.log.Warn("write: server variable denied", "name", msg.Name, "user", c.user)
		c.sendError(ctx, "ACCESS_DENIED", "writing server variables requires control-logic access")
		return
	}

	ds, ok := c.broker.Source(msg.DS)
	if !ok {
		c.log.Warn("write: unknown data source", "ds", msg.DS, "name", msg.Name)
		c.sendError(ctx, "NOT_FOUND", "unknown data source: "+msg.DS)
		return
	}

	// Enforce write permission before touching the value payload.
	meta, err := ds.GetMetadata(ctx, msg.Name)
	if err != nil {
		c.log.Warn("write: GetMetadata failed", "ds", msg.DS, "name", msg.Name, "err", err)
		c.sendError(ctx, "NOT_FOUND", "unknown signal: "+msg.Name)
		return
	}
	if !meta.Writable {
		c.log.Warn("write: signal not writable", "ds", msg.DS, "name", msg.Name)
		c.sendError(ctx, "NOT_WRITABLE", "signal is read-only: "+msg.Name)
		return
	}

	// Decode the JSON value as the appropriate Go type.
	var value any
	if err := json.Unmarshal(msg.Value, &value); err != nil {
		c.log.Warn("write: parse error", "ds", msg.DS, "name", msg.Name, "err", err)
		c.sendError(ctx, "PARSE_ERROR", "invalid value: "+err.Error())
		return
	}

	c.log.Info("write", "ds", msg.DS, "name", msg.Name, "value", value, "user", c.user)
	metrics.IncWrites()
	// Attribute the write to the connecting end-user so data sources (EPICS) can
	// act under that identity rather than the server's.
	ctx = datasource.WithUser(ctx, c.user)
	ev := audit.Event{
		Actor:     c.user,
		ActorType: audit.ActorUser,
		Action:    "signal.write",
		DS:        msg.DS,
		Signal:    msg.Name,
		Value:     formatAuditValue(value),
		IP:        c.ip,
		Outcome:   audit.OutcomeOK,
	}
	if err := ds.Write(ctx, msg.Name, value); err != nil {
		c.log.Warn("write: ds.Write failed", "ds", msg.DS, "name", msg.Name, "err", err)
		ev.Outcome = audit.OutcomeError
		ev.Error = err.Error()
		c.audit.Record(ev)
		c.sendError(ctx, "WRITE_ERROR", err.Error())
		return
	}
	c.audit.Record(ev)
}

// handleDialogResponse routes the answer to a control-logic input dialog. An
// empty/null value means the user dismissed the dialog (drop it without
// writing); otherwise the numeric value is written to the dialog's target.
func (c *wsClient) handleDialogResponse(ctx context.Context, msg inMsg) {
	if c.dialogs == nil || msg.ID == "" {
		return
	}
	if len(msg.Value) == 0 || string(msg.Value) == "null" {
		c.dialogs.cancel(msg.ID)
		return
	}
	var num float64
	if err := json.Unmarshal(msg.Value, &num); err != nil {
		c.dialogs.cancel(msg.ID)
		return
	}
	c.dialogs.respond(ctx, c, msg.ID, num)
}

func (c *wsClient) handleHistory(ctx context.Context, msg inMsg) {
	metrics.IncHistoryReqs()
	ds, ok := c.broker.Source(msg.DS)
	if !ok {
		c.sendError(ctx, "NOT_FOUND", "unknown data source: "+msg.DS)
		return
	}

	maxPts := msg.MaxPoints
	if maxPts <= 0 {
		maxPts = 5000
	}

	vals, err := ds.History(ctx, msg.Name, msg.Start, msg.End, maxPts)
	if err != nil {
		c.sendError(ctx, "HISTORY_ERROR", err.Error())
		return
	}

	pts := make([]histPoint, len(vals))
	for i, v := range vals {
		pts[i] = histPoint{
			TS:    v.Timestamp.UTC().Format(time.RFC3339Nano),
			Value: v.Data,
		}
	}

	resp := outMsg{
		Type:   "history",
		DS:     msg.DS,
		Name:   msg.Name,
		Points: pts,
	}
	b, _ := json.Marshal(resp)
	select {
	case c.outCh <- b:
	case <-ctx.Done():
	}
}

func (c *wsClient) sendMeta(ctx context.Context, ref broker.SignalRef) {
	ds, ok := c.broker.Source(ref.DS)
	if !ok {
		return
	}
	meta, err := ds.GetMetadata(ctx, ref.Name)
	if err != nil {
		return
	}

	msg := outMsg{
		Type: "meta",
		DS:   ref.DS,
		Name: ref.Name,
		Meta: metadataToPayload(meta),
	}
	b, _ := json.Marshal(msg)
	select {
	case c.outCh <- b:
	default:
	}
}

func (c *wsClient) sendError(ctx context.Context, code, message string) {
	msg := outMsg{Type: "error", Code: code, Message: message}
	b, _ := json.Marshal(msg)
	select {
	case c.outCh <- b:
	case <-ctx.Done():
	}
}

// ── Helpers ───────────────────────────────────────────────────────────────────

// clientIP returns the best-effort client address for audit attribution,
// preferring the X-Forwarded-For / X-Real-IP headers set by a reverse proxy and
// falling back to the TCP peer address.
func clientIP(r *http.Request) string {
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		if i := strings.IndexByte(xff, ','); i >= 0 {
			return strings.TrimSpace(xff[:i])
		}
		return strings.TrimSpace(xff)
	}
	if xr := r.Header.Get("X-Real-IP"); xr != "" {
		return strings.TrimSpace(xr)
	}
	if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
		return host
	}
	return r.RemoteAddr
}

// formatAuditValue renders a written value as a compact string for the audit log.
func formatAuditValue(v any) string {
	switch x := v.(type) {
	case string:
		return x
	case float64:
		return strconv.FormatFloat(x, 'g', -1, 64)
	case bool:
		return strconv.FormatBool(x)
	case nil:
		return ""
	default:
		if b, err := json.Marshal(v); err == nil {
			return string(b)
		}
		return ""
	}
}

func metadataToPayload(m datasource.Metadata) *metaPayload {
	return &metaPayload{
		Type:        dataTypeName(m.Type),
		Unit:        m.Unit,
		Description: m.Description,
		DisplayLow:  m.DisplayLow,
		DisplayHigh: m.DisplayHigh,
		DriveLow:    m.DriveLow,
		DriveHigh:   m.DriveHigh,
		EnumStrings: m.EnumStrings,
		Writable:    m.Writable,
		Properties:  m.Properties,
		Tags:        m.Tags,
	}
}

func dataTypeName(t datasource.DataType) string {
	switch t {
	case datasource.TypeFloat64:
		return "float64"
	case datasource.TypeFloat64Array:
		return "float64[]"
	case datasource.TypeString:
		return "string"
	case datasource.TypeInt64:
		return "int64"
	case datasource.TypeBool:
		return "bool"
	case datasource.TypeEnum:
		return "enum"
	default:
		return "unknown"
	}
}