uopi/internal/server/ws.go

package server

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

	"github.com/coder/websocket"

	"github.com/uopi/uopi/internal/access"
	"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"`

	// 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"`

	// 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
}

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))
	}

	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()

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

	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)
	policy *access.Policy // global access-level enforcement

	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(),
			}
			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)
	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
	}

	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)
	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)
		c.sendError(ctx, "WRITE_ERROR", err.Error())
	}
}

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 ───────────────────────────────────────────────────────────────────

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"
	}
}