Edu/services/push-gateway/internal/ws/handler.go

// Package ws implements the WebSocket upgrade endpoint, the /internal/push
// and /internal/broadcast HTTP APIs consumed by msg, and the
// /internal/online/<userID> presence query.
//
// Authentication:
//   - WebSocket /ws: JWT RS256 validated via shared-go/jwks (caching JWKS from
//     iam /.well-known/jwks.json, refreshed every 5 minutes). DevMode accepts
//     the literal "dev-token" returning a synthetic dev-user subject.
//   - /internal/*: X-Internal-Key header matched against PUSH_INTERNAL_TOKEN
//     (ARB-013 alignment with msg). X-Internal-Token is accepted as a
//     backward-compat alias. DevMode skips the check.
//
// Request body field naming: msg sends userId (camelCase); legacy callers may
// still send user_id (snake_case). Both are accepted by /internal/push.
//
// Heartbeat (RFC 6455 control frames, not text messages):
//   - Client sends Ping every 30s; gorilla/websocket auto-replies with Pong.
//   - Server sets a 60s read deadline; if no frame arrives the connection is
//     closed (idle timeout).
//   - Each Ping refreshes the Redis presence TTL via RefreshPresence.
//
// Per-user connection cap: enforced by Hub.Register (default 5). Exceeding it
// returns a 429 PUSH_TOO_MANY_CONNECTIONS for HTTP callers and a close frame
// (1008 policy violation) for WebSocket callers.
package ws

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/edu-cloud/push-gateway/internal/config"
	"github.com/edu-cloud/push-gateway/internal/hub"
	"github.com/edu-cloud/push-gateway/internal/observability"
	"github.com/edu-cloud/push-gateway/internal/redisclient"
	"github.com/edu-cloud/shared-go/jwks"
	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
	"github.com/gorilla/websocket"
)

// internalTokenHeader is the canonical header for /internal/* authentication
// (ARB-013 alignment with msg). legacyInternalTokenHeader is accepted as a
// backward-compat alias for existing callers.
const (
	internalTokenHeader       = "X-Internal-Key"
	legacyInternalTokenHeader = "X-Internal-Token"
)

// readBufferSize / writeBufferSize tune gorilla/websocket's internal buffers.
const (
	readBufferSize  = 4096
	writeBufferSize = 4096
)

// idleTimeout is the maximum interval with no frames before the connection is
// considered dead. Should be at least 2x the client's heartbeat interval.
const idleTimeout = 60 * time.Second

// writeTimeout bounds blocking writes from the writer goroutine.
const writeTimeout = 10 * time.Second

// closeHandshakeTimeout bounds the close-frame exchange before forcing close.
const closeHandshakeTimeout = 5 * time.Second

// pushResponse is the ActionState-shaped envelope returned by /internal/push
// and /internal/broadcast (see 02 §4.2).
type pushResponse struct {
	Success   bool   `json:"success"`
	Delivered bool   `json:"delivered,omitempty"`
	Online    bool   `json:"online,omitempty"`
	Reached   int    `json:"reached,omitempty"`
	Error     *errBody `json:"error,omitempty"`
}

type errBody struct {
	Code    string `json:"code"`
	Message string `json:"message"`
}

var (
	errMissingToken  = errors.New("missing token")
	errInvalidToken  = errors.New("invalid token")
	errInvalidClaims = errors.New("invalid claims")
	errMissingUserID = errors.New("missing user id")
)

// Handler owns the WebSocket upgrade path and the internal HTTP APIs.
type Handler struct {
	hub            *hub.Hub
	cfg            *config.Config
	redis          *redisclient.Client
	jwksFetcher    *jwks.Fetcher
	metrics        *observability.Metrics
	upgrader       websocket.Upgrader
}

// NewHandler constructs a Handler. cfg supplies auth tokens and origin
// whitelist; redis may be nil in DevMode (presence and cross-instance fanout
// are skipped). jwksFetcher may be nil in DevMode (HS256 fallback).
func NewHandler(h *hub.Hub, cfg *config.Config, r *redisclient.Client, m *observability.Metrics) *Handler {
	allowed := make(map[string]struct{}, len(cfg.AllowedOrigins))
	for _, o := range cfg.AllowedOrigins {
		allowed[o] = struct{}{}
	}
	var fetcher *jwks.Fetcher
	if cfg.JWKSURL != "" {
		fetcher = jwks.NewFetcher(cfg.JWKSURL)
	}
	return &Handler{
		hub:         h,
		cfg:         cfg,
		redis:       r,
		jwksFetcher: fetcher,
		metrics:     m,
		upgrader: websocket.Upgrader{
			ReadBufferSize:  readBufferSize,
			WriteBufferSize: writeBufferSize,
			CheckOrigin:     buildOriginChecker(allowed, cfg.DevMode),
		},
	}
}

// buildOriginChecker returns a CheckOrigin function that allows only the
// configured whitelist. In DevMode with an empty whitelist all origins are
// accepted (local development convenience).
func buildOriginChecker(allowed map[string]struct{}, devMode bool) func(r *http.Request) bool {
	return func(r *http.Request) bool {
		origin := r.Header.Get("Origin")
		if origin == "" {
			// Non-browser clients (curl, Postman) have no Origin header.
			return true
		}
		if devMode && len(allowed) == 0 {
			return true
		}
		_, ok := allowed[origin]
		return ok
	}
}

// HandleWebSocket upgrades an HTTP request to a WebSocket connection. The
// caller must provide a valid JWT (RS256 in production, dev-token in DevMode)
// via the ?token= query parameter or the Authorization: Bearer header.
func (h *Handler) HandleWebSocket(c *gin.Context) {
	userID, err := h.authenticate(c)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusUnauthorized, pushResponse{
			Success: false,
			Error:   &errBody{Code: "PUSH_UNAUTHORIZED", Message: err.Error()},
		})
		return
	}

	conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		// Upgrade already wrote the error response via c.Writer.
		observability.Logger().Warn("ws upgrade failed",
			"user_id", userID, "err", err)
		return
	}
	// Apply per-connection limits.
	conn.SetReadLimit(64 * 1024) // 64KB max message size (02 §11)

	connPtr, err := h.hub.Register(userID, conn)
	if err != nil {
		closeCode := websocket.ClosePolicyViolation
		if errors.Is(err, hub.ErrHubClosing) {
			closeCode = websocket.CloseGoingAway
		}
		_ = conn.WriteControl(
			websocket.CloseMessage,
			websocket.FormatCloseMessage(closeCode, err.Error()),
			time.Now().Add(closeHandshakeTimeout),
		)
		_ = conn.Close()
		if errors.Is(err, hub.ErrTooManyConnections) {
			h.metrics.IncDisconnect("too_many_connections")
		}
		return
	}
	defer h.hub.Unregister(connPtr)

	// Writer goroutine: consumes connPtr.Outgoing() and writes to the socket.
	// Single writer to satisfy gorilla/websocket's "no concurrent writers" rule.
	go h.writerLoop(connPtr)

	// Reader loop: handles Ping/Pong control frames and idle timeout.
	h.readerLoop(connPtr)
}

// writerLoop drains the connection's send channel and writes each message as a
// TextMessage. Exits when the channel is closed by Hub.Unregister / CloseAll.
func (h *Handler) writerLoop(c *hub.Connection) {
	conn := c.Conn()
	for msg := range c.Outgoing() {
		_ = conn.SetWriteDeadline(time.Now().Add(writeTimeout))
		if err := conn.WriteMessage(websocket.TextMessage, msg); err != nil {
			observability.Logger().Debug("ws write failed",
				"user_id", c.UserID, "conn_id", c.ConnID, "err", err)
			return
		}
	}
}

// readerLoop reads frames until the connection closes. It enforces the 60s
// idle timeout via SetReadDeadline and refreshes the Redis presence TTL on
// each Pong handler invocation (RFC 6455 control frame heartbeat).
func (h *Handler) readerLoop(c *hub.Connection) {
	conn := c.Conn()
	_ = conn.SetReadDeadline(time.Now().Add(idleTimeout))
	conn.SetPongHandler(func(appData string) error {
		_ = conn.SetReadDeadline(time.Now().Add(idleTimeout))
		c.TouchPing()
		h.metrics.Heartbeats.Inc()
		if h.redis != nil {
			ctx, cancel := contextWithTimeout(500 * time.Millisecond)
			defer cancel()
			_ = h.redis.RefreshPresence(ctx, c.UserID)
		}
		return nil
	})

	for {
		messageType, _, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
				h.metrics.IncDisconnect("closed")
			} else if websocket.IsUnexpectedCloseError(err) {
				h.metrics.IncDisconnect("error")
				observability.Logger().Debug("ws read failed",
					"user_id", c.UserID, "conn_id", c.ConnID, "err", err)
			} else {
				h.metrics.IncDisconnect("idle")
			}
			return
		}
		// We accept only Ping/Pong control frames (handled above) and ignore
		// any client-sent Text/Binary messages. Clients must NOT send text
		// "ping" messages (legacy protocol removed in P5).
		if messageType == websocket.TextMessage || messageType == websocket.BinaryMessage {
			// Acknowledge application-level ack messages with a no-op; ignore
			// everything else. Future P6 reconnect protocol may use this path.
			continue
		}
	}
}

// PushHandler implements POST /internal/push: directed push to a single user.
// Response carries delivered + online so msg (ai10) can decide offline fallback.
//
// Request body accepts both camelCase (userId, msg's preferred format) and
// snake_case (user_id, legacy format) field names for the user identifier.
func (h *Handler) PushHandler(c *gin.Context) {
	if !h.checkInternalToken(c) {
		return
	}

	var req struct {
		UserID    string         `json:"userId"`  // camelCase (msg canonical, ARB-013)
		UserIDLeg string         `json:"user_id"` // snake_case (legacy)
		Event     string         `json:"event" binding:"required"`
		Data      map[string]any `json:"data"`
		TTL       *int           `json:"ttl,omitempty"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, pushResponse{
			Success: false,
			Error:   &errBody{Code: "PUSH_INVALID_REQUEST", Message: err.Error()},
		})
		return
	}
	// Prefer userId (camelCase); fall back to user_id (snake_case, legacy).
	if req.UserID == "" {
		req.UserID = req.UserIDLeg
	}
	if req.UserID == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, pushResponse{
			Success: false,
			Error:   &errBody{Code: "PUSH_INVALID_REQUEST", Message: "userId (or user_id) is required"},
		})
		return
	}

	message, err := buildMessage(req.Event, req.Data)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, pushResponse{
			Success: false,
			Error:   &errBody{Code: "PUSH_INVALID_PAYLOAD", Message: err.Error()},
		})
		return
	}

	// Try local delivery first.
	delivered := h.hub.SendToUser(req.UserID, message)
	if delivered > 0 {
		h.metrics.IncPushed(req.Event, "delivered")
		c.JSON(http.StatusOK, pushResponse{Success: true, Delivered: true, Online: true})
		return
	}

	// Fall back to cross-instance Redis Pub/Sub fanout.
	if h.redis != nil {
		ctx, cancel := contextWithTimeout(2 * time.Second)
		defer cancel()
		online, _, err := h.redis.IsOnline(ctx, req.UserID)
		if err != nil {
			observability.Logger().Warn("redis IsOnline failed",
				"user_id", req.UserID, "err", err)
			// Treat Redis errors as "online unknown" — try publishing anyway.
			online = true
		}
		if online {
			err := h.redis.PublishUser(ctx, req.UserID, redisclient.CrossInstanceMessage{
				Event: req.Event,
				Data:  mustMarshal(req.Data),
			})
			if err != nil {
				observability.Logger().Warn("redis PublishUser failed",
					"user_id", req.UserID, "err", err)
				c.JSON(http.StatusOK, pushResponse{Success: false, Delivered: false, Online: true})
				return
			}
			// Pub/Sub is fire-and-forget: assume delivered if any instance held the user.
			h.metrics.IncPushed(req.Event, "delivered")
			c.JSON(http.StatusOK, pushResponse{Success: true, Delivered: true, Online: true})
			return
		}
		// Not online anywhere — msg should fall back to offline push (SMS/email).
		h.metrics.IncPushed(req.Event, "offline")
		c.JSON(http.StatusOK, pushResponse{Success: true, Delivered: false, Online: false})
		return
	}

	// No Redis available (DevMode): report offline so msg falls back.
	h.metrics.IncPushed(req.Event, "offline")
	c.JSON(http.StatusOK, pushResponse{Success: true, Delivered: false, Online: false})
}

// BroadcastHandler implements POST /internal/broadcast: push to every online
// client across all instances.
func (h *Handler) BroadcastHandler(c *gin.Context) {
	if !h.checkInternalToken(c) {
		return
	}

	var req struct {
		Event  string         `json:"event" binding:"required"`
		Data   map[string]any `json:"data"`
		Filter map[string]any `json:"filter,omitempty"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, pushResponse{
			Success: false,
			Error:   &errBody{Code: "PUSH_INVALID_REQUEST", Message: err.Error()},
		})
		return
	}

	message, err := buildMessage(req.Event, req.Data)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, pushResponse{
			Success: false,
			Error:   &errBody{Code: "PUSH_INVALID_PAYLOAD", Message: err.Error()},
		})
		return
	}

	local := h.hub.Broadcast(message)

	if h.redis != nil {
		ctx, cancel := contextWithTimeout(2 * time.Second)
		defer cancel()
		_ = h.redis.PublishBroadcast(ctx, redisclient.CrossInstanceMessage{
			Event: req.Event,
			Data:  mustMarshal(req.Data),
		})
	}

	h.metrics.IncPushed(req.Event, "delivered")
	c.JSON(http.StatusOK, pushResponse{Success: true, Reached: local})
}

// OnlineHandler implements GET /internal/online/<userID>: returns whether the
// user has any live connection across the cluster.
func (h *Handler) OnlineHandler(c *gin.Context) {
	if !h.checkInternalToken(c) {
		return
	}
	userID := c.Param("userID")
	if userID == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, pushResponse{
			Success: false,
			Error:   &errBody{Code: "PUSH_INVALID_REQUEST", Message: "missing userID"},
		})
		return
	}

	// Local shortcut.
	if h.hub.HasUser(userID) {
		c.JSON(http.StatusOK, gin.H{"online": true, "instances": []string{h.cfg.InstanceID}})
		return
	}

	if h.redis != nil {
		ctx, cancel := contextWithTimeout(1 * time.Second)
		defer cancel()
		online, members, err := h.redis.IsOnline(ctx, userID)
		if err != nil {
			observability.Logger().Warn("redis IsOnline failed",
				"user_id", userID, "err", err)
			c.JSON(http.StatusOK, gin.H{"online": false, "instances": []string{}})
			return
		}
		c.JSON(http.StatusOK, gin.H{"online": online, "instances": members})
		return
	}
	c.JSON(http.StatusOK, gin.H{"online": false, "instances": []string{}})
}

// checkInternalToken validates the internal token header. The canonical header
// is X-Internal-Key (ARB-013 alignment with msg); X-Internal-Token is accepted
// as a backward-compat alias. Returns false when the request has been aborted
// (caller should return immediately).
func (h *Handler) checkInternalToken(c *gin.Context) bool {
	if h.cfg.DevMode {
		return true
	}
	if h.cfg.InternalAPIToken == "" {
		c.AbortWithStatusJSON(http.StatusUnauthorized, pushResponse{
			Success: false,
			Error:   &errBody{Code: "PUSH_UNAUTHORIZED", Message: "internal token not configured"},
		})
		return false
	}
	token := c.GetHeader(internalTokenHeader)
	if token == "" {
		token = c.GetHeader(legacyInternalTokenHeader)
	}
	if token != h.cfg.InternalAPIToken {
		c.AbortWithStatusJSON(http.StatusUnauthorized, pushResponse{
			Success: false,
			Error:   &errBody{Code: "PUSH_UNAUTHORIZED", Message: "invalid or missing internal token"},
		})
		return false
	}
	return true
}

// authenticate validates the WebSocket JWT. Production uses RS256 via the
// shared-go/jwks Fetcher; DevMode accepts the literal "dev-token" string.
func (h *Handler) authenticate(c *gin.Context) (string, error) {
	tokenStr := c.Query("token")
	if tokenStr == "" {
		if auth := c.GetHeader("Authorization"); strings.HasPrefix(auth, "Bearer ") {
			tokenStr = strings.TrimPrefix(auth, "Bearer ")
		}
	}
	if tokenStr == "" {
		return "", errMissingToken
	}

	// DevMode shortcut: dev-token maps to a synthetic local user.
	if h.cfg.DevMode && tokenStr == "dev-token" {
		return "dev-user", nil
	}

	// RS256 via shared-go/jwks (preferred, requires iam /.well-known/jwks.json).
	if h.jwksFetcher != nil {
		claims, err := h.jwksFetcher.ValidateToken(tokenStr)
		if err != nil {
			return "", errInvalidToken
		}
		if claims.UserID == "" {
			return "", errMissingUserID
		}
		return claims.UserID, nil
	}

	// DevMode without JWKS: fall back to HS256 with the configured secret so
	// local integration tests can still sign tokens.
	if h.cfg.DevMode && h.cfg.JWTSecret != "" {
		return validateHS256(tokenStr, h.cfg.JWTSecret)
	}
	return "", errInvalidToken
}

// buildMessage constructs the application-layer WebSocket message envelope
// (02 §4.1): {type:"message", event, data, timestamp}.
func buildMessage(event string, data map[string]any) ([]byte, error) {
	return json.Marshal(map[string]any{
		"type":      "message",
		"event":     event,
		"data":      data,
		"timestamp": time.Now().UTC().Format(time.RFC3339),
	})
}

// mustMarshal is a best-effort JSON encoder for the Pub/Sub payload; on error
// it returns null, which the subscriber handles as empty data.
func mustMarshal(data map[string]any) []byte {
	if data == nil {
		return []byte("null")
	}
	b, err := json.Marshal(data)
	if err != nil {
		return []byte("null")
	}
	return b
}

// contextWithTimeout returns a fresh context + cancel pair with the given
// timeout. Used by HTTP handlers and the reader loop for short-lived Redis
// round-trips (presence refresh, IsOnline, Publish).
func contextWithTimeout(timeout time.Duration) (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), timeout)
}

// validateHS256 parses and validates an HS256 JWT signed with secret. Used
// only in DevMode when no JWKS endpoint is configured (local integration
// tests). Returns the user_id claim on success.
func validateHS256(tokenStr, secret string) (string, error) {
	claims := &jwks.Claims{}
	parsed, err := jwt.ParseWithClaims(tokenStr, claims, func(t *jwt.Token) (any, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("ws: unexpected signing method %v", t.Header["alg"])
		}
		return []byte(secret), nil
	})
	if err != nil {
		return "", errInvalidToken
	}
	if !parsed.Valid {
		return "", errInvalidToken
	}
	if claims.UserID == "" {
		return "", errMissingUserID
	}
	return claims.UserID, nil
}