chore(push-gateway): ai02 module updates - config, hub, ws, kafka, health, docs
This commit is contained in:
SpecialX
120
services/push-gateway/internal/config/config_test.go
Normal file
120
services/push-gateway/internal/config/config_test.go
Normal file
@@ -0,0 +1,120 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestLoadDevMode verifies DevMode defaults and that required vars are optional.
|
||||
func TestLoadDevMode(t *testing.T) {
|
||||
t.Setenv("DEV_MODE", "true")
|
||||
t.Setenv("JWT_SECRET", "")
|
||||
t.Setenv("INTERNAL_API_TOKEN", "")
|
||||
|
||||
cfg := Load()
|
||||
|
||||
if !cfg.DevMode {
|
||||
t.Error("DevMode = false, want true")
|
||||
}
|
||||
if cfg.JWTSecret != devJWTSecret {
|
||||
t.Errorf("JWTSecret = %q, want dev default %q", cfg.JWTSecret, devJWTSecret)
|
||||
}
|
||||
if cfg.InternalAPIToken != "" {
|
||||
t.Errorf("InternalAPIToken = %q, want empty in DevMode", cfg.InternalAPIToken)
|
||||
}
|
||||
if cfg.Port != "8081" {
|
||||
t.Errorf("Port = %q, want 8081", cfg.Port)
|
||||
}
|
||||
if cfg.MaxConnsPerUser != 5 {
|
||||
t.Errorf("MaxConnsPerUser = %d, want 5", cfg.MaxConnsPerUser)
|
||||
}
|
||||
if cfg.HeartbeatInterval != 30 {
|
||||
t.Errorf("HeartbeatInterval = %d, want 30", cfg.HeartbeatInterval)
|
||||
}
|
||||
if cfg.KafkaNotificationTopic != "edu.notification.requested" {
|
||||
t.Errorf("KafkaNotificationTopic = %q, want edu.notification.requested", cfg.KafkaNotificationTopic)
|
||||
}
|
||||
if cfg.KafkaConsumerGroup != "push-gateway" {
|
||||
t.Errorf("KafkaConsumerGroup = %q, want push-gateway", cfg.KafkaConsumerGroup)
|
||||
}
|
||||
if cfg.InstanceID == "" {
|
||||
t.Error("InstanceID should default to hostname, got empty")
|
||||
}
|
||||
}
|
||||
|
||||
// TestLoadProdMode verifies production requires JWT_SECRET + INTERNAL_API_TOKEN.
|
||||
func TestLoadProdMode(t *testing.T) {
|
||||
// Ensure the env vars are unset so Load fails. We can't call os.Exit in
|
||||
// a test, so we verify the required behavior by setting DEV_MODE=false
|
||||
// and checking that Load triggers log.Fatal. Since log.Fatal calls
|
||||
// os.Exit, we instead test the happy path: all required vars set.
|
||||
t.Setenv("DEV_MODE", "false")
|
||||
t.Setenv("JWT_SECRET", "prod-secret")
|
||||
t.Setenv("INTERNAL_API_TOKEN", "prod-token")
|
||||
|
||||
cfg := Load()
|
||||
|
||||
if cfg.DevMode {
|
||||
t.Error("DevMode = true, want false")
|
||||
}
|
||||
if cfg.JWTSecret != "prod-secret" {
|
||||
t.Errorf("JWTSecret = %q, want prod-secret", cfg.JWTSecret)
|
||||
}
|
||||
if cfg.InternalAPIToken != "prod-token" {
|
||||
t.Errorf("InternalAPIToken = %q, want prod-token", cfg.InternalAPIToken)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseOrigins verifies comma-separated origin parsing.
|
||||
func TestParseOrigins(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
raw string
|
||||
want []string
|
||||
}{
|
||||
{"empty", "", nil},
|
||||
{"single", "http://localhost:3000", []string{"http://localhost:3000"}},
|
||||
{"multi", "http://a.com,http://b.com,http://c.com", []string{"http://a.com", "http://b.com", "http://c.com"}},
|
||||
{"with-spaces", " http://a.com , http://b.com ", []string{"http://a.com", "http://b.com"}},
|
||||
{"trailing-comma", "http://a.com,", []string{"http://a.com"}},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := parseOrigins(tc.raw)
|
||||
if len(got) != len(tc.want) {
|
||||
t.Errorf("parseOrigins(%q) = %v, want %v", tc.raw, got, tc.want)
|
||||
return
|
||||
}
|
||||
for i := range got {
|
||||
if got[i] != tc.want[i] {
|
||||
t.Errorf("parseOrigins(%q)[%d] = %q, want %q", tc.raw, i, got[i], tc.want[i])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseBrokers verifies comma-separated broker parsing.
|
||||
func TestParseBrokers(t *testing.T) {
|
||||
got := parseBrokers("kafka1:9092,kafka2:9092,kafka3:9092")
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("parseBrokers returned %d items, want 3", len(got))
|
||||
}
|
||||
if got[0] != "kafka1:9092" || got[1] != "kafka2:9092" || got[2] != "kafka3:9092" {
|
||||
t.Errorf("parseBrokers = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGenerateInstanceID verifies fallback to hostname.
|
||||
func TestGenerateInstanceID(t *testing.T) {
|
||||
id := generateInstanceID()
|
||||
if id == "" {
|
||||
t.Error("generateInstanceID returned empty string")
|
||||
}
|
||||
host, err := os.Hostname()
|
||||
if err == nil && host != "" {
|
||||
if id != host {
|
||||
t.Errorf("generateInstanceID = %q, want hostname %q", id, host)
|
||||
}
|
||||
}
|
||||
}
|
||||
186
services/push-gateway/internal/health/readyz.go
Normal file
186
services/push-gateway/internal/health/readyz.go
Normal file
@@ -0,0 +1,186 @@
|
||||
// Package health implements the /healthz (liveness) and /readyz (readiness)
|
||||
// endpoints for push-gateway.
|
||||
//
|
||||
// /healthz is a trivial liveness probe: it returns 200 + {status:ok} as long
|
||||
// as the process is running and the Gin router is serving.
|
||||
//
|
||||
// /readyz reports readiness based on downstream dependency health. Per
|
||||
// ARB-015 §17.4 / ISSUE-058 the readiness probe uses SOFT FAILURE semantics:
|
||||
// when Redis or Kafka is unavailable the endpoint still returns HTTP 200 but
|
||||
// carries `degraded: true` plus a `dependencies` block describing which
|
||||
// component failed. This prevents Kubernetes from evicting the pod when a
|
||||
// transient dependency blip occurs, at the cost of accepting some degraded
|
||||
// behavior (no cross-instance fanout, no Kafka consumption) during the blip.
|
||||
// Only the Hub being in a shutting-down state returns a non-200 (503).
|
||||
//
|
||||
// The hard-failure case is limited to:
|
||||
// - Hub.closing == true (process is shutting down) -> 503
|
||||
// - Internal misconfiguration (both Redis and Kafka missing in non-DevMode)
|
||||
// -> still 200 + degraded, since the pod can still serve local WebSocket
|
||||
// traffic.
|
||||
package health
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/edu-cloud/push-gateway/internal/hub"
|
||||
"github.com/edu-cloud/push-gateway/internal/kafkaconsumer"
|
||||
"github.com/edu-cloud/push-gateway/internal/redisclient"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// probeTimeout bounds each dependency check so a hung dependency cannot stall
|
||||
// the readiness probe. Kubernetes typically expects /readyz to respond within
|
||||
// 5s; we use 1s per probe to leave headroom.
|
||||
const probeTimeout = 1 * time.Second
|
||||
|
||||
// dependencyStatus is the per-component health snapshot included in the
|
||||
// /readyz response body.
|
||||
type dependencyStatus struct {
|
||||
Ok bool `json:"ok"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Latency int64 `json:"latency_ms,omitempty"`
|
||||
}
|
||||
|
||||
// readyzResponse is the ActionState-shaped envelope returned by /readyz.
|
||||
// `degraded` is true when at least one non-critical dependency is unhealthy.
|
||||
type readyzResponse struct {
|
||||
Status string `json:"status"`
|
||||
Service string `json:"service"`
|
||||
InstanceID string `json:"instance_id"`
|
||||
Degraded bool `json:"degraded"`
|
||||
Connections int `json:"connections"`
|
||||
Users int `json:"users"`
|
||||
Dependencies map[string]*dependencyStatus `json:"dependencies"`
|
||||
}
|
||||
|
||||
// Healthz returns a trivial liveness handler: 200 + {status:ok, service}.
|
||||
// Liveness never depends on downstream components — if the process can serve
|
||||
// the request, it is alive.
|
||||
func Healthz(service string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"status": "ok",
|
||||
"service": service,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Readyzer builds the /readyz handler. It probes Redis (PING) and Kafka
|
||||
// (reader lag / connectivity) and reports degraded state per ARB-015 §17.4.
|
||||
//
|
||||
// The Hub is used to report local connection counts and to detect the
|
||||
// shutting-down state (which triggers a hard 503). The redisClient may be nil
|
||||
// in DevMode; the kafkaConsumer may be nil when KAFKA_BROKERS is unset. Both
|
||||
// nil cases are reported as degraded rather than failing the probe.
|
||||
type Readyzer struct {
|
||||
hub *hub.Hub
|
||||
redis *redisclient.Client
|
||||
kafka *kafkaconsumer.Consumer
|
||||
instance string
|
||||
service string
|
||||
}
|
||||
|
||||
// NewReadyzer constructs a Readyzer. redis and kafka may be nil; the resulting
|
||||
// probe will mark the missing dependency as degraded.
|
||||
func NewReadyzer(h *hub.Hub, r *redisclient.Client, k *kafkaconsumer.Consumer, service, instanceID string) *Readyzer {
|
||||
return &Readyzer{
|
||||
hub: h,
|
||||
redis: r,
|
||||
kafka: k,
|
||||
instance: instanceID,
|
||||
service: service,
|
||||
}
|
||||
}
|
||||
|
||||
// Handler is the gin.HandlerFunc for GET /readyz.
|
||||
func (rz *Readyzer) Handler(c *gin.Context) {
|
||||
// Hard failure: Hub is shutting down — reject new traffic.
|
||||
// (Hub.CloseAll sets closing=true; we treat this as 503 so the load
|
||||
// balancer stops sending WebSocket upgrades during drain.)
|
||||
if rz.hub.IsClosing() {
|
||||
c.JSON(http.StatusServiceUnavailable, readyzResponse{
|
||||
Status: "shutting_down",
|
||||
Service: rz.service,
|
||||
InstanceID: rz.instance,
|
||||
Degraded: true,
|
||||
Connections: rz.hub.ActiveConnections(),
|
||||
Users: rz.hub.UserCount(),
|
||||
Dependencies: map[string]*dependencyStatus{
|
||||
"hub": {Ok: false, Error: "hub is closing"},
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
deps := make(map[string]*dependencyStatus, 2)
|
||||
degraded := false
|
||||
|
||||
// Redis probe (soft failure).
|
||||
if rz.redis == nil {
|
||||
deps["redis"] = &dependencyStatus{Ok: false, Error: "not configured"}
|
||||
degraded = true
|
||||
} else {
|
||||
deps["redis"] = probeRedis(rz.redis)
|
||||
if !deps["redis"].Ok {
|
||||
degraded = true
|
||||
}
|
||||
}
|
||||
|
||||
// Kafka probe (soft failure).
|
||||
if rz.kafka == nil {
|
||||
deps["kafka"] = &dependencyStatus{Ok: false, Error: "not configured"}
|
||||
degraded = true
|
||||
} else {
|
||||
deps["kafka"] = probeKafka(rz.kafka)
|
||||
if !deps["kafka"].Ok {
|
||||
degraded = true
|
||||
}
|
||||
}
|
||||
|
||||
// Always 200 (unless shutting down) per ISSUE-058/006 soft-failure rule.
|
||||
c.JSON(http.StatusOK, readyzResponse{
|
||||
Status: statusText(degraded),
|
||||
Service: rz.service,
|
||||
InstanceID: rz.instance,
|
||||
Degraded: degraded,
|
||||
Connections: rz.hub.ActiveConnections(),
|
||||
Users: rz.hub.UserCount(),
|
||||
Dependencies: deps,
|
||||
})
|
||||
}
|
||||
|
||||
// statusText returns "ok" when healthy, "degraded" when at least one
|
||||
// dependency is unhealthy.
|
||||
func statusText(degraded bool) string {
|
||||
if degraded {
|
||||
return "degraded"
|
||||
}
|
||||
return "ok"
|
||||
}
|
||||
|
||||
// probeRedis issues a PING with a short timeout and reports latency.
|
||||
func probeRedis(r *redisclient.Client) *dependencyStatus {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), probeTimeout)
|
||||
defer cancel()
|
||||
start := time.Now()
|
||||
if err := r.Ping(ctx); err != nil {
|
||||
return &dependencyStatus{Ok: false, Error: err.Error()}
|
||||
}
|
||||
return &dependencyStatus{Ok: true, Latency: time.Since(start).Milliseconds()}
|
||||
}
|
||||
|
||||
// probeKafka checks that the consumer reader is still reachable. We use a
|
||||
// lightweight Lag() call (segmentio/kafka-go Client API); on failure the
|
||||
// consumer is marked degraded. Note: a degraded Kafka does NOT block WebSocket
|
||||
// traffic — it only pauses notification consumption until recovery.
|
||||
func probeKafka(k *kafkaconsumer.Consumer) *dependencyStatus {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), probeTimeout)
|
||||
defer cancel()
|
||||
if err := k.HealthCheck(ctx); err != nil {
|
||||
return &dependencyStatus{Ok: false, Error: err.Error()}
|
||||
}
|
||||
return &dependencyStatus{Ok: true}
|
||||
}
|
||||
@@ -172,6 +172,15 @@ func (h *Hub) HasUser(userID string) bool {
|
||||
return ok && len(conns) > 0
|
||||
}
|
||||
|
||||
// IsClosing reports whether CloseAll has been invoked. The /readyz probe uses
|
||||
// this to return 503 during graceful shutdown so the load balancer stops
|
||||
// sending new WebSocket upgrades while existing connections drain.
|
||||
func (h *Hub) IsClosing() bool {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
return h.closing
|
||||
}
|
||||
|
||||
// Register adds a new connection for userID. Returns ErrTooManyConnections
|
||||
// when the per-user cap is exceeded and ErrHubClosing when the Hub is shutting
|
||||
// down. The returned Connection must be Unregistered by the caller on close.
|
||||
@@ -262,7 +271,13 @@ func (h *Hub) SendToUser(userID string, message []byte) int {
|
||||
// Cross-instance broadcast is handled separately via Redis Pub/Sub.
|
||||
func (h *Hub) Broadcast(message []byte) int {
|
||||
h.mu.RLock()
|
||||
snapshot := make([]*Connection, 0, h.ActiveConnections())
|
||||
// Count inline: calling ActiveConnections() here would re-enter RLock
|
||||
// while we already hold it, which can deadlock if a writer is waiting.
|
||||
total := 0
|
||||
for _, conns := range h.clients {
|
||||
total += len(conns)
|
||||
}
|
||||
snapshot := make([]*Connection, 0, total)
|
||||
for _, conns := range h.clients {
|
||||
for _, c := range conns {
|
||||
snapshot = append(snapshot, c)
|
||||
@@ -285,7 +300,13 @@ func (h *Hub) Broadcast(message []byte) int {
|
||||
func (h *Hub) CloseAll() {
|
||||
h.mu.Lock()
|
||||
h.closing = true
|
||||
all := make([]*Connection, 0, h.ActiveConnections())
|
||||
// Count inline instead of calling ActiveConnections() to avoid reentrant
|
||||
// RLock while holding the write lock (Go's sync.RWMutex is non-reentrant).
|
||||
total := 0
|
||||
for _, conns := range h.clients {
|
||||
total += len(conns)
|
||||
}
|
||||
all := make([]*Connection, 0, total)
|
||||
for _, conns := range h.clients {
|
||||
for _, c := range conns {
|
||||
all = append(all, c)
|
||||
|
||||
238
services/push-gateway/internal/hub/hub_test.go
Normal file
238
services/push-gateway/internal/hub/hub_test.go
Normal file
@@ -0,0 +1,238 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// newTestConn returns a real *websocket.Conn (client side) backed by a local
|
||||
// httptest server. The server-side handler blocks on a done channel until the
|
||||
// test ends; cleanup closes the client conn first, then the server, avoiding
|
||||
// the httptest.Server.Close() hang that occurs when the handler is blocked on
|
||||
// ReadMessage.
|
||||
func newTestConn(t *testing.T) *websocket.Conn {
|
||||
t.Helper()
|
||||
done := make(chan struct{})
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
up := websocket.Upgrader{}
|
||||
c, err := up.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
t.Errorf("upgrade: %v", err)
|
||||
return
|
||||
}
|
||||
<-done
|
||||
_ = c.Close()
|
||||
}))
|
||||
cli, _, err := websocket.DefaultDialer.Dial(strings.Replace(srv.URL, "http:", "ws:", 1), nil)
|
||||
if err != nil {
|
||||
close(done)
|
||||
srv.Close()
|
||||
t.Fatalf("dial: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
close(done)
|
||||
_ = cli.Close()
|
||||
srv.Close()
|
||||
})
|
||||
return cli
|
||||
}
|
||||
|
||||
// TestRegisterAndUnregister verifies basic Hub register/unregister bookkeeping.
|
||||
func TestRegisterAndUnregister(t *testing.T) {
|
||||
h := NewHub(5)
|
||||
conn := newTestConn(t)
|
||||
|
||||
c, err := h.Register("user-1", conn)
|
||||
if err != nil {
|
||||
t.Fatalf("Register: %v", err)
|
||||
}
|
||||
if c.UserID != "user-1" {
|
||||
t.Errorf("UserID = %q, want user-1", c.UserID)
|
||||
}
|
||||
if c.ConnID == "" {
|
||||
t.Error("ConnID should be non-empty")
|
||||
}
|
||||
if h.ActiveConnections() != 1 {
|
||||
t.Errorf("ActiveConnections = %d, want 1", h.ActiveConnections())
|
||||
}
|
||||
if h.UserCount() != 1 {
|
||||
t.Errorf("UserCount = %d, want 1", h.UserCount())
|
||||
}
|
||||
if !h.HasUser("user-1") {
|
||||
t.Error("HasUser(user-1) = false, want true")
|
||||
}
|
||||
|
||||
h.Unregister(c)
|
||||
if h.ActiveConnections() != 0 {
|
||||
t.Errorf("after Unregister ActiveConnections = %d, want 0", h.ActiveConnections())
|
||||
}
|
||||
if h.HasUser("user-1") {
|
||||
t.Error("HasUser(user-1) = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegisterTooManyConnections verifies the per-user cap is enforced.
|
||||
func TestRegisterTooManyConnections(t *testing.T) {
|
||||
h := NewHub(2)
|
||||
for i := 0; i < 2; i++ {
|
||||
conn := newTestConn(t)
|
||||
if _, err := h.Register("user-cap", conn); err != nil {
|
||||
t.Fatalf("Register #%d: %v", i, err)
|
||||
}
|
||||
}
|
||||
conn := newTestConn(t)
|
||||
_, err := h.Register("user-cap", conn)
|
||||
if err != ErrTooManyConnections {
|
||||
t.Errorf("Register over cap err = %v, want ErrTooManyConnections", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegisterWhenClosing verifies that CloseAll rejects further Register.
|
||||
func TestRegisterWhenClosing(t *testing.T) {
|
||||
h := NewHub(5)
|
||||
h.CloseAll()
|
||||
conn := newTestConn(t)
|
||||
_, err := h.Register("user-x", conn)
|
||||
if err != ErrHubClosing {
|
||||
t.Errorf("Register after CloseAll err = %v, want ErrHubClosing", err)
|
||||
}
|
||||
if !h.IsClosing() {
|
||||
t.Error("IsClosing = false, want true")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSendToUser verifies message delivery to a single user.
|
||||
func TestSendToUser(t *testing.T) {
|
||||
h := NewHub(5)
|
||||
conn := newTestConn(t)
|
||||
c, err := h.Register("user-send", conn)
|
||||
if err != nil {
|
||||
t.Fatalf("Register: %v", err)
|
||||
}
|
||||
defer h.Unregister(c)
|
||||
|
||||
msg := []byte(`{"type":"message","event":"test"}`)
|
||||
delivered := h.SendToUser("user-send", msg)
|
||||
if delivered != 1 {
|
||||
t.Errorf("SendToUser delivered = %d, want 1", delivered)
|
||||
}
|
||||
|
||||
// Unknown user: 0 delivered.
|
||||
delivered = h.SendToUser("user-absent", msg)
|
||||
if delivered != 0 {
|
||||
t.Errorf("SendToUser absent delivered = %d, want 0", delivered)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBroadcast verifies message delivery to all connections.
|
||||
func TestBroadcast(t *testing.T) {
|
||||
h := NewHub(5)
|
||||
conn1 := newTestConn(t)
|
||||
conn2 := newTestConn(t)
|
||||
c1, _ := h.Register("u1", conn1)
|
||||
c2, _ := h.Register("u2", conn2)
|
||||
defer h.Unregister(c1)
|
||||
defer h.Unregister(c2)
|
||||
|
||||
msg := []byte(`{"type":"message","event":"broadcast"}`)
|
||||
reached := h.Broadcast(msg)
|
||||
if reached != 2 {
|
||||
t.Errorf("Broadcast reached = %d, want 2", reached)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPresenceHooks verifies that Register/Unregister call the hooks.
|
||||
func TestPresenceHooks(t *testing.T) {
|
||||
h := NewHub(5)
|
||||
var mu sync.Mutex
|
||||
registered := []string{}
|
||||
unregistered := []string{}
|
||||
h.SetPresenceHooks(
|
||||
func(userID, connID string) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
registered = append(registered, userID)
|
||||
},
|
||||
func(userID, connID string) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
unregistered = append(unregistered, userID)
|
||||
},
|
||||
)
|
||||
|
||||
conn := newTestConn(t)
|
||||
c, _ := h.Register("hook-user", conn)
|
||||
h.Unregister(c)
|
||||
|
||||
if len(registered) != 1 || registered[0] != "hook-user" {
|
||||
t.Errorf("registered = %v, want [hook-user]", registered)
|
||||
}
|
||||
if len(unregistered) != 1 || unregistered[0] != "hook-user" {
|
||||
t.Errorf("unregistered = %v, want [hook-user]", unregistered)
|
||||
}
|
||||
}
|
||||
|
||||
// TestForEachUser verifies iteration over online users (used by ISSUE-058 rebuild).
|
||||
func TestForEachUser(t *testing.T) {
|
||||
h := NewHub(5)
|
||||
conn1 := newTestConn(t)
|
||||
conn2 := newTestConn(t)
|
||||
c1, _ := h.Register("u-a", conn1)
|
||||
c2, _ := h.Register("u-b", conn2)
|
||||
defer h.Unregister(c1)
|
||||
defer h.Unregister(c2)
|
||||
|
||||
seen := map[string]bool{}
|
||||
h.ForEachUser(func(userID string) {
|
||||
seen[userID] = true
|
||||
})
|
||||
if len(seen) != 2 || !seen["u-a"] || !seen["u-b"] {
|
||||
t.Errorf("ForEachUser seen = %v, want {u-a, u-b}", seen)
|
||||
}
|
||||
}
|
||||
|
||||
// TestConnectionSendOnClosed verifies Send returns false after Close.
|
||||
func TestConnectionSendOnClosed(t *testing.T) {
|
||||
h := NewHub(5)
|
||||
conn := newTestConn(t)
|
||||
c, _ := h.Register("u-close", conn)
|
||||
c.Close()
|
||||
if ok := c.Send([]byte("x")); ok {
|
||||
t.Error("Send after Close = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
// TestConnectionSendChannelFull verifies Send returns false when channel is full.
|
||||
func TestConnectionSendChannelFull(t *testing.T) {
|
||||
h := NewHub(5)
|
||||
conn := newTestConn(t)
|
||||
c, _ := h.Register("u-full", conn)
|
||||
defer h.Unregister(c)
|
||||
// Fill the send channel (cap 64).
|
||||
for i := 0; i < sendBufferSize; i++ {
|
||||
if !c.Send([]byte("x")) {
|
||||
t.Fatalf("Send #%d returned false unexpectedly", i)
|
||||
}
|
||||
}
|
||||
// Next Send should drop.
|
||||
if ok := c.Send([]byte("overflow")); ok {
|
||||
t.Error("Send on full channel = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewHubDefaultMaxConns verifies NewHub applies a default when maxConns <= 0.
|
||||
func TestNewHubDefaultMaxConns(t *testing.T) {
|
||||
h := NewHub(0)
|
||||
if h.maxConnsPerUser != 5 {
|
||||
t.Errorf("NewHub(0) maxConnsPerUser = %d, want default 5", h.maxConnsPerUser)
|
||||
}
|
||||
h2 := NewHub(-1)
|
||||
if h2.maxConnsPerUser != 5 {
|
||||
t.Errorf("NewHub(-1) maxConnsPerUser = %d, want default 5", h2.maxConnsPerUser)
|
||||
}
|
||||
}
|
||||
@@ -120,6 +120,39 @@ func (c *Consumer) Close() error {
|
||||
return c.reader.Close()
|
||||
}
|
||||
|
||||
// HealthCheck verifies that the Kafka broker is reachable. Used by /readyz as
|
||||
// a soft-failure probe: a failing check marks the consumer as degraded but
|
||||
// does not return 503 (per ARB-015 §17.4 / ISSUE-058).
|
||||
//
|
||||
// The check issues a metadata request for the consumer topic; on success the
|
||||
// broker is considered reachable. This is cheaper than a Lag() call which
|
||||
// requires partition assignment to have completed.
|
||||
func (c *Consumer) HealthCheck(ctx context.Context) error {
|
||||
conn, err := kafka.DialContext(ctx, "tcp", c.reader.Config().Brokers[0])
|
||||
if err != nil {
|
||||
return fmt.Errorf("kafka: dial broker: %w", err)
|
||||
}
|
||||
defer conn.Close()
|
||||
partitions, err := conn.ReadPartitions(c.topic)
|
||||
if err != nil {
|
||||
return fmt.Errorf("kafka: read partitions: %w", err)
|
||||
}
|
||||
if len(partitions) == 0 {
|
||||
return fmt.Errorf("kafka: topic %q has no partitions", c.topic)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Topic returns the configured topic name. Exposed for diagnostics.
|
||||
func (c *Consumer) Topic() string {
|
||||
return c.topic
|
||||
}
|
||||
|
||||
// GroupID returns the configured consumer group name.
|
||||
func (c *Consumer) GroupID() string {
|
||||
return c.groupID
|
||||
}
|
||||
|
||||
// processMessage handles a single Kafka message with retry and idempotency.
|
||||
func (c *Consumer) processMessage(ctx context.Context, msg kafka.Message) {
|
||||
var event NotificationRequested
|
||||
|
||||
@@ -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