feat: 修复 WebSocket 越权与 AI 流 SSRF

server-go/internal/ws/gateway_test.go

@@ -0,0 +1,120 @@
+package ws
+
+import (
+	"errors"
+	"testing"
+
+	"silk-server-go/internal/service"
+)
+
+type fakeAuthorizer struct {
+	allowed map[string]bool
+}
+
+func (f fakeAuthorizer) CanReadDevice(userID, deviceKey string) (bool, error) {
+	if f.allowed == nil {
+		return false, nil
+	}
+	return f.allowed[userID+"|"+deviceKey], nil
+}
+
+func newTestHub() *Hub {
+	return NewHub(
+		"test-secret",
+		fakeAuthorizer{allowed: map[string]bool{"user-1|device-a": true}},
+		[]string{"http://localhost:5174"},
+	)
+}
+
+func TestWebSocketOriginRejectsUnknownOrigin(t *testing.T) {
+	hub := newTestHub()
+	if hub.originAllowed("http://evil.example") {
+		t.Fatal("unknown origin should be rejected")
+	}
+	if !hub.originAllowed("http://localhost:5174") {
+		t.Fatal("configured origin should be allowed")
+	}
+}
+
+func TestWebSocketTicketIsOneTimeAndExpires(t *testing.T) {
+	hub := newTestHub()
+	ticket := hub.IssueTicket("user-1", "admin", "User One")
+	if ticket == "" {
+		t.Fatal("expected non-empty ticket")
+	}
+
+	info, ok := hub.consumeTicket(ticket)
+	if !ok || info.userID != "user-1" {
+		t.Fatalf("first ticket consume failed: ok=%v info=%+v", ok, info)
+	}
+	if _, ok := hub.consumeTicket(ticket); ok {
+		t.Fatal("ticket should be single-use")
+	}
+}
+
+func TestWebSocketRejectsLongLivedJWTQuery(t *testing.T) {
+	hub := newTestHub()
+	if hub.ticketFromQuery("") != "" {
+		t.Fatal("missing ticket should not be accepted")
+	}
+	if hub.ticketFromQuery("token=eyJhbGciOiJIUzI1NiJ9.abc") != "" {
+		t.Fatal("long-lived token query must not be accepted")
+	}
+}
+
+func TestWebSocketDeviceAuthorizerRejectsUnauthorized(t *testing.T) {
+	hub := newTestHub()
+	if err := hub.authorizeSubscription("user-1", "device-a"); err != nil {
+		t.Fatalf("authorized device rejected: %v", err)
+	}
+	if err := hub.authorizeSubscription("user-1", "device-b"); err == nil {
+		t.Fatal("unauthorized device should be rejected")
+	}
+}
+
+func TestBroadcastTelemetryDoesNotLeakGlobal(t *testing.T) {
+	hub := newTestHub()
+	subscribed := &client{rooms: map[string]bool{"device:device-a": true}, send: make(chan []byte, 1)}
+	unsubscribed := &client{rooms: map[string]bool{}, send: make(chan []byte, 1)}
+	hub.clients[subscribed] = true
+	hub.clients[unsubscribed] = true
+
+	hub.BroadcastTelemetry("device-a", map[string]interface{}{"value": 1})
+
+	if len(subscribed.send) == 0 {
+		t.Fatal("subscribed client should receive telemetry")
+	}
+	if len(unsubscribed.send) != 0 {
+		t.Fatal("unsubscribed client must not receive global telemetry")
+	}
+}
+
+func TestBroadcastAlarmDoesNotLeakGlobal(t *testing.T) {
+	hub := newTestHub()
+	subscribed := &client{rooms: map[string]bool{"device:device-a": true}, send: make(chan []byte, 1)}
+	unsubscribed := &client{rooms: map[string]bool{}, send: make(chan []byte, 1)}
+	hub.clients[subscribed] = true
+	hub.clients[unsubscribed] = true
+
+	hub.BroadcastAlarm(service.AlarmEvent{DeviceKey: "device-a", Code: "high"})
+
+	if len(unsubscribed.send) != 0 {
+		t.Fatal("unsubscribed client must not receive global alarm")
+	}
+}
+
+type brokenAuthorizer struct{}
+
+func (brokenAuthorizer) CanReadDevice(userID, deviceKey string) (bool, error) {
+	return false, errors.New("authorizer unavailable")
+}
+
+func TestWebSocketAuthorizerErrorIsDenied(t *testing.T) {
+	hub := NewHub("test-secret", brokenAuthorizer{}, nil)
+	if err := hub.authorizeSubscription("user-1", "device-a"); err == nil {
+		t.Fatal("authorizer error should deny subscription")
+	}
+	if _, ok := hub.consumeTicket("missing"); ok {
+		t.Fatal("missing ticket should not be consumed")
+	}
+}