chore(push-gateway): ai02 module updates - config, hub, ws, kafka, health, docs
This commit is contained in:
SpecialX
@@ -1,18 +1,79 @@
|
||||
// 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-Token header matched against INTERNAL_API_TOKEN
|
||||
// (ARB-015 §17.3 / president §7.2). DevMode skips the check.
|
||||
//
|
||||
// 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"
|
||||
"log"
|
||||
"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-015 §17.3 / president §7.2). Was X-Internal-Key in pre-P5 builds.
|
||||
const internalTokenHeader = "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")
|
||||
@@ -20,148 +81,349 @@ var (
|
||||
errMissingUserID = errors.New("missing user id")
|
||||
)
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true // P5 骨架，生产环境需校验 origin
|
||||
},
|
||||
}
|
||||
|
||||
// internalAPIKeyHeader 是内部 API 鉴权请求头名称
|
||||
const internalAPIKeyHeader = "X-Internal-Key"
|
||||
|
||||
// Handler 处理 WebSocket 升级与内部推送 API
|
||||
// Handler owns the WebSocket upgrade path and the internal HTTP APIs.
|
||||
type Handler struct {
|
||||
hub *hub.Hub
|
||||
jwtSecret string
|
||||
devMode bool
|
||||
internalAPIKey string
|
||||
cfg *config.Config
|
||||
redis *redisclient.Client
|
||||
jwksFetcher *jwks.Fetcher
|
||||
metrics *observability.Metrics
|
||||
upgrader websocket.Upgrader
|
||||
}
|
||||
|
||||
// NewHandler 创建 Handler 实例
|
||||
func NewHandler(h *hub.Hub, jwtSecret string, devMode bool, internalAPIKey string) *Handler {
|
||||
return &Handler{hub: h, jwtSecret: jwtSecret, devMode: devMode, internalAPIKey: internalAPIKey}
|
||||
// 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),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// HandleWebSocket 升级 HTTP 为 WebSocket，保持长连接并处理心跳
|
||||
// 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.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, pushResponse{
|
||||
Success: false,
|
||||
Error: &errBody{Code: "PUSH_UNAUTHORIZED", Message: err.Error()},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
log.Printf("[ws] upgrade failed, user=%s: %v", userID, err)
|
||||
// Upgrade already wrote the error response via c.Writer.
|
||||
observability.Logger().Warn("ws upgrade failed",
|
||||
"user_id", userID, "err", err)
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
// Apply per-connection limits.
|
||||
conn.SetReadLimit(64 * 1024) // 64KB max message size (02 §11)
|
||||
|
||||
connPtr := h.hub.Register(userID, conn)
|
||||
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)
|
||||
|
||||
// 写协程：从 send chan 读取并写入 WebSocket，避免并发写
|
||||
go func() {
|
||||
for msg := range connPtr.Outgoing() {
|
||||
if err := conn.WriteMessage(websocket.TextMessage, msg); err != nil {
|
||||
log.Printf("[ws] write failed, user=%s: %v", userID, err)
|
||||
// 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.
|
||||
func (h *Handler) PushHandler(c *gin.Context) {
|
||||
if !h.checkInternalToken(c) {
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
UserID string `json:"user_id" binding:"required"`
|
||||
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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}()
|
||||
|
||||
// 读取循环（保持连接，处理心跳）
|
||||
for {
|
||||
_, msg, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
|
||||
log.Printf("[ws] read failed, user=%s: %v", userID, err)
|
||||
}
|
||||
break
|
||||
}
|
||||
if strings.ToLower(string(msg)) == "ping" {
|
||||
connPtr.Send([]byte("pong"))
|
||||
}
|
||||
// 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})
|
||||
}
|
||||
|
||||
// PushHandler 接收来自 Msg 服务的定向推送请求
|
||||
func (h *Handler) PushHandler(c *gin.Context) {
|
||||
if !h.checkInternalAPIKey(c) {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
|
||||
"success": false,
|
||||
"error": gin.H{"code": "UNAUTHORIZED", "message": "invalid or missing internal api key"},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
UserID string `json:"userId"`
|
||||
Event string `json:"event"`
|
||||
Data map[string]any `json:"data"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": gin.H{"code": "INVALID_REQUEST", "message": err.Error()}})
|
||||
return
|
||||
}
|
||||
|
||||
message, err := buildMessage(req.Event, req.Data)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": gin.H{"code": "INVALID_PAYLOAD", "message": err.Error()}})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.hub.SendToUser(req.UserID, message); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": gin.H{"code": "PUSH_FAILED", "message": err.Error()}})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
}
|
||||
|
||||
// BroadcastHandler 接收来自 Msg 服务的广播请求
|
||||
// BroadcastHandler implements POST /internal/broadcast: push to every online
|
||||
// client across all instances.
|
||||
func (h *Handler) BroadcastHandler(c *gin.Context) {
|
||||
if !h.checkInternalAPIKey(c) {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
|
||||
"success": false,
|
||||
"error": gin.H{"code": "UNAUTHORIZED", "message": "invalid or missing internal api key"},
|
||||
})
|
||||
if !h.checkInternalToken(c) {
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
Event string `json:"event"`
|
||||
Data map[string]any `json:"data"`
|
||||
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.JSON(http.StatusBadRequest, gin.H{"success": false, "error": gin.H{"code": "INVALID_REQUEST", "message": err.Error()}})
|
||||
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.JSON(http.StatusBadRequest, gin.H{"success": false, "error": gin.H{"code": "INVALID_PAYLOAD", "message": err.Error()}})
|
||||
c.AbortWithStatusJSON(http.StatusBadRequest, pushResponse{
|
||||
Success: false,
|
||||
Error: &errBody{Code: "PUSH_INVALID_PAYLOAD", Message: err.Error()},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
h.hub.Broadcast(message)
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
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})
|
||||
}
|
||||
|
||||
// checkInternalAPIKey 校验内部 API 请求头 X-Internal-Key。
|
||||
// DevMode 下跳过校验；非 DevMode 下要求 INTERNAL_API_KEY 已配置且与请求头匹配。
|
||||
func (h *Handler) checkInternalAPIKey(c *gin.Context) bool {
|
||||
if h.devMode {
|
||||
// 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 X-Internal-Token header. 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.internalAPIKey == "" {
|
||||
log.Println("[ws] internal api key not configured, rejecting request")
|
||||
if h.cfg.InternalAPIToken == "" {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, pushResponse{
|
||||
Success: false,
|
||||
Error: &errBody{Code: "PUSH_UNAUTHORIZED", Message: "internal token not configured"},
|
||||
})
|
||||
return false
|
||||
}
|
||||
return c.GetHeader(internalAPIKeyHeader) == h.internalAPIKey
|
||||
if c.GetHeader(internalTokenHeader) != 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 校验 WebSocket 连接的 JWT 或 dev token
|
||||
// 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 == "" {
|
||||
@@ -173,37 +435,81 @@ func (h *Handler) authenticate(c *gin.Context) (string, error) {
|
||||
return "", errMissingToken
|
||||
}
|
||||
|
||||
// DEV_MODE 下接受 dev-token，便于本地联调
|
||||
if h.devMode && tokenStr == "dev-token" {
|
||||
// DevMode shortcut: dev-token maps to a synthetic local user.
|
||||
if h.cfg.DevMode && tokenStr == "dev-token" {
|
||||
return "dev-user", nil
|
||||
}
|
||||
|
||||
token, err := jwt.Parse(tokenStr, func(t *jwt.Token) (any, error) {
|
||||
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
|
||||
return nil, jwt.ErrSignatureInvalid
|
||||
// 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
|
||||
}
|
||||
return []byte(h.jwtSecret), nil
|
||||
})
|
||||
if err != nil || !token.Valid {
|
||||
return "", errInvalidToken
|
||||
if claims.UserID == "" {
|
||||
return "", errMissingUserID
|
||||
}
|
||||
return claims.UserID, nil
|
||||
}
|
||||
|
||||
claims, ok := token.Claims.(jwt.MapClaims)
|
||||
if !ok {
|
||||
return "", errInvalidClaims
|
||||
// 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)
|
||||
}
|
||||
|
||||
userID, ok := claims["sub"].(string)
|
||||
if !ok {
|
||||
return "", errMissingUserID
|
||||
}
|
||||
return userID, nil
|
||||
return "", errInvalidToken
|
||||
}
|
||||
|
||||
// buildMessage 构造推送 JSON 消息体
|
||||
// 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{
|
||||
"event": event,
|
||||
"data": data,
|
||||
"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
|
||||
}
|
||||
|
||||
196
services/push-gateway/internal/ws/handler_test.go
Normal file
196
services/push-gateway/internal/ws/handler_test.go
Normal file
@@ -0,0 +1,196 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/edu-cloud/shared-go/jwks"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
// TestBuildMessage verifies the WebSocket message envelope format (02 §4.1).
|
||||
func TestBuildMessage(t *testing.T) {
|
||||
data := map[string]any{"title": "hello", "body": "world"}
|
||||
msg, err := buildMessage("notification.created", data)
|
||||
if err != nil {
|
||||
t.Fatalf("buildMessage: %v", err)
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(msg, &payload); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if payload["type"] != "message" {
|
||||
t.Errorf("type = %v, want message", payload["type"])
|
||||
}
|
||||
if payload["event"] != "notification.created" {
|
||||
t.Errorf("event = %v, want notification.created", payload["event"])
|
||||
}
|
||||
if payload["timestamp"] == nil {
|
||||
t.Error("timestamp missing")
|
||||
}
|
||||
dataField, ok := payload["data"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("data should be a map, got %T", payload["data"])
|
||||
}
|
||||
if dataField["title"] != "hello" {
|
||||
t.Errorf("data.title = %v, want hello", dataField["title"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildMessageNilData verifies nil data is handled gracefully.
|
||||
func TestBuildMessageNilData(t *testing.T) {
|
||||
msg, err := buildMessage("test.event", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("buildMessage: %v", err)
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(msg, &payload); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if payload["data"] != nil {
|
||||
t.Errorf("data = %v, want nil", payload["data"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestMustMarshal verifies best-effort JSON encoding.
|
||||
func TestMustMarshal(t *testing.T) {
|
||||
// nil -> "null"
|
||||
if got := mustMarshal(nil); string(got) != "null" {
|
||||
t.Errorf("mustMarshal(nil) = %q, want null", got)
|
||||
}
|
||||
// valid map -> JSON
|
||||
got := mustMarshal(map[string]any{"a": 1})
|
||||
var m map[string]any
|
||||
if err := json.Unmarshal(got, &m); err != nil {
|
||||
t.Fatalf("mustMarshal result unmarshal: %v", err)
|
||||
}
|
||||
if m["a"].(float64) != 1 {
|
||||
t.Errorf("mustMarshal.a = %v, want 1", m["a"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildOriginChecker verifies origin whitelist logic.
|
||||
func TestBuildOriginChecker(t *testing.T) {
|
||||
allowed := map[string]struct{}{
|
||||
"http://localhost:3000": {},
|
||||
"https://app.example.com": {},
|
||||
}
|
||||
|
||||
// Non-devMode: strict whitelist.
|
||||
checker := buildOriginChecker(allowed, false)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
origin string
|
||||
want bool
|
||||
}{
|
||||
{"empty origin (non-browser)", "", true},
|
||||
{"allowed origin", "http://localhost:3000", true},
|
||||
{"another allowed", "https://app.example.com", true},
|
||||
{"disallowed origin", "http://evil.com", false},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
req := httptest.NewRequest("GET", "/ws", nil)
|
||||
if tc.origin != "" {
|
||||
req.Header.Set("Origin", tc.origin)
|
||||
}
|
||||
if got := checker(req); got != tc.want {
|
||||
t.Errorf("checker(origin=%q) = %v, want %v", tc.origin, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// DevMode with empty whitelist: allow all.
|
||||
devChecker := buildOriginChecker(map[string]struct{}{}, true)
|
||||
req := httptest.NewRequest("GET", "/ws", nil)
|
||||
req.Header.Set("Origin", "http://anything.com")
|
||||
if !devChecker(req) {
|
||||
t.Error("devMode empty whitelist should allow all origins")
|
||||
}
|
||||
}
|
||||
|
||||
// TestValidateHS256 verifies HS256 JWT validation in DevMode fallback.
|
||||
func TestValidateHS256(t *testing.T) {
|
||||
secret := "test-secret"
|
||||
|
||||
// Valid token with user_id claim.
|
||||
claims := jwks.Claims{
|
||||
UserID: "user-123",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenStr, err := token.SignedString([]byte(secret))
|
||||
if err != nil {
|
||||
t.Fatalf("SignedString: %v", err)
|
||||
}
|
||||
|
||||
userID, err := validateHS256(tokenStr, secret)
|
||||
if err != nil {
|
||||
t.Fatalf("validateHS256: %v", err)
|
||||
}
|
||||
if userID != "user-123" {
|
||||
t.Errorf("userID = %q, want user-123", userID)
|
||||
}
|
||||
|
||||
// Wrong secret -> error.
|
||||
_, err = validateHS256(tokenStr, "wrong-secret")
|
||||
if err == nil {
|
||||
t.Error("validateHS256 with wrong secret should fail")
|
||||
}
|
||||
|
||||
// Expired token -> error.
|
||||
expiredClaims := jwks.Claims{
|
||||
UserID: "user-exp",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(-1 * time.Hour)),
|
||||
},
|
||||
}
|
||||
expiredToken := jwt.NewWithClaims(jwt.SigningMethodHS256, expiredClaims)
|
||||
expiredStr, _ := expiredToken.SignedString([]byte(secret))
|
||||
_, err = validateHS256(expiredStr, secret)
|
||||
if err == nil {
|
||||
t.Error("validateHS256 with expired token should fail")
|
||||
}
|
||||
|
||||
// Missing user_id -> error.
|
||||
noUserClaims := jwks.Claims{
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
},
|
||||
}
|
||||
noUserToken := jwt.NewWithClaims(jwt.SigningMethodHS256, noUserClaims)
|
||||
noUserStr, _ := noUserToken.SignedString([]byte(secret))
|
||||
_, err = validateHS256(noUserStr, secret)
|
||||
if err != errMissingUserID {
|
||||
t.Errorf("validateHS256 without user_id err = %v, want errMissingUserID", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestContextWithTimeout verifies the helper returns a working context.
|
||||
func TestContextWithTimeout(t *testing.T) {
|
||||
ctx, cancel := contextWithTimeout(100 * time.Millisecond)
|
||||
defer cancel()
|
||||
if ctx == nil {
|
||||
t.Fatal("context is nil")
|
||||
}
|
||||
// Context should not be done immediately.
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
t.Fatal("context done immediately")
|
||||
default:
|
||||
}
|
||||
// Wait for timeout.
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// expected
|
||||
default:
|
||||
t.Error("context not done after timeout")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user