diff --git a/SECURITY.md b/SECURITY.md index b863623b8ebf..b5b53f26fd99 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -90,9 +90,9 @@ Asynchronous execution paths do not create a separate data-access capability. A Task metadata is itself a request-scoped resource. Non-admin users and embedded guests may read or cancel only tasks they are subscribed to or that otherwise represent work they are entitled to observe; Admin may observe and manage tasks as part of the trusted operational boundary. A bug that lets a principal create, read, join, cancel, or receive task state for work outside the role and capability matrix is in scope. -Realtime transports, including WebSocket delivery backed by Redis or Valkey Pub/Sub, are notification mechanisms rather than authorization sources. WebSocket notification access is controlled by `can_read` on the `Realtime` resource. The broadcast scope is authenticated-global, not public: it reaches every authenticated realtime socket, and an anonymous request receives no realtime principal, no JWT cookie, and therefore no socket, so it never receives these messages (true anonymous/Public-role realtime is not offered and would require a separate, restricted model). Broadcast Pub/Sub messages, such as list-view entity-change events, must be context-free nudges; sensitive or authoritative state must not be published on the broadcast scope. Targeted Pub/Sub messages may carry task state only when the producer names routing keys derived from an authorized identity, such as a task subscriber's principal channel (or a per-tab channel derived from it); the producer validates every routing key against the task's own subscriber principals before publishing, and the websocket server forwards the payload only to sockets bound to those keys. Full data and result payloads must still be fetched through the normal protected REST API or cache-read path. Redis Streams used for task completion, dependency, and lock-release signalling are likewise coordination signals; the metastore or cache entry they wake a consumer to read remains the source of truth. +Realtime transports, including WebSocket delivery backed by Redis or Valkey Pub/Sub, are notification mechanisms rather than authorization sources. WebSocket notification access is controlled by `can_read` on the `Realtime` resource. The broadcast scope is authenticated-global, not public: it reaches every authenticated realtime socket, and an anonymous request receives no realtime principal, no connection ticket, and therefore no socket, so it never receives these messages (true anonymous/Public-role realtime is not offered and would require a separate, restricted model). Broadcast Pub/Sub messages, such as list-view entity-change events, must be context-free nudges; sensitive or authoritative state must not be published on the broadcast scope. Targeted Pub/Sub messages may carry task state only when the producer names routing keys derived from an authorized identity, such as a task subscriber's principal channel (or a per-tab channel derived from it); the producer validates every routing key against the task's own subscriber principals before publishing, and the websocket server forwards the payload only to sockets bound to those keys. Full data and result payloads must still be fetched through the normal protected REST API or cache-read path. Redis Streams used for task completion, dependency, and lock-release signalling are likewise coordination signals; the metastore or cache entry they wake a consumer to read remains the source of truth. -The realtime notification permission is distinct from the permission to read the underlying object. It controls whether a principal receives push notifications, not whether they may read the object once they call the protected REST API. Existing websocket connections are authorized by the JWT accepted at upgrade time; permission revocation after token minting is bounded by `WEBSOCKET_JWT_EXPIRATION_SECONDS` plus the websocket server's socket-check interval. Redis Streams are internal server-to-server coordination primitives and should not be directly exposed as an end-user subscription surface. +The realtime notification permission is distinct from the permission to read the underlying object. It controls whether a principal receives push notifications, not whether they may read the object once they call the protected REST API. An open websocket connection is authorized by the most recent connection ticket presented on it (the first `auth` frame, then each in-band refresh); permission revocation after a ticket is minted is bounded by `WEBSOCKET_JWT_EXPIRATION_SECONDS` plus the websocket server's socket-check interval, because the next refresh is refused. Redis Streams are internal server-to-server coordination primitives and should not be directly exposed as an end-user subscription surface. ### Version History and Change Metadata diff --git a/UPDATING.md b/UPDATING.md index fe7c1c26132d..4b88c71c69dd 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -1248,14 +1248,26 @@ WEBSOCKET_JWT_SECRET = "" The built-in Gamma role receives `can_read Realtime`; grant that permission to custom roles that should receive websocket notifications. -Run the `superset-websocket` Node server on the **same browser-visible host** -(so its JWT channel cookie is shared) and point its `redis` config at the same -instance as `DISTRIBUTED_COORDINATION_CONFIG`, plus `jwtSecret` / -`jwtCookieName` matching the Flask config (`WEBSOCKET_JWT_SECRET` / -`WEBSOCKET_JWT_COOKIE_NAME`, default `superset-ws-token`). During websocket JWT +Run the `superset-websocket` Node server where browsers can reach it and point +its `redis` config at the same instance as `DISTRIBUTED_COORDINATION_CONFIG`, +plus a `jwtSecret` matching the Flask config (`WEBSOCKET_JWT_SECRET`). The +browser authenticates with a short-lived connection ticket fetched from +`GET /api/v1/realtime/token/` and sent as the first frame of the socket, then +refreshes it on the open socket before it expires, so the server no longer needs +to share a host with the web application and no credential rides the handshake or +the URL. Upgrade the web servers and the websocket servers together, and reload +open browser tabs: a tab still running the previous cookie-based client is +refused by the new server (it closes the socket after the 10 second auth +timeout, repeatedly) until it is reloaded, and the same applies to replicas rolled +at different times. The `WEBSOCKET_JWT_COOKIE_NAME`, +`WEBSOCKET_JWT_COOKIE_SECURE`, `WEBSOCKET_JWT_COOKIE_SAMESITE` and +`WEBSOCKET_JWT_COOKIE_DOMAIN` Flask settings and the websocket server's +`jwtCookieName` / `JWT_COOKIE_NAME` option are removed. During websocket JWT secret rotation, set the websocket server's `previousJwtSecret` / -`PREVIOUS_JWT_SECRET` to the old key while Flask continues minting cookies with -`WEBSOCKET_JWT_SECRET`. The server is bundled in the official Superset image +`PREVIOUS_JWT_SECRET` to the old key while Flask continues minting tickets with +`WEBSOCKET_JWT_SECRET`. After a logout or session revocation, an open socket stays +authenticated until its current ticket expires, because the next refresh fails. +The server is bundled in the official Superset image and launched via an alternate entrypoint — no separate image is required: `docker run /app/docker/entrypoints/run-websocket.sh` (or the opt-in `websocket` profile in `docker compose`). It **subscribes** to a single diff --git a/docker-compose-non-dev.yml b/docker-compose-non-dev.yml index 8032c44b2317..ae159d4c7c42 100644 --- a/docker-compose-non-dev.yml +++ b/docker-compose-non-dev.yml @@ -159,7 +159,6 @@ services: REDIS_HOST: redis REDIS_PORT: 6379 PORT: 8080 - JWT_COOKIE_NAME: superset-ws-token # Dev-only default; must match the app's WEBSOCKET_JWT_SECRET and be # replaced with a strong secret (>= 32 bytes) outside local development. JWT_SECRET: ${WEBSOCKET_JWT_SECRET:-dev-only-websocket-secret-change-me!} diff --git a/docker/entrypoints/run-websocket.sh b/docker/entrypoints/run-websocket.sh index 80724805ae28..10b5be1f2bc9 100755 --- a/docker/entrypoints/run-websocket.sh +++ b/docker/entrypoints/run-websocket.sh @@ -26,7 +26,6 @@ # the authoritative, complete set (Redis connection, logging, connection limits, # StatsD, etc.). The values that MUST match the Flask app's config are: # JWT_SECRET == WEBSOCKET_JWT_SECRET -# JWT_COOKIE_NAME == WEBSOCKET_JWT_COOKIE_NAME (default superset-ws-token) # REALTIME_CHANNEL_PREFIX == Flask REALTIME_CHANNEL_PREFIX (default empty; set a # per-deployment value on both sides to isolate a shared Redis/Valkey) # Optional rotation setting: diff --git a/docker/superset-websocket/config.example.json b/docker/superset-websocket/config.example.json index 93578d17372b..33a29a30dca9 100644 --- a/docker/superset-websocket/config.example.json +++ b/docker/superset-websocket/config.example.json @@ -17,6 +17,5 @@ }, "jwtAlgorithms": ["HS256"], "jwtSecret": "CHANGE-ME-IN-PRODUCTION-GOTTA-BE-LONG-AND-SECRET", - "previousJwtSecret": "", - "jwtCookieName": "superset-ws-token" + "previousJwtSecret": "" } diff --git a/docs/admin_docs/security/securing_superset.mdx b/docs/admin_docs/security/securing_superset.mdx index 19c13e42735f..6dad74efc95b 100644 --- a/docs/admin_docs/security/securing_superset.mdx +++ b/docs/admin_docs/security/securing_superset.mdx @@ -180,14 +180,14 @@ https://superset.apache.org/admin-docs/configuration/configuring-superset/#rotat | --------------------------------- | ------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------- | ----------------------------------------- | | `SUPERSET_SECRET_KEY` | Signs session cookies; key material for encrypting stored DB credentials (Fernet/AES) | Forged sessions (auth bypass / privilege escalation); decryption of exfiltrated metadata-DB secrets | Quarterly + post-incident | | `GUEST_TOKEN_JWT_SECRET` | Signs embedded-dashboard guest tokens | Forged guest tokens → unauthorized dashboard/data access | Quarterly + post-incident | -| `WEBSOCKET_JWT_SECRET` | Signs the realtime websocket channel JWT cookie | Forged websocket tokens → unauthorized realtime notifications | Quarterly + post-incident | +| `WEBSOCKET_JWT_SECRET` | Signs the realtime websocket connection tickets | Forged websocket tokens → unauthorized realtime notifications | Quarterly + post-incident | | SMTP password | Outbound email for alerts & reports | Email relay abuse / spoofing | Per organizational policy + post-incident | | Database connection passwords | Access to analytical databases and the metadata DB | Direct database access | Per organizational policy + post-incident | Notes: - Rotating `GUEST_TOKEN_JWT_SECRET` or `WEBSOCKET_JWT_SECRET` invalidates outstanding tokens of that type; schedule rotations accordingly. -- `WEBSOCKET_JWT_SECRET` can be rotated without disconnecting live sockets: set the outgoing value as `PREVIOUS_JWT_SECRET` on the websocket server so it keeps verifying old cookies, then remove it once they have aged out. +- `WEBSOCKET_JWT_SECRET` can be rotated without disconnecting live sockets: set the outgoing value as `PREVIOUS_JWT_SECRET` on the websocket server so it keeps verifying old tickets, then remove it once they have aged out. - After a suspected compromise, rotate **all** of the above, not only `SUPERSET_SECRET_KEY`. - Keep the register under change control so new secrets introduced by future features are added to the rotation schedule. diff --git a/superset-frontend/src/embedded/index.test.tsx b/superset-frontend/src/embedded/index.test.tsx index b9e8b1caa7d7..feeb4704b73d 100644 --- a/superset-frontend/src/embedded/index.test.tsx +++ b/superset-frontend/src/embedded/index.test.tsx @@ -87,6 +87,11 @@ jest.mock('@superset-ui/switchboard', () => ({ const mockSetupClient = jest.fn(); jest.mock('src/setup/setupClient', () => mockSetupClient, { virtual: true }); +const mockConnectRealtime = jest.fn(); +jest.mock('src/middleware/realtime', () => ({ + connectRealtime: () => mockConnectRealtime(), +})); + jest.mock('src/views/store', () => ({ store: { dispatch: jest.fn(), @@ -267,6 +272,22 @@ test('oversized refresh updates diagnostics without restarting successful auth', ); }); +test('reconnects realtime once the guest client is configured', async () => { + require('./index'); + await flush(); + sendHandshake(); + mockSetupClient.mockClear(); + mockConnectRealtime.mockClear(); + + mockSwitchboard.handler!({ guestToken: 'short' }); + await flush(); + + expect(mockConnectRealtime).toHaveBeenCalledTimes(1); + expect(mockSetupClient.mock.invocationCallOrder[0]).toBeLessThan( + mockConnectRealtime.mock.invocationCallOrder[0], + ); +}); + test('refresh during pending authentication does not misattribute size evidence', async () => { let rejectRequest: (error: unknown) => void = () => {}; mockGetMeWithRole.mockReturnValue( diff --git a/superset-frontend/src/embedded/index.tsx b/superset-frontend/src/embedded/index.tsx index d23786aceced..6f3317614551 100644 --- a/superset-frontend/src/embedded/index.tsx +++ b/superset-frontend/src/embedded/index.tsx @@ -36,6 +36,7 @@ import getBootstrapData, { applicationRoot } from 'src/utils/getBootstrapData'; import initPreamble from 'src/preamble'; import { setupAGGridModules } from '@superset-ui/core/components/ThemedAgGridReact'; import setupClient from 'src/setup/setupClient'; +import { connectRealtime } from 'src/middleware/realtime'; import { useUiConfig } from 'src/components/UiConfigContext'; import { store, USER_LOADED } from 'src/views/store'; import { Loading } from '@superset-ui/core/components'; @@ -284,6 +285,10 @@ function setupGuestClient(guestToken: string) { guestTokenHeaderName: bootstrapData.config?.GUEST_TOKEN_HEADER_NAME, unauthorizedHandler: guestUnauthorizedHandler, }); + // The realtime socket first connects at module load, before a guest token + // exists, and its ticket request is refused. Reconnect with the guest + // credential; a no-op while a socket is already open. + connectRealtime(); } window.addEventListener('message', function embeddedPageInitializer(event) { diff --git a/superset-frontend/src/middleware/asyncEvent.test.ts b/superset-frontend/src/middleware/asyncEvent.test.ts index 6872fe35da2d..c1c9829df07a 100644 --- a/superset-frontend/src/middleware/asyncEvent.test.ts +++ b/superset-frontend/src/middleware/asyncEvent.test.ts @@ -43,9 +43,11 @@ jest.mock('src/hooks/useTabId', () => { /* eslint-disable no-var, vars-on-top */ var mockRealtimeOpenListener: (() => void) | undefined; var mockRealtimeStateListener: ((state: string) => void) | undefined; +var mockRealtimeState: string; /* eslint-enable no-var, vars-on-top */ jest.mock('src/middleware/realtime', () => ({ connectRealtime: jest.fn(), + getRealtimeState: () => mockRealtimeState ?? 'open', subscribeRealtime: jest.fn(() => () => {}), subscribeRealtimeOpen: (listener: () => void) => { mockRealtimeOpenListener = listener; @@ -55,6 +57,7 @@ jest.mock('src/middleware/realtime', () => ({ }, subscribeRealtimeState: (listener: (state: string) => void) => { mockRealtimeStateListener = listener; + mockRealtimeState = 'open'; return () => { mockRealtimeStateListener = undefined; }; @@ -857,6 +860,42 @@ test('WS mode: an unhealthy socket settles pending waiters with a bounded error' expect(refetch).not.toHaveBeenCalled(); }); +test('WS mode: an unauthorized socket settles pending waiters with a bounded error', async () => { + queueStatuses(); + asyncEvent.init(wsConfig); + + const refetch = jest.fn(); + const promise = asyncEvent.waitForAsyncData( + { task_ids: ['task-1'] }, + refetch, + ); + await new Promise(resolve => { + setTimeout(resolve, 0); + }); + + // The session lost the realtime permission: no socket will ever settle this. + mockRealtimeStateListener?.('unauthorized'); + + await expect(promise).rejects.toThrow('Realtime connection unavailable'); + expect(refetch).not.toHaveBeenCalled(); +}); + +test.each(['unauthorized', 'unhealthy'])( + 'WS mode: a waiter registered while the socket is %s settles immediately', + async state => { + queueStatuses(); + asyncEvent.init(wsConfig); + mockRealtimeState = state; + + const refetch = jest.fn(); + await expect( + asyncEvent.waitForAsyncData({ task_ids: ['task-1'] }, refetch), + ).rejects.toThrow('Realtime connection unavailable'); + expect(refetch).not.toHaveBeenCalled(); + mockRealtimeState = 'open'; + }, +); + test('WS mode: coalesces many same-tick registrations into one catch-up', async () => { jest.useFakeTimers(); queueStatuses(); // all catch-ups see no change diff --git a/superset-frontend/src/middleware/asyncEvent.ts b/superset-frontend/src/middleware/asyncEvent.ts index 42849ab127ae..3335aed131c1 100644 --- a/superset-frontend/src/middleware/asyncEvent.ts +++ b/superset-frontend/src/middleware/asyncEvent.ts @@ -44,6 +44,7 @@ import { getTabId } from 'src/hooks/useTabId'; import { getAsyncQueryError } from './asyncQueryError'; import { connectRealtime, + getRealtimeState, subscribeRealtime, subscribeRealtimeOpen, subscribeRealtimeState, @@ -292,6 +293,11 @@ const abandonRealtimeWaiters = () => { ); }; +const isRealtimeDown = (): boolean => { + const state = getRealtimeState(); + return state === 'unhealthy' || state === 'unauthorized'; +}; + const finishStatus = (taskId: string, status: string, message?: string) => { const waiters = waitersByTaskId.get(taskId); if (!waiters) return; @@ -507,7 +513,9 @@ subscribeRealtimeState(state => { if (!wsEnabled || !waitersByTaskId.size) return; if (state === 'reconnecting') { scheduleCatchUp(); - } else if (state === 'unhealthy') { + } else if (state === 'unhealthy' || state === 'unauthorized') { + // `unauthorized` is terminal (the session lost the realtime permission or + // ended), so no socket will ever settle these waiters. abandonRealtimeWaiters(); } }); @@ -607,6 +615,14 @@ export const waitForAsyncData = async ( } waiters.add(waiter); }); + // The realtime connection state is the source of truth: a waiter registering + // while the transport is already terminally down (`unauthorized`) or + // `unhealthy` is settled at once, not left to wait on a transition that + // already happened. + if (wsEnabled && isRealtimeDown()) { + settle(waiter, new Error('Realtime connection unavailable')); + return; + } if (wsEnabled) { // WS is the transport (no polling). Bound a genuinely-lost completion (a // dropped socket message with no reconnect) so a chart can't spin forever; diff --git a/superset-frontend/src/middleware/realtime.test.ts b/superset-frontend/src/middleware/realtime.test.ts index 9523544d65ee..aa5990e06084 100644 --- a/superset-frontend/src/middleware/realtime.test.ts +++ b/superset-frontend/src/middleware/realtime.test.ts @@ -20,6 +20,7 @@ import { SupersetClient } from '@superset-ui/core'; import { connectRealtime, dispatchRealtimeMessage, + getRealtimeState, resetRealtimeForTests, subscribeRealtime, subscribeRealtimeOpen, @@ -27,7 +28,7 @@ import { } from 'src/middleware/realtime'; // Controllable useTabId mock (overrides the global shim for this file) so we can -// drive both the current tab id baked into the connect URL and tab-id-change +// drive both the current tab id sent in the `auth` frame and tab-id-change // notifications. `var` + lazy init sidesteps jest.mock hoisting/TDZ. /* eslint-disable no-var, vars-on-top */ var mockTabIdValue: string; @@ -54,7 +55,7 @@ const fireTabIdChange = (newId: string) => { // A minimal fake WebSocket so connectRealtime can "open" a socket without a // real network connection. Instances register themselves so a test can drive -// onmessage/onclose. +// onopen/onmessage/onclose and inspect what the client sent. class FakeWebSocket { static instances: FakeWebSocket[] = []; @@ -64,10 +65,16 @@ class FakeWebSocket { onmessage: ((event: { data: string }) => void) | null = null; - onclose: (() => void) | null = null; + onclose: ((event: { code: number }) => void) | null = null; onerror: (() => void) | null = null; + sent: string[] = []; + + send = jest.fn((data: string) => { + this.sent.push(data); + }); + close = jest.fn(() => { this.readyState = 3; }); @@ -75,26 +82,84 @@ class FakeWebSocket { constructor(public url: string) { FakeWebSocket.instances.push(this); } + + /** The `auth` frames the client has sent, parsed. */ + get authFrames(): Array<{ type: string; token: string; tab_id: string }> { + return this.sent.map(frame => JSON.parse(frame)); + } } const originalWebSocket = global.WebSocket; +const ENABLED = { + WEBSOCKET_ENABLE: true, + WEBSOCKET_URL: 'ws://localhost:8080/', +}; + +let getSpy: jest.SpyInstance; +let ticketCount: number; + +/** Resolve the ticket endpoint with a distinct token and the given lifetime. */ +const serveTickets = (expiresIn = 60) => { + ticketCount = 0; + getSpy.mockImplementation(async () => { + ticketCount += 1; + return { json: { token: `ticket-${ticketCount}`, expires_in: expiresIn } }; + }); +}; + +/** Reject the ticket endpoint like SupersetClient does: with the Response. */ +const rejectTickets = (status: number) => { + getSpy.mockRejectedValue({ status }); +}; + +// Let pending promise continuations (the ticket fetch) run. +const flush = async () => { + for (let i = 0; i < 10; i += 1) { + // eslint-disable-next-line no-await-in-loop + await Promise.resolve(); + } +}; + +const lastSocket = () => + FakeWebSocket.instances[FakeWebSocket.instances.length - 1]; + +/** The server's acknowledgement of an `auth` frame. */ +const ack = (ws: FakeWebSocket, expiresIn = 60) => + ws.onmessage?.({ + data: JSON.stringify({ + topic: 'realtime.auth', + payload: { status: 'ok', expires_in: expiresIn }, + }), + }); + +/** Connect, open the socket, and have the server acknowledge the first frame. */ +const connectAuthenticated = async (expiresIn = 60) => { + connectRealtime(ENABLED); + await flush(); + const ws = lastSocket(); + ws.readyState = 1; + ws.onopen?.(); + ack(ws, expiresIn); + return ws; +}; + beforeEach(() => { + jest.useFakeTimers(); FakeWebSocket.instances = []; global.WebSocket = FakeWebSocket as unknown as typeof WebSocket; + getSpy = jest.spyOn(SupersetClient, 'get'); + serveTickets(); }); afterEach(() => { resetRealtimeForTests(); global.WebSocket = originalWebSocket; mockTabIdValue = 'test-tab-id'; + getSpy.mockRestore(); + jest.useRealTimers(); }); -const ENABLED = { - WEBSOCKET_ENABLE: true, - WEBSOCKET_URL: 'ws://localhost:8080/', -}; - test('dispatches a parsed envelope to every handler of its topic', () => { const a = jest.fn(); const b = jest.fn(); @@ -161,230 +226,567 @@ test('a throwing handler does not break other handlers', () => { expect(good).toHaveBeenCalledTimes(1); }); -test('notifies open-listeners when the socket connects', () => { - const onOpen = jest.fn(); - const unsubscribe = subscribeRealtimeOpen(onOpen); - connectRealtime(ENABLED); +test('does not open a socket when disabled', async () => { + connectRealtime({ WEBSOCKET_ENABLE: false, WEBSOCKET_URL: 'ws://x/' }); + await flush(); - expect(FakeWebSocket.instances).toHaveLength(1); - FakeWebSocket.instances[0].onopen?.(); - expect(onOpen).toHaveBeenCalledTimes(1); + expect(FakeWebSocket.instances).toHaveLength(0); + expect(getSpy).not.toHaveBeenCalled(); +}); - // A reopen (openSocket wires onopen on every socket) notifies again. - FakeWebSocket.instances[0].onopen?.(); - expect(onOpen).toHaveBeenCalledTimes(2); +test('fetches a ticket, keeps it out of the URL, and authenticates in-band', async () => { + const states: string[] = []; + subscribeRealtimeState(state => states.push(state)); + connectRealtime(ENABLED); + await flush(); - unsubscribe(); - FakeWebSocket.instances[0].onopen?.(); - expect(onOpen).toHaveBeenCalledTimes(2); // no longer notified + // A refused ticket must not trigger the client's login redirect. + expect(getSpy).toHaveBeenCalledWith({ + endpoint: '/api/v1/realtime/token/', + ignoreUnauthorized: true, + }); + const ws = lastSocket(); + expect(ws.url).toBe(ENABLED.WEBSOCKET_URL); + + ws.readyState = 1; + ws.onopen?.(); + // The first frame carries the ticket and this tab's id. + expect(ws.authFrames).toEqual([ + { type: 'auth', token: 'ticket-1', tab_id: 'test-tab-id' }, + ]); + // Not reported open until the server acknowledges the ticket. + expect(states).toEqual(['connecting']); + + ack(ws); + expect(states).toEqual(['connecting', 'open']); }); -test('does not open a socket when disabled', () => { - connectRealtime({ WEBSOCKET_ENABLE: false, WEBSOCKET_URL: 'ws://x/' }); - expect(FakeWebSocket.instances).toHaveLength(0); +test('resolves a root-relative ws url against the page', async () => { + connectRealtime({ WEBSOCKET_ENABLE: true, WEBSOCKET_URL: '/superset-ws' }); + await flush(); + + const u = new URL(lastSocket().url); + expect(u.protocol).toBe('ws:'); // jsdom page is http:// -> ws:// + expect(u.pathname).toBe('/superset-ws'); + expect(u.search).toBe(''); }); -test('opens a socket when enabled and routes its messages to handlers', () => { +test('routes socket messages to handlers only after the server acknowledges auth', async () => { const handler = jest.fn(); subscribeRealtime('task.status', handler); connectRealtime(ENABLED); + await flush(); + const ws = lastSocket(); + const message = { + data: JSON.stringify({ topic: 'task.status', payload: { ok: true } }), + }; + + ws.onmessage?.(message); + expect(handler).not.toHaveBeenCalled(); + + ack(ws); + ws.onmessage?.(message); + expect(handler).toHaveBeenCalledWith({ ok: true }); +}); + +test('notifies open-listeners once per connection, when auth is acknowledged', async () => { + const reasons: string[] = []; + subscribeRealtimeOpen(reason => reasons.push(reason)); + connectRealtime(ENABLED); + await flush(); + const ws = lastSocket(); + ws.readyState = 1; + ws.onopen?.(); + expect(reasons).toEqual([]); + + ack(ws); + ack(ws); // a refresh acknowledgement is not a new connection + + expect(reasons).toEqual(['initial']); +}); + +test('keeps the active socket when reconnecting with the same config', async () => { + const ws = await connectAuthenticated(); + + connectRealtime(ENABLED); + await flush(); expect(FakeWebSocket.instances).toHaveLength(1); - const ws = FakeWebSocket.instances[0]; - // The connect URL is the configured base plus this tab's id, so the server - // can bind a per-tab channel. getTabId is mocked to 'test-tab-id' (shim). - const wsUrl = new URL(ws.url); - expect(`${wsUrl.protocol}//${wsUrl.host}${wsUrl.pathname}`).toBe( - ENABLED.WEBSOCKET_URL, - ); - expect(wsUrl.searchParams.get('tab_id')).toBe('test-tab-id'); + expect(ws.close).not.toHaveBeenCalled(); +}); - ws.onmessage?.({ - data: JSON.stringify({ topic: 'task.status', payload: { ok: true } }), - }); +test('refreshes the ticket on the open socket at 75% of its lifetime', async () => { + const ws = await connectAuthenticated(60); + expect(getSpy).toHaveBeenCalledTimes(1); - expect(handler).toHaveBeenCalledWith({ ok: true }); + jest.advanceTimersByTime(44_000); + await flush(); + expect(getSpy).toHaveBeenCalledTimes(1); + + jest.advanceTimersByTime(1_000); + await flush(); + + // A fresh ticket goes out on the same socket: no teardown, no reconnect. + expect(getSpy).toHaveBeenCalledTimes(2); + expect(ws.authFrames.map(frame => frame.token)).toEqual([ + 'ticket-1', + 'ticket-2', + ]); + expect(ws.close).not.toHaveBeenCalled(); + expect(FakeWebSocket.instances).toHaveLength(1); + + // The server's reply schedules the next refresh from its reported lifetime. + ack(ws, 60); + jest.advanceTimersByTime(45_000); + await flush(); + expect(ws.authFrames.map(frame => frame.token)).toEqual([ + 'ticket-1', + 'ticket-2', + 'ticket-3', + ]); + expect(FakeWebSocket.instances).toHaveLength(1); }); -test('resolves a root-relative ws url against the page and adds tab_id', () => { - // A root-relative endpoint (e.g. behind a same-origin proxy) is valid for - // WebSocket; it must resolve against the page URL and normalize to ws(s), - // still carrying tab_id — not silently drop the tab id. - connectRealtime({ WEBSOCKET_ENABLE: true, WEBSOCKET_URL: '/superset-ws' }); +test('a socket survives several ticket lifetimes with zero reconnects', async () => { + const ws = await connectAuthenticated(60); + + for (let i = 0; i < 4; i += 1) { + jest.advanceTimersByTime(45_000); + // eslint-disable-next-line no-await-in-loop + await flush(); + ack(ws, 60); + } expect(FakeWebSocket.instances).toHaveLength(1); - const u = new URL(FakeWebSocket.instances[0].url); - expect(u.protocol).toBe('ws:'); // jsdom page is http:// → ws:// - expect(u.pathname).toBe('/superset-ws'); - expect(u.searchParams.get('tab_id')).toBe('test-tab-id'); + expect(ws.close).not.toHaveBeenCalled(); + expect(ws.authFrames).toHaveLength(5); }); -test('reconnects with the new tab id when the tab id changes', () => { +test('retries a transient ticket failure during refresh until expiry', async () => { + const ws = await connectAuthenticated(60); + getSpy.mockRejectedValue({ status: 503 }); + + jest.advanceTimersByTime(45_000); + await flush(); + expect(getSpy).toHaveBeenCalledTimes(2); + + jest.advanceTimersByTime(1_000); + await flush(); + expect(getSpy).toHaveBeenCalledTimes(3); + + // Service recovers: the next retry sends the fresh ticket. + serveTickets(60); + jest.advanceTimersByTime(2_000); + await flush(); + expect(ws.authFrames).toHaveLength(2); + expect(ws.close).not.toHaveBeenCalled(); +}); + +test('stops retrying a failing refresh once the ticket has expired', async () => { + await connectAuthenticated(60); + getSpy.mockRejectedValue({ status: 503 }); + + jest.advanceTimersByTime(45_000 + 15_000); + await flush(); + const callsAtExpiry = getSpy.mock.calls.length; + + jest.advanceTimersByTime(60_000); + await flush(); + + expect(getSpy.mock.calls.length).toBe(callsAtExpiry); +}); + +test('an unauthorized ticket request is terminal: no socket, no retries', async () => { + rejectTickets(401); + const states: string[] = []; + subscribeRealtimeState(state => states.push(state)); + + connectRealtime(ENABLED); + await flush(); + jest.advanceTimersByTime(5 * 60_000); + await flush(); + + expect(states).toEqual(['connecting', 'unauthorized']); + expect(FakeWebSocket.instances).toHaveLength(0); + expect(getSpy).toHaveBeenCalledTimes(1); +}); + +test('connecting again after an unauthorized ticket starts a new attempt', async () => { + rejectTickets(401); connectRealtime(ENABLED); + await flush(); + expect(getRealtimeState()).toBe('unauthorized'); + + // e.g. an embedded page once its guest token has been configured + serveTickets(); + connectRealtime(ENABLED); + await flush(); + expect(FakeWebSocket.instances).toHaveLength(1); - expect( - new URL(FakeWebSocket.instances[0].url).searchParams.get('tab_id'), - ).toBe('test-tab-id'); + expect(getSpy).toHaveBeenCalledTimes(2); +}); - // A duplicated tab is reassigned a new id: the socket must reconnect so it - // re-registers under the new per-tab channel (else tab-targeted events miss - // it and only polling notices completion). - fireTabIdChange('tab-2'); +test('a forbidden ticket request is terminal too', async () => { + rejectTickets(403); + const states: string[] = []; + subscribeRealtimeState(state => states.push(state)); + + connectRealtime(ENABLED); + await flush(); + + expect(states[states.length - 1]).toBe('unauthorized'); + expect(FakeWebSocket.instances).toHaveLength(0); +}); + +test('losing authorization at refresh closes the socket and stops connecting', async () => { + const ws = await connectAuthenticated(60); + const states: string[] = []; + subscribeRealtimeState(state => states.push(state)); + rejectTickets(401); + jest.advanceTimersByTime(45_000); + await flush(); + jest.advanceTimersByTime(5 * 60_000); + await flush(); + + expect(states).toEqual(['unauthorized']); + expect(ws.close).toHaveBeenCalled(); + expect(FakeWebSocket.instances).toHaveLength(1); + expect(getSpy).toHaveBeenCalledTimes(2); +}); + +test('reconnects immediately with a fresh ticket when the server closes with 4401', async () => { + const ws = await connectAuthenticated(); + + ws.readyState = 3; + ws.onclose?.({ code: 4401 }); + await flush(); + + // No backoff timer was needed. expect(FakeWebSocket.instances).toHaveLength(2); - expect(FakeWebSocket.instances[0].close).toHaveBeenCalled(); - expect( - new URL(FakeWebSocket.instances[1].url).searchParams.get('tab_id'), - ).toBe('tab-2'); + expect(getSpy).toHaveBeenCalledTimes(2); + const next = lastSocket(); + next.readyState = 1; + next.onopen?.(); + expect(next.authFrames[0].token).toBe('ticket-2'); }); -test('keeps the active socket when reconnecting with the same config', () => { - connectRealtime(ENABLED); - connectRealtime(ENABLED); +test('reports a reconnect after 4401 as reason `reconnect`', async () => { + const reasons: string[] = []; + subscribeRealtimeOpen(reason => reasons.push(reason)); + const ws = await connectAuthenticated(); + + ws.readyState = 3; + ws.onclose?.({ code: 4401 }); + await flush(); + const next = lastSocket(); + next.readyState = 1; + next.onopen?.(); + ack(next); + + expect(reasons).toEqual(['initial', 'reconnect']); +}); +test('backs off after other closes, fetching a fresh ticket on each attempt', async () => { + const ws = await connectAuthenticated(); + + ws.readyState = 3; + ws.onclose?.({ code: 1006 }); + await flush(); expect(FakeWebSocket.instances).toHaveLength(1); - expect(FakeWebSocket.instances[0].close).not.toHaveBeenCalled(); + + jest.advanceTimersByTime(5_000); + await flush(); + + expect(FakeWebSocket.instances).toHaveLength(2); + expect(getSpy).toHaveBeenCalledTimes(2); }); -test('reconnects after the socket closes', () => { - jest.useFakeTimers(); +test('a 4401 on a connection that never authenticated backs off', async () => { connectRealtime(ENABLED); + await flush(); + const ws = lastSocket(); + ws.readyState = 1; + ws.onopen?.(); + + ws.readyState = 3; + ws.onclose?.({ code: 4401 }); + await flush(); expect(FakeWebSocket.instances).toHaveLength(1); - FakeWebSocket.instances[0].onclose?.(); - jest.advanceTimersByTime(5000); + jest.advanceTimersByTime(5_000); + await flush(); + expect(FakeWebSocket.instances).toHaveLength(2); +}); + +test('abandons and retries a connection the server never acknowledges', async () => { + connectRealtime(ENABLED); + await flush(); + const ws = lastSocket(); + ws.readyState = 1; + ws.onopen?.(); + + jest.advanceTimersByTime(10_000); + expect(ws.close).toHaveBeenCalled(); + jest.advanceTimersByTime(5_000); + await flush(); expect(FakeWebSocket.instances).toHaveLength(2); - jest.useRealTimers(); }); -test('reports reconnecting then unhealthy after repeated failures, and open resets it', () => { - jest.useFakeTimers(); +test('reports reconnecting then unhealthy after repeated failures, and open resets it', async () => { const states: string[] = []; subscribeRealtimeState(s => states.push(s)); connectRealtime(ENABLED); + await flush(); // Three consecutive closes with no successful OPEN in between. for (let i = 0; i < 3; i += 1) { - FakeWebSocket.instances[FakeWebSocket.instances.length - 1].onclose?.(); + lastSocket().onclose?.({ code: 1006 }); // Advance past the max backoff so the scheduled reconnect fires. - jest.advanceTimersByTime(30000); + jest.advanceTimersByTime(30_000); + // eslint-disable-next-line no-await-in-loop + await flush(); } expect(states).toContain('reconnecting'); expect(states).toContain('unhealthy'); - // A successful OPEN clears the failure state back to healthy. - FakeWebSocket.instances[FakeWebSocket.instances.length - 1].onopen?.(); + // A successful, acknowledged connection clears the failure state. + const ws = lastSocket(); + ws.readyState = 1; + ws.onopen?.(); + ack(ws); expect(states[states.length - 1]).toBe('open'); - jest.useRealTimers(); }); -test('a synchronous WebSocket constructor failure still schedules a reconnect', () => { - jest.useFakeTimers(); +test('a failed ticket fetch is retried with backoff and counts toward unhealthy', async () => { + getSpy.mockRejectedValue({ status: 500 }); + const states: string[] = []; + subscribeRealtimeState(s => states.push(s)); + + connectRealtime(ENABLED); + await flush(); + for (let i = 0; i < 3; i += 1) { + jest.advanceTimersByTime(30_000); + // eslint-disable-next-line no-await-in-loop + await flush(); + } + + expect(FakeWebSocket.instances).toHaveLength(0); + expect(getSpy.mock.calls.length).toBeGreaterThan(1); + expect(states).toContain('unhealthy'); +}); + +test('a synchronous WebSocket constructor failure still schedules a reconnect', async () => { const Broken = jest.fn(() => { throw new Error('construct failed'); }); global.WebSocket = Broken as unknown as typeof WebSocket; connectRealtime(ENABLED); + await flush(); // The throw is caught; a reconnect is scheduled rather than dead-ending, so the // constructor is retried (repeatedly, with backoff) instead of only once. - jest.advanceTimersByTime(30000); + for (let i = 0; i < 3; i += 1) { + jest.advanceTimersByTime(30_000); + // eslint-disable-next-line no-await-in-loop + await flush(); + } expect(Broken.mock.calls.length).toBeGreaterThan(1); - jest.useRealTimers(); }); -test('disconnect closes the socket and stops reconnecting', () => { - jest.useFakeTimers(); +test('re-authenticates in place with the new tab id when the tab id changes', async () => { + const ws = await connectAuthenticated(); + expect(ws.authFrames[0].tab_id).toBe('test-tab-id'); + + // A duplicated tab is reassigned a new id: the server must rebind the per-tab + // channel (else tab-targeted events miss it and only polling notices + // completion), and it does so from an `auth` frame on the live socket. + const reasons: string[] = []; + subscribeRealtimeOpen(reason => reasons.push(reason)); + fireTabIdChange('tab-2'); + await flush(); + + expect(FakeWebSocket.instances).toHaveLength(1); + expect(ws.close).not.toHaveBeenCalled(); + expect(ws.authFrames[1]).toEqual({ + type: 'auth', + token: 'ticket-2', + tab_id: 'tab-2', + }); + + // The rebind acknowledgement raises exactly one reconciliation event, so a + // message sent to the old tab channel in the meantime is not lost. + expect(reasons).toEqual([]); + ack(ws); + expect(reasons).toEqual(['rebind']); +}); + +test('an ordinary ticket refresh acknowledgement raises no open event', async () => { + const reasons: string[] = []; + subscribeRealtimeOpen(reason => reasons.push(reason)); + const ws = await connectAuthenticated(60); + expect(reasons).toEqual(['initial']); + + jest.advanceTimersByTime(45_000); + await flush(); + ack(ws, 60); + + expect(reasons).toEqual(['initial']); +}); + +test('a tab id change before the socket is open is picked up by the first frame', async () => { connectRealtime(ENABLED); - const ws = FakeWebSocket.instances[0]; + await flush(); + const ws = lastSocket(); + + fireTabIdChange('tab-2'); + await flush(); + ws.readyState = 1; + ws.onopen?.(); + + expect(ws.authFrames).toEqual([ + { type: 'auth', token: 'ticket-1', tab_id: 'tab-2' }, + ]); +}); + +test('disconnect closes the socket and stops reconnecting', async () => { + const ws = await connectAuthenticated(); // disconnect via a superseding reset; the prior socket must not reconnect. resetRealtimeForTests(); expect(ws.close).toHaveBeenCalled(); - ws.onclose?.(); // a late close from the torn-down socket - jest.advanceTimersByTime(5000); + ws.onclose?.({ code: 1006 }); // a late close from the torn-down socket + jest.advanceTimersByTime(5_000); + await flush(); expect(FakeWebSocket.instances).toHaveLength(1); // no new socket - jest.useRealTimers(); }); -test('refreshes the cookie and reconnects before the token expires', async () => { - jest.useFakeTimers(); - const getSpy = jest - .spyOn(SupersetClient, 'get') - .mockResolvedValue({} as never); - try { - // A 900s token → the keepalive fires at 0.6 of its life (540s), issuing the - // authed cookie-refresh GET, then reconnecting so the new handshake rides a - // freshly minted token before the server terminates the socket at expiry. - connectRealtime({ ...ENABLED, WEBSOCKET_JWT_EXPIRATION_SECONDS: 900 }); - expect(FakeWebSocket.instances).toHaveLength(1); - // Simulate the socket having connected — the keepalive only hands off a - // socket that is still OPEN (readyState 1). - FakeWebSocket.instances[0].readyState = 1; +test('a ticket that arrives after disconnect does not open a socket', async () => { + connectRealtime(ENABLED); + resetRealtimeForTests(); + await flush(); - jest.advanceTimersByTime(540_000); - expect(getSpy).toHaveBeenCalledWith({ endpoint: '/api/v1/me/' }); + expect(FakeWebSocket.instances).toHaveLength(0); +}); - // Flush the refresh promise chain so the .finally reconnect runs. - await Promise.resolve(); - await Promise.resolve(); - expect(FakeWebSocket.instances).toHaveLength(2); - } finally { - getSpy.mockRestore(); - jest.useRealTimers(); - } +/** A ticket endpoint whose responses the test releases by hand. */ +const deferTickets = () => { + const pending: Array< + (ticket: { token: string; expires_in: number }) => void + > = []; + getSpy.mockImplementation( + () => + new Promise(resolve => { + pending.push(json => resolve({ json })); + }), + ); + return pending; +}; + +test('a ticket fetched for a refresh is dropped if the client disconnected meanwhile', async () => { + const ws = await connectAuthenticated(60); + const pending = deferTickets(); + + jest.advanceTimersByTime(45_000); // refresh starts; its fetch is in flight + expect(pending).toHaveLength(1); + resetRealtimeForTests(); + pending[0]({ token: 'late', expires_in: 60 }); + await flush(); + + expect(ws.authFrames.map(frame => frame.token)).toEqual(['ticket-1']); + expect(FakeWebSocket.instances).toHaveLength(1); }); -test('a socket that drops just before keepalive reconnects (not a keepalive handoff)', async () => { - jest.useFakeTimers(); - const getSpy = jest - .spyOn(SupersetClient, 'get') - .mockResolvedValue({} as never); +test('a superseded connect attempt cannot open a socket after a reconfigure', async () => { + const pending = deferTickets(); + connectRealtime(ENABLED); + await flush(); + connectRealtime({ ...ENABLED, WEBSOCKET_URL: 'ws://other:8080/' }); + await flush(); + expect(pending).toHaveLength(2); + + // The older attempt's ticket arrives last and must be ignored. + pending[1]({ token: 'new', expires_in: 60 }); + await flush(); + pending[0]({ token: 'old', expires_in: 60 }); + await flush(); + + expect(FakeWebSocket.instances).toHaveLength(1); + expect(lastSocket().url).toBe('ws://other:8080/'); +}); + +test('an auth ack from an obsolete socket does not authenticate the current one', async () => { + const states: string[] = []; const reasons: string[] = []; + const handler = jest.fn(); + subscribeRealtime('task.status', handler); + const old = await connectAuthenticated(60); + subscribeRealtimeState(state => states.push(state)); subscribeRealtimeOpen(reason => reasons.push(reason)); - try { - connectRealtime({ ...ENABLED, WEBSOCKET_JWT_EXPIRATION_SECONDS: 900 }); - const ws = FakeWebSocket.instances[0]; - ws.readyState = 1; // OPEN - ws.onopen?.(); - expect(reasons).toEqual(['initial']); - - // The socket really drops just before the keepalive (540s) fires; onclose - // schedules a real reconnect at +5s. - jest.advanceTimersByTime(539_000); - ws.readyState = 3; // CLOSED - ws.onclose?.(); - - // The keepalive timer fires while the socket is closed → it must bail (no - // refresh GET, no keepalive handoff), leaving the real reconnect to run. - jest.advanceTimersByTime(2_000); - await Promise.resolve(); - expect(getSpy).not.toHaveBeenCalled(); - // The scheduled reconnect opens with reason 'reconnect' — so list views still - // reconcile the outage rather than treating it as a seamless keepalive. - jest.advanceTimersByTime(5_000); - FakeWebSocket.instances[1].readyState = 1; - FakeWebSocket.instances[1].onopen?.(); - expect(reasons).toEqual(['initial', 'reconnect']); - } finally { - getSpy.mockRestore(); - jest.useRealTimers(); - } + old.readyState = 3; + old.onclose?.({ code: 1006 }); + jest.advanceTimersByTime(5_000); + await flush(); + const current = lastSocket(); + expect(current).not.toBe(old); + current.readyState = 1; + current.onopen?.(); + + // A delayed ack, a close and a message from the dead socket arrive late. + ack(old, 60); + old.onclose?.({ code: 4401 }); + old.onmessage?.({ + data: JSON.stringify({ topic: 'task.status', payload: { stale: true } }), + }); + await flush(); + + expect(states).not.toContain('open'); + expect(reasons).toEqual([]); + expect(handler).not.toHaveBeenCalled(); + expect(FakeWebSocket.instances).toHaveLength(2); + + // The current socket still authenticates normally. + ack(current, 60); + expect(states[states.length - 1]).toBe('open'); + expect(reasons).toEqual(['reconnect']); }); -test('does not schedule a keepalive without a configured token lifetime', () => { - jest.useFakeTimers(); - const getSpy = jest.spyOn(SupersetClient, 'get'); - try { - connectRealtime(ENABLED); // no WEBSOCKET_JWT_EXPIRATION_SECONDS - jest.advanceTimersByTime(60 * 60 * 1000); - expect(getSpy).not.toHaveBeenCalled(); +test.each([ + ['404 from an older web server', { status: 404 }], + ['500', { status: 500 }], + ['502', { status: 502 }], + ['a network error', new TypeError('Failed to fetch')], +])( + '%s is transient: it backs off and retries, never unauthorized', + async (_label, error) => { + getSpy.mockRejectedValue(error); + const states: string[] = []; + subscribeRealtimeState(state => states.push(state)); + + connectRealtime(ENABLED); + await flush(); + jest.advanceTimersByTime(5_000); + await flush(); + + expect(states).not.toContain('unauthorized'); + expect(states).toContain('reconnecting'); + expect(getSpy.mock.calls.length).toBeGreaterThan(1); + + // Once the endpoint appears (the rollout finished) the client connects. + serveTickets(); + jest.advanceTimersByTime(30_000); + await flush(); expect(FakeWebSocket.instances).toHaveLength(1); - } finally { - getSpy.mockRestore(); - jest.useRealTimers(); - } + }, +); + +test('getRealtimeState reflects the latest transition', async () => { + rejectTickets(401); + connectRealtime(ENABLED); + await flush(); + + expect(getRealtimeState()).toBe('unauthorized'); }); diff --git a/superset-frontend/src/middleware/realtime.ts b/superset-frontend/src/middleware/realtime.ts index 46da8110c11f..146c98c3dbf4 100644 --- a/superset-frontend/src/middleware/realtime.ts +++ b/superset-frontend/src/middleware/realtime.ts @@ -30,8 +30,14 @@ * * The transport is strictly best-effort: it is only an acceleration layer over * each feature's own authorized fetch/poll, so a socket outage costs latency, - * not correctness. Connection auth is the `superset-ws-token` JWT cookie, which - * rides the handshake automatically (same-host requirement). + * not correctness. + * + * Authentication is in-band. Every connection attempt first fetches a short-lived + * ticket from `/api/v1/realtime/token/`, opens the socket, and sends it as the + * first `auth` frame; the socket is reported `open` only once the server has + * acknowledged it (`realtime.auth`). Before the ticket expires, a fresh one is + * sent on the open socket, so a long-lived connection never tears down to + * re-authenticate. No credential rides the handshake or the URL. */ import { SupersetClient } from '@superset-ui/core'; import { logging } from '@apache-superset/core/utils'; @@ -50,7 +56,6 @@ type RealtimeHandler = (payload: unknown) => void; type RealtimeConfig = { WEBSOCKET_ENABLE?: boolean; WEBSOCKET_URL?: string; - WEBSOCKET_JWT_EXPIRATION_SECONDS?: number; }; // Reconnect backoff after the socket closes. The socket is a first-class @@ -63,41 +68,51 @@ const RECONNECT_MAX_MS = 30000; // reported `unhealthy` so waiters can stop waiting on it (bounded, rather than a // silent multi-minute hang) instead of assuming a message is merely in flight. const UNHEALTHY_AFTER_ATTEMPTS = 3; -const WS_CONNECTING = 0; const WS_OPEN = 1; -// Fraction of the token lifetime after which the client proactively refreshes the -// channel cookie and reconnects, so the socket rides a fresh token before the -// server terminates it at expiry (see Finding 6 / superset/websocket/channel.py's -// sliding-window re-mint). Past the 0.5 sliding-window boundary (and clear of it, -// so timer jitter can't land exactly on the boundary and miss the re-mint), while -// still leaving ample margin before expiry. -const KEEPALIVE_FRACTION = 0.6; -// An authed, same-origin GET whose response passes through the Flask -// `after_request` hook that (re)mints the ws cookie. Any authed endpoint works; -// `me` is the cheapest stable one. -const COOKIE_REFRESH_ENDPOINT = '/api/v1/me/'; +const TICKET_ENDPOINT = '/api/v1/realtime/token/'; +// Topic of the server's acknowledgement of an `auth` frame. +const AUTH_TOPIC = 'realtime.auth'; +// The server closes with this code when the connection never authenticated, its +// ticket was invalid, or its ticket expired; a fresh ticket fixes it, so the +// client reconnects immediately rather than backing off. +const AUTH_CLOSE_CODE = 4401; +// The server closes with this code when a refresh named another principal. +const FORBIDDEN_CLOSE_CODE = 4403; +// How long to wait for the server to acknowledge the first `auth` frame. +const AUTH_ACK_TIMEOUT_MS = 10000; +// Fraction of the ticket lifetime after which a fresh ticket is sent on the open +// socket, leaving the remaining quarter as margin for retries. +const REFRESH_FRACTION = 0.75; +const REFRESH_RETRY_BASE_MS = 1000; +const REFRESH_RETRY_MAX_MS = 30000; let socket: WebSocket | undefined; let reconnectTimeoutId: number; -let keepaliveTimeoutId: number; -// Token lifetime (ms) from bootstrap config; 0 disables the proactive keepalive. -let tokenLifetimeMs = 0; +let refreshTimeoutId: number; +let authTimeoutId: number; +// When the ticket the server last acknowledged expires (epoch ms). +let ticketExpiresAtMs = 0; let enabled = false; let url: string | undefined; let started = false; -// Bumped on every (re)configuration so an in-flight reconnect scheduled against -// a superseded socket detects it is stale and stops. +// Bumped on every (re)configuration so an in-flight reconnect, ticket fetch, or +// refresh scheduled against a superseded socket detects it is stale and stops. let generation = 0; +// Identifies the latest connection attempt within a generation. Every attempt +// takes a new value, so a ticket fetch, ack, timer or close event belonging to an +// earlier attempt can tell it has been superseded and must not touch the +// current connection. +let attemptSeq = 0; const handlersByTopic = new Map>(); -// Fired every time a socket transitions to OPEN (initial connect and each -// reconnect), so a feature can reconcile state it may have missed while the -// socket was down (see asyncEvent's reconnect catch-up). + // Why a socket became OPEN, passed to open-listeners so a consumer can decide // whether to reconcile: `initial` (first connect), `reconnect` (recovered from a -// drop — a message may have been missed), or `keepalive` (a planned pre-expiry -// token refresh — seamless, so an expensive reconcile isn't warranted). -export type RealtimeOpenReason = 'initial' | 'reconnect' | 'keepalive'; +// drop — a message may have been missed), or `rebind` (the server acknowledged a +// changed tab id on the open socket — a message sent to the old per-tab channel +// may have been missed). Ordinary ticket refreshes happen on the open socket and +// raise no event. +export type RealtimeOpenReason = 'initial' | 'reconnect' | 'rebind'; type RealtimeOpenListener = (reason: RealtimeOpenReason) => void; @@ -107,26 +122,38 @@ const openListeners = new Set(); // Connection health, surfaced so a feature can reconcile via its authorized REST // API on a reconnect and, crucially, stop waiting on a genuinely-down socket -// (`unhealthy`) instead of hanging until a long give-up timeout. +// (`unhealthy`) instead of hanging until a long give-up timeout. `unauthorized` +// is terminal: the ticket endpoint refused this session (logged out, or the +// realtime permission was removed), so the client stops connecting. export type RealtimeConnectionState = | 'connecting' | 'open' | 'reconnecting' - | 'unhealthy'; + | 'unhealthy' + | 'unauthorized'; type RealtimeStateListener = (state: RealtimeConnectionState) => void; const stateListeners = new Set(); +let currentState: RealtimeConnectionState = 'connecting'; + +/** + * The current connection state. Consumers that register after a transition + * (e.g. a chart request started while the transport is `unauthorized`) read it + * here instead of relying on having observed the transition. + */ +export const getRealtimeState = (): RealtimeConnectionState => currentState; + // Consecutive socket closes (or connect failures) with no intervening OPEN. Reset // to 0 on OPEN. Drives the reconnect backoff and the `unhealthy` transition. let reconnectAttempts = 0; /** * Register a listener fired on every connection-state transition - * (`connecting`/`open`/`reconnecting`/`unhealthy`); returns an unsubscribe - * function. Use it to reconcile on reconnect and to give up promptly when the - * socket is `unhealthy` rather than waiting out a long timeout. + * (`connecting`/`open`/`reconnecting`/`unhealthy`/`unauthorized`); returns an + * unsubscribe function. Use it to reconcile on reconnect and to give up promptly + * when the socket is `unhealthy` rather than waiting out a long timeout. */ export const subscribeRealtimeState = ( listener: RealtimeStateListener, @@ -138,6 +165,7 @@ export const subscribeRealtimeState = ( }; const emitState = (state: RealtimeConnectionState): void => { + currentState = state; stateListeners.forEach(listener => { try { listener(state); @@ -178,8 +206,8 @@ const scheduleReconnect = (thisGeneration: number): void => { /** * Register a listener fired when the socket (re)connects (transitions to OPEN); * returns an unsubscribe function. The listener receives a `RealtimeOpenReason` so - * it can skip an expensive reconcile on a seamless `keepalive` reconnect and only - * reconcile on a real `reconnect` after a drop. + * it can skip an expensive reconcile on the `initial` connect and only reconcile + * on a real `reconnect` after a drop. */ export const subscribeRealtimeOpen = ( listener: RealtimeOpenListener, @@ -190,27 +218,23 @@ export const subscribeRealtimeOpen = ( }; }; -/** - * Parse a raw socket message and dispatch the `{topic, payload}` envelope to the - * handlers registered for its topic. Exported for tests. Strictly best-effort: - * malformed data and a throwing handler are logged and swallowed, never - * propagated. - */ -export const dispatchRealtimeMessage = (rawData: string): void => { - let message: RealtimeMessage; +const parseRealtimeMessage = (rawData: string): RealtimeMessage | null => { try { const parsed: unknown = JSON.parse(rawData) ?? {}; - if (!parsed || typeof parsed !== 'object') return; + if (!parsed || typeof parsed !== 'object') return null; const { topic, payload } = parsed as { topic?: unknown; payload?: unknown; }; - if (typeof topic !== 'string') return; - message = { topic, payload }; + if (typeof topic !== 'string') return null; + return { topic, payload }; } catch (err) { logging.warn('Failed to parse realtime message', err); - return; + return null; } +}; + +const deliverRealtimeMessage = (message: RealtimeMessage): void => { const topicHandlers = handlersByTopic.get(message.topic); if (!topicHandlers) return; topicHandlers.forEach(handler => { @@ -222,12 +246,22 @@ export const dispatchRealtimeMessage = (rawData: string): void => { }); }; +/** + * Parse a raw socket message and dispatch the `{topic, payload}` envelope to the + * handlers registered for its topic. Exported for tests. Strictly best-effort: + * malformed data and a throwing handler are logged and swallowed, never + * propagated. + */ +export const dispatchRealtimeMessage = (rawData: string): void => { + const message = parseRealtimeMessage(rawData); + if (message) deliverRealtimeMessage(message); +}; + // Build the per-connection ws URL. The configured endpoint may be absolute // (`ws://host/`) or root-relative (`/superset-ws`, behind a same-origin proxy), // so resolve it against the page URL and normalize an http(s) result to ws(s). -// The tab id is advertised as a query param so the server binds a per-tab channel -// for tab-targeted delivery. The stored base `url` is left untouched so -// connectRealtime's idempotency compare stays stable. +// The stored base `url` is left untouched so connectRealtime's idempotency +// compare stays stable. const buildConnectUrl = (baseUrl: string): string => { try { const base = @@ -235,22 +269,22 @@ const buildConnectUrl = (baseUrl: string): string => { const parsed = new URL(baseUrl, base); if (parsed.protocol === 'http:') parsed.protocol = 'ws:'; else if (parsed.protocol === 'https:') parsed.protocol = 'wss:'; - parsed.searchParams.set('tab_id', getTabId()); return parsed.toString(); } catch { - // Last resort (unparseable and no page URL): still advertise the tab id so - // per-tab delivery works rather than silently falling back to polling. - const sep = baseUrl.includes('?') ? '&' : '?'; - return `${baseUrl}${sep}tab_id=${encodeURIComponent(getTabId())}`; + return baseUrl; } }; -const hasActiveSocket = (): boolean => - socket?.readyState === WS_CONNECTING || socket?.readyState === WS_OPEN; +const hasActiveSocket = (): boolean => socket?.readyState === WS_OPEN; -const teardownSocket = (): void => { +const clearTimers = (): void => { if (reconnectTimeoutId) clearTimeout(reconnectTimeoutId); - if (keepaliveTimeoutId) clearTimeout(keepaliveTimeoutId); + if (refreshTimeoutId) clearTimeout(refreshTimeoutId); + if (authTimeoutId) clearTimeout(authTimeoutId); +}; + +const teardownSocket = (): void => { + clearTimers(); if (socket) { // Detach handlers first so the closing socket can't schedule a reconnect // against the superseded generation. @@ -267,16 +301,140 @@ const teardownSocket = (): void => { } }; +interface Ticket { + token: string; + expiresInMs: number; +} + +/** + * Fetch a fresh connection ticket for the current session. A 401 must not run + * the client's unauthorized handler (a login redirect, or the embedded guest + * handler): the realtime transport is optional, and the rejection is handled by + * `isUnauthorizedError`. + */ +const fetchTicket = async (): Promise => { + const { json } = await SupersetClient.get({ + endpoint: TICKET_ENDPOINT, + ignoreUnauthorized: true, + }); + const { token, expires_in: expiresIn } = (json ?? {}) as { + token?: unknown; + expires_in?: unknown; + }; + if (typeof token !== 'string' || typeof expiresIn !== 'number') { + throw new Error('Malformed realtime ticket response'); + } + return { token, expiresInMs: expiresIn * 1000 }; +}; + +/** + * Whether a ticket request failed because the session may not use the transport + * (no session, or no realtime permission). Retrying cannot succeed. + */ +const isUnauthorizedError = (err: unknown): boolean => { + const status = (err as { status?: unknown } | null)?.status; + return status === 401 || status === 403; +}; + +// Tab ids carried by `auth` frames the server has not acknowledged yet, oldest +// first, and the tab id the server last acknowledged. The server answers every +// frame in order, so comparing the two tells whether an ack rebound the tab. +let unackedTabIds: string[] = []; +let ackedTabId: string | undefined; + +const sendAuthFrame = (ws: WebSocket, token: string): void => { + const tabId = getTabId(); + ws.send(JSON.stringify({ type: 'auth', token, tab_id: tabId })); + unackedTabIds.push(tabId); +}; + +const notifyOpen = (reason: RealtimeOpenReason): void => { + openListeners.forEach(listener => { + try { + listener(reason); + } catch (err) { + logging.warn('Realtime open-listener error', err); + } + }); +}; + +// Stop connecting after the server refused this session. Closes any live socket +// without scheduling a reconnect. +const enterUnauthorized = (): void => { + teardownSocket(); + emitState('unauthorized'); +}; + +const refreshRetryDelayMs = (attempt: number): number => + Math.min(REFRESH_RETRY_BASE_MS * 2 ** attempt, REFRESH_RETRY_MAX_MS); + +// Send a fresh ticket on the open socket `ws`. A transient ticket failure is +// retried with backoff until the current ticket expires; past it the server closes the socket +// with AUTH_CLOSE_CODE and the reconnect path takes over. The server's +// `realtime.auth` reply schedules the next refresh. +function refreshTicket( + thisGeneration: number, + ws: WebSocket, + attempt = 0, +): void { + fetchTicket().then( + ({ token }) => { + if (thisGeneration !== generation || socket !== ws) return; + if (ws.readyState !== WS_OPEN) return; + try { + sendAuthFrame(ws, token); + } catch (err) { + logging.warn('Failed to send realtime auth refresh', err); + } + }, + err => { + if (thisGeneration !== generation || socket !== ws) return; + if (isUnauthorizedError(err)) { + enterUnauthorized(); + return; + } + const delay = refreshRetryDelayMs(attempt); + if (Date.now() + delay >= ticketExpiresAtMs) return; + refreshTimeoutId = window.setTimeout( + () => refreshTicket(thisGeneration, ws, attempt + 1), + delay, + ); + }, + ); +} + // Declared as a hoisted function so scheduleReconnect (above) can reference it // for the mutual reconnect recursion without a use-before-define cycle. -function openSocket(thisGeneration: number, reason: RealtimeOpenReason): void { +async function openSocket( + thisGeneration: number, + reason: RealtimeOpenReason, +): Promise { if (thisGeneration !== generation) return; if (!enabled || !url || typeof WebSocket === 'undefined') return; - const connectUrl = buildConnectUrl(url); + attemptSeq += 1; + const attempt = attemptSeq; + const isCurrent = () => + thisGeneration === generation && attempt === attemptSeq; emitState('connecting'); + + let ticket: Ticket; + try { + ticket = await fetchTicket(); + } catch (err) { + if (!isCurrent()) return; + if (isUnauthorizedError(err)) { + enterUnauthorized(); + return; + } + logging.warn('Failed to fetch realtime ticket', err); + scheduleReconnect(thisGeneration); + return; + } + if (!isCurrent()) return; + let ws: WebSocket; try { - ws = new WebSocket(connectUrl); + ws = new WebSocket(buildConnectUrl(url)); } catch (err) { logging.warn('Failed to open realtime WebSocket', err); // A synchronous constructor throw would otherwise dead-end this generation @@ -287,67 +445,83 @@ function openSocket(thisGeneration: number, reason: RealtimeOpenReason): void { return; } socket = ws; + unackedTabIds = []; + ackedTabId = undefined; + // Set once the server acknowledges the first `auth` frame. + let authenticated = false; - // Proactively refresh the channel cookie and reconnect before this - // connection's token expires, so an idle realtime surface (one making no other - // HTTP requests, e.g. a quiet list view) does not silently lose the socket at - // JWT expiry. The GET re-mints the httponly cookie via the Flask after_request - // hook (inside its sliding window); the reconnect then rides the fresh token. - // Best-effort: reconnect regardless of the GET's outcome (a stale cookie just - // fails the handshake, and onclose retries). This reconnect is `keepalive` (a - // seamless handover), so open-listeners can skip an expensive reconcile. - if (enabled && tokenLifetimeMs > 0) { - keepaliveTimeoutId = window.setTimeout(() => { - // Only hand off a socket that is still OPEN. Besides a superseded - // generation or a replaced socket, guard on readyState: `onclose` merely - // *schedules* a delayed reconnect (it doesn't clear `socket` or bump the - // generation), so a just-closed socket still satisfies `socket === ws`. A - // seamless keepalive handoff of a closed socket would cancel that real - // `reconnect` and reopen as `keepalive` — and list views skip reconcile on - // `keepalive`, so they'd miss entity changes from the actual outage. - if ( - thisGeneration !== generation || - socket !== ws || - ws.readyState !== WS_OPEN - ) { + ws.onopen = () => { + if (thisGeneration !== generation || socket !== ws) return; + try { + sendAuthFrame(ws, ticket.token); + } catch (err) { + logging.warn('Failed to send realtime auth frame', err); + } + // A connection the server never acknowledges (a proxy that swallows frames, + // an overloaded server) is abandoned and retried rather than reported open. + authTimeoutId = window.setTimeout(() => { + if (thisGeneration !== generation || socket !== ws || authenticated) { return; } - SupersetClient.get({ endpoint: COOKIE_REFRESH_ENDPOINT }) - .catch(() => {}) - .finally(() => { - if ( - thisGeneration !== generation || - socket !== ws || - ws.readyState !== WS_OPEN - ) { - return; - } - generation += 1; - teardownSocket(); - openSocket(generation, 'keepalive'); - }); - }, tokenLifetimeMs * KEEPALIVE_FRACTION); - } + teardownSocket(); + scheduleReconnect(thisGeneration); + }, AUTH_ACK_TIMEOUT_MS); + }; - ws.onopen = () => { - if (thisGeneration !== generation) return; + const handleAuthAck = (payload: unknown): void => { + const { status, expires_in: expiresIn } = (payload ?? {}) as { + status?: unknown; + expires_in?: unknown; + }; + if (status !== 'ok' || typeof expiresIn !== 'number') return; + if (refreshTimeoutId) clearTimeout(refreshTimeoutId); + ticketExpiresAtMs = Date.now() + expiresIn * 1000; + refreshTimeoutId = window.setTimeout( + () => refreshTicket(thisGeneration, ws), + expiresIn * 1000 * REFRESH_FRACTION, + ); + const tabId = unackedTabIds.shift(); + if (authenticated) { + if (tabId !== undefined && tabId !== ackedTabId) { + ackedTabId = tabId; + notifyOpen('rebind'); + } + return; + } + authenticated = true; + ackedTabId = tabId; + if (authTimeoutId) clearTimeout(authTimeoutId); // Recovered: reset the failure counter and report a healthy connection. reconnectAttempts = 0; emitState('open'); - openListeners.forEach(listener => { - try { - listener(reason); - } catch (err) { - logging.warn('Realtime open-listener error', err); - } - }); + notifyOpen(reason); }; + ws.onmessage = (event: MessageEvent) => { - if (thisGeneration === generation) - dispatchRealtimeMessage(String(event.data)); + if (thisGeneration !== generation || socket !== ws) return; + const message = parseRealtimeMessage(String(event.data)); + if (!message) return; + if (message.topic === AUTH_TOPIC) { + handleAuthAck(message.payload); + } else if (authenticated) { + deliverRealtimeMessage(message); + } }; - ws.onclose = () => { - if (thisGeneration !== generation) return; + + ws.onclose = (event: CloseEvent) => { + if (thisGeneration !== generation || socket !== ws) return; + clearTimers(); + socket = undefined; + if ( + authenticated && + (event.code === AUTH_CLOSE_CODE || event.code === FORBIDDEN_CLOSE_CODE) + ) { + // The ticket expired (or was refused) on a socket that had been healthy: + // reconnect right away with a fresh ticket instead of backing off. + emitState('reconnecting'); + openSocket(thisGeneration, 'reconnect'); + return; + } scheduleReconnect(thisGeneration); }; // Errors surface as a subsequent close; let onclose own the reconnect and @@ -383,7 +557,6 @@ export const connectRealtime = (config?: RealtimeConfig): void => { started = true; enabled = nextEnabled; url = nextUrl; - tokenLifetimeMs = (conf?.WEBSOCKET_JWT_EXPIRATION_SECONDS ?? 0) * 1000; openSocket(generation, 'initial'); }; @@ -422,13 +595,13 @@ export const subscribeRealtime = ( // Keep the socket's per-tab channel in sync with the tab id. If this tab's id is // reassigned (a duplicated tab resolving a TAB_ID_DENIED collision), the live // socket stays bound to the old per-tab channel while new subscriptions use the -// new id, so tab-targeted status events would miss it. Reconnect to re-register -// under the current id (openSocket re-reads getTabId). No-op unless active. +// new id, so tab-targeted status events would miss it. Re-authenticate on the +// open socket so the server rebinds the tab channel; no reconnect is needed. +// A socket still connecting reads the current id when it sends its first frame. subscribeTabIdChange(() => { - if (!enabled || !hasActiveSocket()) return; - generation += 1; - teardownSocket(); - openSocket(generation, 'reconnect'); + if (!enabled || !socket || socket.readyState !== WS_OPEN) return; + if (refreshTimeoutId) clearTimeout(refreshTimeoutId); + refreshTicket(generation, socket); }); // Test-only: reset module state between cases. @@ -437,6 +610,7 @@ export const resetRealtimeForTests = (): void => { handlersByTopic.clear(); openListeners.clear(); stateListeners.clear(); + currentState = 'connecting'; reconnectAttempts = 0; started = false; enabled = false; diff --git a/superset-frontend/src/views/CRUD/hooks.test.tsx b/superset-frontend/src/views/CRUD/hooks.test.tsx index 191033706749..a1cfc3e67517 100644 --- a/superset-frontend/src/views/CRUD/hooks.test.tsx +++ b/superset-frontend/src/views/CRUD/hooks.test.tsx @@ -737,9 +737,9 @@ test('useListViewResource: realtime reconciles displayed rows on reconnect', asy expect(endpoint).toContain('opr:in'); }); -test('useListViewResource: realtime does NOT reconcile on a keepalive reconnect', async () => { - // A seamless keepalive token refresh must not trigger a full displayed-row - // re-fetch every cycle — only a real `reconnect` (drop recovery) reconciles. +test('useListViewResource: realtime does NOT reconcile on the initial connect', async () => { + // The rows were just fetched, so the first connect must not trigger a full + // displayed-row re-fetch — only a real `reconnect` (drop recovery) reconciles. const initial = [{ id: 1, name: 'A' }]; const getSpy = jest .spyOn(SupersetClient, 'get') @@ -775,9 +775,9 @@ test('useListViewResource: realtime does NOT reconcile on a keepalive reconnect' const callsBefore = getSpy.mock.calls.length; act(() => { - emitRealtimeOpenForTests('keepalive'); + emitRealtimeOpenForTests('initial'); }); - // No reconcile fetch for a keepalive reconnect. + // No reconcile fetch for the initial connect. await new Promise(resolve => { setTimeout(resolve, 50); }); diff --git a/superset-frontend/src/views/CRUD/hooks.ts b/superset-frontend/src/views/CRUD/hooks.ts index addc9aeb3d81..e3bce270c0a9 100644 --- a/superset-frontend/src/views/CRUD/hooks.ts +++ b/superset-frontend/src/views/CRUD/hooks.ts @@ -363,10 +363,8 @@ export function useListViewResource( // Nudges are lossy and not replayed, so after recovering from a dropped // socket refetch the currently-displayed rows once to reconcile anything - // missed while disconnected. Only on a real `reconnect` — not the initial - // connect (rows were just fetched) or a seamless `keepalive` token refresh — - // so a healthy socket doesn't trigger a full re-fetch every keepalive cycle. - // In-place merge, no full-list redraw. + // missed while disconnected. Only on a real `reconnect`, not the initial + // connect (rows were just fetched). In-place merge, no full-list redraw. const unsubscribeOpen = subscribeRealtimeOpen(reason => { if (reason !== 'reconnect') return; patchRows(collectionRef.current.map(rowId)); diff --git a/superset-websocket/README.md b/superset-websocket/README.md index bf1dafc4d52b..a1c7cc30ed6b 100644 --- a/superset-websocket/README.md +++ b/superset-websocket/README.md @@ -32,8 +32,9 @@ events out to JWT-bound socket routing keys. - Redis or Valkey, reachable for Pub/Sub To use realtime push, enable it in the Superset backend (`WEBSOCKET_ENABLE`, -`WEBSOCKET_URL`, `WEBSOCKET_JWT_SECRET`; see below) and run this server on the -same browser-visible host. +`WEBSOCKET_URL`, `WEBSOCKET_JWT_SECRET`; see below) and run this server where +browsers can reach it. It need not share a host with the web application: no +credential travels as a host-scoped cookie. ## Architecture @@ -62,7 +63,7 @@ is from **who** receives it: - **`scope`** — the delivery breadth the server routes by: - `authenticated_global` — broadcast to **every authenticated** realtime socket. This is authenticated-global, **not** public: anonymous users get no realtime - principal, no JWT cookie, and therefore no socket, so they never receive these + principal, no connection ticket, and therefore no socket, so they never receive these messages. (True anonymous/Public-role realtime would need a separate, restricted model and is out of scope.) - `principal` / `tab` — targeted to the routing keys in `routes`. The server @@ -89,7 +90,7 @@ The two topics in use: default (reaching all of a principal's tabs) but a task type may narrow them to a per-tab channel (async chart-data does this), so only the tab watching a task is notified. A socket is bound to its per-tab channel only when the browser - advertises a `tab_id` on the connect URL (see Connection); it is always also + advertises a `tab_id` in its `auth` frame (see Connection); it is always also reachable on its principal channel. The server forwards each browser message as `{topic, payload}` — the browser @@ -100,50 +101,72 @@ before publishing, so a key can never target another principal. ### Connection -When a user's browser connects, it does so over HTTP, including the JWT -authentication cookie set by the Flask app (`WEBSOCKET_JWT_COOKIE_NAME`, default -`superset-ws-token`). _Because authentication is cookie-based, the WebSocket -server must be served from the same browser-visible host as the web -application._ In Kubernetes this is expected to run cleanly as a separate -Deployment/Service behind an ingress that supports WebSocket upgrades. The -server verifies the JWT with the shared secret (`jwtSecret` / `JWT_SECRET`). -Superset mints the token only after the request principal has `can_read` on the -`Realtime` resource. The token carries `aud`, `iss`, `sub`, `principal_type`, -`channel`, and `exp`; the server rejects tokens whose channel does not match the -principal identity. The permission itself is not serialized into the token: a -valid token signed with the websocket secret is the proof that Superset already -authorized the transport. -The socket is bound to the `channel` claim (its principal channel), which is how -principal-grain task-status fanout selects recipient sockets. A principal may -have multiple sockets (e.g. several browser tabs); a message to the principal -channel is sent to all of them. When a browser also advertises a `tab_id` on the -connect URL (`wss://…/?tab_id=`), its socket is additionally bound to a -per-tab channel (`:`, derived from the authorized principal -channel so it can never cross principals), letting the producer target one tab. -Because permission is checked when Superset mints the JWT and when the websocket -server accepts the upgrade, revocation after minting is bounded by the token -lifetime; the Superset default is 15 minutes. To keep a long-lived surface -connected without waiting for a hard disconnect at expiry, Superset re-mints the -cookie inside a sliding window (any request in the second half of the token's -life gets a fresh cookie), and the browser client proactively refreshes the -cookie and reconnects before expiry — so an active or idle-but-open realtime -surface stays connected while revocation remains bounded by the same lifetime. +A browser authenticates in-band with a short-lived **connection ticket**, a JWT +that Superset mints from `GET /api/v1/realtime/token/`. The endpoint returns +`{"token": "", "expires_in": }` and is gated by `can_read` on the +`Realtime` resource (401 without a session, 403 without the permission). The +ticket carries `aud`, `iss`, `sub`, `principal_type`, `channel`, and `exp`; the +server verifies it with the shared secret (`jwtSecret` / `JWT_SECRET`) and +rejects tokens whose channel does not match the principal identity. The +permission itself is not serialized into the ticket: a valid ticket signed with +the websocket secret is the proof that Superset already authorized the transport. +Because the ticket is never sent as a cookie or in the URL, the WebSocket server +can be served from any host the browser can reach, and proxies never log it. In +Kubernetes this is expected to run cleanly as a separate Deployment/Service +behind an ingress that supports WebSocket upgrades. + +The handshake itself carries no credential. The server checks only the `Origin` +allowlist and subscriber health, answering a refusal with a real HTTP status +(403 for a disallowed origin, 503 while the Redis subscriber is unhealthy). +Every connection then moves through two states: + +- **`pending_auth`**: the initial state. The socket is registered on no channel + and is routed no message. It is closed with code `4401` if no valid `auth` + frame arrives within 10 seconds. +- **authenticated**: entered on the first valid `auth` frame, + `{"type": "auth", "token": "", "tab_id": ""}`. The socket is bound to + the ticket's `channel` claim (its principal channel), which is how + principal-grain task-status fanout selects recipient sockets, and, when a valid + `tab_id` is present, to a per-tab channel (`:`, derived from + the authorized principal channel so it can never cross principals). The server + replies `{"topic": "realtime.auth", "payload": {"status": "ok", "expires_in": }}`. + +A principal may have multiple sockets (e.g. several browser tabs); a message to +the principal channel is sent to all of them. + +An authenticated client refreshes its ticket on the open socket by sending +another `auth` frame before the current one expires; there is no reconnect and no +message gap. The server requires the refreshed ticket to name the same `channel` +as the socket's bound one (otherwise it closes the socket with `4403`), extends +the socket's expiry, and sends the same `realtime.auth` reply. A changed `tab_id` +rebinds the per-tab channel. A socket whose ticket expires without a refresh is +closed with `4401`, and a frame carrying an invalid or expired ticket closes the +socket with `4401`, so the client knows to reconnect with a fresh ticket. + +Because permission is checked when Superset mints a ticket, revocation after +minting is bounded by the ticket lifetime (`WEBSOCKET_JWT_EXPIRATION_SECONDS`; +the Superset default is 15 minutes). After a logout or session revocation, an +open socket stays authenticated until its current ticket expires, because the +next refresh fails (the ticket endpoint returns 401 or 403); the server then +closes the socket at expiry. During websocket JWT secret rotation, the websocket service can accept both the current key (`jwtSecret` / `JWT_SECRET`) and one previous verify-only key (`previousJwtSecret` / `PREVIOUS_JWT_SECRET`). The Flask app does not need the -previous key: it keeps minting new cookies with `WEBSOCKET_JWT_SECRET`, and any -cookie signed by the old key is replaced on the next HTTP response. Keep the -previous key configured on the websocket service until old cookies and open +previous key: it keeps minting new tickets with `WEBSOCKET_JWT_SECRET`, and any +ticket signed by the old key is replaced at the client's next refresh. Keep the +previous key configured on the websocket service until old tickets and open sockets have aged out, then remove it. The service is stateless apart from process-local live socket handles. Each replica subscribes to the same Redis Pub/Sub channel, receives the same events, and forwards only to matching sockets connected to that replica. Sticky -sessions are not required; reconnecting to another replica reuses the JWT cookie -and binds the socket to the same principal channel. Short websocket JWT -lifetimes keep connection lifetime bounded for pod recycling; if a pod is -terminated before expiry, connected browsers reconnect and a one-shot catch-up +sessions are not required; reconnecting to another replica presents a fresh +ticket and binds the socket to the same principal channel. Because clients +refresh their ticket on the open socket, the ticket lifetime does not bound +connection lifetime: draining a replica requires closing its sockets explicitly +or terminating the pod. When a pod is terminated, connected browsers reconnect +and a one-shot catch-up fetch on reconnect reconciles any events missed while disconnected (chart-data waiters re-check `/task/status_changes` from their retained cursor; list views refetch their displayed rows). There is no recurring poll while the socket is @@ -156,7 +179,8 @@ The server uses standard WebSocket to detect dead connections. Active sockets are pinged regularly (config: `pingSocketsIntervalMs`) and the internal registry records the last _pong_ timestamp; a socket that has not responded within `socketResponseTimeoutMs` is -terminated. Sockets are also terminated after the JWT `exp` time passes. The +terminated. Sockets are closed with code `4401` after the ticket `exp` time +passes without a refresh. The channel registry is periodically cleaned (`gcChannelsIntervalMs`) to release stale references. @@ -191,7 +215,7 @@ The `Origin` header of each upgrade request must exactly match one of the configured values. When `allowedOrigins` is empty (the default) the check is skipped and any origin is accepted; a single `"*"` entry explicitly allows any origin. Setting this is recommended for production deployments, especially when -the JWT cookie uses `SameSite=None`. +connections can originate from untrusted pages. ### Production hardening checklist @@ -200,16 +224,15 @@ image, review these before exposing it publicly: - **`ALLOWED_ORIGINS`** — set it to your Superset origin(s) (see above). The default (empty) accepts any origin. -- **Secure, scoped JWT cookie** — on the Flask side set - `WEBSOCKET_JWT_COOKIE_SECURE = True` (HTTPS only) and choose an explicit - `WEBSOCKET_JWT_COOKIE_SAMESITE` (`"Lax"` when the app and websocket share a - site; `"None"` only for cross-site, which then *requires* `Secure` and an - `ALLOWED_ORIGINS` allowlist). Scope with `WEBSOCKET_JWT_COOKIE_DOMAIN` if - needed. Cookies are short-lived and JWT-bound, but these flags close the - cross-site surface. +- **TLS** — serve the websocket over `wss://` so the `auth` frame is encrypted + in transit. Tickets are short-lived and never appear in URLs, but they are + readable by page JavaScript, as is the REST API session. - **Connection caps** — the server accepts unlimited connections by default (`0`). Set `maxTotalConnections` and `maxConnectionsPerChannel` (and optionally `maxSocketBufferBytes`) to bound resource use and blunt abuse. + Upgraded sockets that have not yet authenticated count toward + `maxTotalConnections` while they wait out the 10 second auth window, and + `maxPendingConnections` (`MAX_PENDING_CONNECTIONS`) can cap them separately. - **Keepalive/timeouts** — `pingSocketsIntervalMs` / `socketResponseTimeoutMs` govern dead-connection reaping; the defaults are sane, tune per environment. @@ -226,14 +249,7 @@ WEBSOCKET_JWT_SECRET = "= 32 bytes>" The built-in `Gamma` role receives `can_read` on `Realtime`; grant that permission to any additional roles that should receive websocket notifications. Without that permission, Superset masks `WEBSOCKET_ENABLE` to `False` for the -request and does not mint the websocket JWT cookie. - -Note that the WebSocket server must be run on the same hostname (different port) -for the JWT cookie to be shared between the Flask app and the WebSocket server. - -Note also that `localhost` and `127.0.0.1` are not considered the same host. For -example, if you're pointing your browser to `localhost:` for Superset, -then the WebSocket url will need to be configured as `localhost:`. +request and refuses to issue a connection ticket. The following values must be coordinated between the Flask app config and this server's `config.json` (or its environment-variable overrides): @@ -242,7 +258,6 @@ this server's `config.json` (or its environment-variable overrides): | --------------------------- | --------------------------- | ------------------------------------------- | | Current signing/verify key | `WEBSOCKET_JWT_SECRET` | `jwtSecret` / `JWT_SECRET` | | Previous verify-only key | not needed | `previousJwtSecret` / `PREVIOUS_JWT_SECRET` | -| Cookie name | `WEBSOCKET_JWT_COOKIE_NAME` | `jwtCookieName` / `JWT_COOKIE_NAME` | | Realtime channel prefix | `REALTIME_CHANNEL_PREFIX` | `realtimeChannelPrefix` / `REALTIME_CHANNEL_PREFIX` | The Redis connection (`redis` / `REDIS_*`) must point at the same Redis instance @@ -327,8 +342,8 @@ docker run --rm -p 8080:8080 \ ``` Configure it with the same environment variables as the standalone server (see -`src/config.ts`); `JWT_SECRET` / `JWT_COOKIE_NAME` must match the Flask app's -`WEBSOCKET_JWT_SECRET` / `WEBSOCKET_JWT_COOKIE_NAME`, +`src/config.ts`); `JWT_SECRET` must match the Flask app's +`WEBSOCKET_JWT_SECRET`, `PREVIOUS_JWT_SECRET` can hold the prior websocket JWT secret during rotation, and the Redis connection must point at the same instance as `DISTRIBUTED_COORDINATION_CONFIG`. diff --git a/superset-websocket/config.example.json b/superset-websocket/config.example.json index aeb961f4a7eb..407cb7aaad23 100644 --- a/superset-websocket/config.example.json +++ b/superset-websocket/config.example.json @@ -26,6 +26,5 @@ "jwtAlgorithms": ["HS256"], "jwtSecret": "CHANGE-ME", "previousJwtSecret": "", - "jwtCookieName": "superset-ws-token", "allowedOrigins": [] } diff --git a/superset-websocket/config.test.json b/superset-websocket/config.test.json index 07d66dbd6a6d..a82e564dd345 100644 --- a/superset-websocket/config.test.json +++ b/superset-websocket/config.test.json @@ -9,6 +9,5 @@ "globalTags": [] }, "jwtSecret": "test123-test123-test123-test123-test123-test123-test123", - "previousJwtSecret": "old123-old123-old123-old123-old123-old123-old123", - "jwtCookieName": "test-ws-token" + "previousJwtSecret": "old123-old123-old123-old123-old123-old123-old123" } diff --git a/superset-websocket/package-lock.json b/superset-websocket/package-lock.json index 6f635813d32e..49a774832eb9 100644 --- a/superset-websocket/package-lock.json +++ b/superset-websocket/package-lock.json @@ -9,7 +9,6 @@ "version": "0.0.1", "license": "Apache-2.0", "dependencies": { - "cookie": "^2.0.1", "hot-shots": "^17.1.1", "ioredis": "^6.0.0", "jsonwebtoken": "^9.0.3", @@ -2187,19 +2186,6 @@ "node": ">=18" } }, - "node_modules/cookie": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/cookie/-/cookie-2.0.1.tgz", - "integrity": "sha512-yuToqVvRrj6pfDXREyQAAv8SkAEk/8GS3jQRTiUMm66TVtBYmqQeoEjL2Lmq8Rpo6271vH76InTChTitEAm65w==", - "license": "MIT", - "engines": { - "node": ">=22" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/express" - } - }, "node_modules/debug": { "version": "4.4.3", "resolved": "https://registry.npmjs.org/debug/-/debug-4.4.3.tgz", diff --git a/superset-websocket/package.json b/superset-websocket/package.json index c72df2c72019..b9a73e23b9e1 100644 --- a/superset-websocket/package.json +++ b/superset-websocket/package.json @@ -17,7 +17,6 @@ }, "license": "Apache-2.0", "dependencies": { - "cookie": "^2.0.1", "hot-shots": "^17.1.1", "ioredis": "^6.0.0", "jsonwebtoken": "^9.0.3", diff --git a/superset-websocket/spec/auth.test.ts b/superset-websocket/spec/auth.test.ts new file mode 100644 index 000000000000..4c499d0e502a --- /dev/null +++ b/superset-websocket/spec/auth.test.ts @@ -0,0 +1,362 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * End-to-end tests of the in-band authentication state machine over real + * WebSocket connections on a loopback port: pending_auth -> authenticated, + * in-band refresh, and the close codes that tell a client what to do next. + */ +import jwt from 'jsonwebtoken'; +import * as http from 'http'; +import { duplexPair, type Duplex } from 'stream'; +import { afterAll, beforeAll, beforeEach, expect, test, vi } from 'vitest'; +import { WebSocket } from 'ws'; +import config from '../config.test.json' with { type: 'json' }; +import * as server from '../src/index'; + +vi.mock('ioredis', () => ({ + Redis: vi.fn().mockImplementation(function () { + return { subscribe: vi.fn(), on: vi.fn() }; + }), +})); + +type Frame = { topic: string; payload: Record }; + +const signTicket = ( + channel = 'user:5', + expiresInSeconds = 3600, + secret = config.jwtSecret, +) => + jwt.sign( + { + channel, + sub: channel.replace('user:', ''), + principal_type: 'user', + aud: 'superset-websocket', + iss: 'superset', + exp: Math.floor(Date.now() / 1000) + expiresInSeconds, + }, + secret, + ); + +// The server under test is driven over in-memory duplex streams rather than a +// TCP listener: the HTTP parser, the upgrade handshake and the WebSocket framing +// are all real, only the transport is a pipe. +let httpServer: http.Server; +const clients: WebSocket[] = []; +const serverStreams: Duplex[] = []; + +/** Give a stream the socket methods `http` and `ws` call on their transport. */ +const asSocket = (stream: Duplex): Duplex => + Object.assign(stream, { + setTimeout: () => stream, + setNoDelay: () => stream, + setKeepAlive: () => stream, + ref: () => stream, + unref: () => stream, + remoteAddress: '127.0.0.1', + }); + +/** A client socket whose peer is the server under test. */ +const connectToServer = (): Duplex => { + const [clientSide, serverSide] = duplexPair(); + serverStreams.push(serverSide); + httpServer.emit('connection', asSocket(serverSide)); + return asSocket(clientSide); +}; + +const newClient = ( + options: ConstructorParameters[2] = {}, +) => { + const ws = new WebSocket('ws://localhost/', { + createConnection: () => connectToServer() as never, + ...options, + }); + clients.push(ws); + return ws; +}; + +beforeAll(() => { + httpServer = http.createServer(); + httpServer.on('upgrade', server.httpUpgrade); +}); + +beforeEach(async () => { + // Close the previous test's connections on both ends and let their + // server-side close handlers run before resetting, so a late release can't + // decrement the freshly reset pending-connection count. Terminating a client + // alone doesn't reach the server over the in-memory pipe. + clients.splice(0).forEach(client => client.terminate()); + serverStreams.splice(0).forEach(stream => stream.destroy()); + await vi.waitFor(() => expect(server.pendingConnectionCount()).toBe(0)); + server.resetState(); +}); + +afterAll(() => { + clients.forEach(client => client.terminate()); +}); + +/** Open a socket and buffer everything it receives, in order. */ +const open = async () => { + const ws = newClient(); + const frames: Frame[] = []; + const closed = new Promise<{ code: number; reason: string }>(resolve => { + ws.on('close', (code, reason) => + resolve({ code, reason: reason.toString() }), + ); + }); + ws.on('message', data => frames.push(JSON.parse(data.toString()))); + await new Promise((resolve, reject) => { + ws.on('open', resolve); + ws.on('error', reject); + }); + return { ws, frames, closed }; +}; + +const authFrame = (token: string, tabId?: string) => + JSON.stringify({ type: 'auth', token, tab_id: tabId }); + +test('a new connection is pending_auth: registered on no channel', async () => { + const { frames } = await open(); + + await new Promise(resolve => setTimeout(resolve, 50)); + expect(frames).toEqual([]); + expect(server.channels).toEqual({}); + expect(server.sockets).toEqual({}); +}); + +test('routed messages reach no socket until it authenticates', async () => { + const { ws, frames } = await open(); + const nudge = JSON.stringify({ + topic: 'entity.changed', + scope: 'authenticated_global', + payload: { entity_type: 'task', id: 'a' }, + }); + + server.routeRedisMessage('realtime', nudge); + ws.send(authFrame(signTicket())); + await vi.waitFor(() => expect(frames).toHaveLength(1)); + server.routeRedisMessage('realtime', nudge); + await vi.waitFor(() => expect(frames).toHaveLength(2)); + + expect(frames[0].topic).toBe('realtime.auth'); + expect(frames[1].topic).toBe('entity.changed'); +}); + +test('the first valid auth frame binds the principal and tab channels', async () => { + const { ws, frames } = await open(); + + ws.send(authFrame(signTicket('user:5', 600), 'tab3')); + await vi.waitFor(() => expect(frames).toHaveLength(1)); + + expect(frames[0]).toEqual({ + topic: 'realtime.auth', + payload: { status: 'ok', expires_in: expect.any(Number) }, + }); + expect(frames[0].payload.expires_in).toBeGreaterThan(590); + expect(frames[0].payload.expires_in).toBeLessThanOrEqual(600); + expect(server.channels['user:5'].sockets).toHaveLength(1); + expect(server.channels['user:5:tab3'].sockets).toEqual( + server.channels['user:5'].sockets, + ); +}); + +test('a refresh extends the expiry on the open socket without reconnecting', async () => { + const { ws, frames } = await open(); + ws.send(authFrame(signTicket('user:5', 60))); + await vi.waitFor(() => expect(frames).toHaveLength(1)); + const [socketId] = Object.keys(server.sockets); + const before = server.sockets[socketId].identity.tokenExpiresAtMs; + + ws.send(authFrame(signTicket('user:5', 3600))); + await vi.waitFor(() => expect(frames).toHaveLength(2)); + + expect(frames[1].payload.expires_in).toBeGreaterThan(3500); + expect(Object.keys(server.sockets)).toEqual([socketId]); + expect(server.sockets[socketId].identity.tokenExpiresAtMs).toBeGreaterThan( + before, + ); + expect(ws.readyState).toBe(WebSocket.OPEN); +}); + +test('a refresh with a changed tab_id rebinds the tab channel', async () => { + const { ws, frames } = await open(); + ws.send(authFrame(signTicket(), 'old')); + await vi.waitFor(() => expect(frames).toHaveLength(1)); + + ws.send(authFrame(signTicket(), 'new')); + await vi.waitFor(() => expect(frames).toHaveLength(2)); + + expect(server.channels['user:5:old']).toBeUndefined(); + expect(server.channels['user:5:new'].sockets).toEqual( + server.channels['user:5'].sockets, + ); +}); + +test('a refresh without a tab_id keeps the existing tab binding', async () => { + const { ws, frames } = await open(); + ws.send(authFrame(signTicket(), 'tab3')); + await vi.waitFor(() => expect(frames).toHaveLength(1)); + + ws.send(authFrame(signTicket())); + await vi.waitFor(() => expect(frames).toHaveLength(2)); + + expect(server.channels['user:5:tab3'].sockets).toHaveLength(1); +}); + +test('a refresh naming another principal closes the socket with 4403', async () => { + const { ws, frames, closed } = await open(); + ws.send(authFrame(signTicket('user:5'))); + await vi.waitFor(() => expect(frames).toHaveLength(1)); + + ws.send(authFrame(signTicket('user:6'))); + + expect(await closed).toEqual({ code: 4403, reason: 'principal mismatch' }); +}); + +test('a refresh ack over the send buffer cap terminates the socket', async () => { + const { ws, frames } = await open(); + ws.send(authFrame(signTicket())); + await vi.waitFor(() => expect(frames).toHaveLength(1)); + const [socketId] = Object.keys(server.sockets); + const serverWs = server.sockets[socketId].ws; + vi.spyOn(serverWs, 'bufferedAmount', 'get').mockReturnValue(2048); + const sendSpy = vi.spyOn(serverWs, 'send'); + const terminateSpy = vi.spyOn(serverWs, 'terminate'); + server.opts.maxSocketBufferBytes = 1024; + + try { + ws.send(authFrame(signTicket())); + + await vi.waitFor(() => expect(terminateSpy).toHaveBeenCalled()); + expect(sendSpy).not.toHaveBeenCalled(); + expect(server.sockets).toEqual({}); + } finally { + server.opts.maxSocketBufferBytes = 0; + } +}); + +test.each([ + [ + 'signed with an unknown secret', + () => signTicket('user:5', 60, 'x'.repeat(40)), + ], + ['already expired', () => signTicket('user:5', -5)], +])('a ticket %s closes the socket with 4401', async (_label, makeToken) => { + const { ws, closed } = await open(); + + ws.send(authFrame(makeToken())); + + expect(await closed).toEqual({ code: 4401, reason: 'invalid token' }); + expect(server.channels).toEqual({}); +}); + +test('frames that are not auth frames are ignored', async () => { + const { ws, frames } = await open(); + + ws.send('not json'); + ws.send(JSON.stringify({ type: 'ping' })); + ws.send(JSON.stringify({ type: 'auth' })); + ws.send(authFrame(signTicket())); + await vi.waitFor(() => expect(frames).toHaveLength(1)); + + expect(frames[0].topic).toBe('realtime.auth'); +}); + +test('an upgrade from a disallowed origin is answered with 403', async () => { + server.opts.allowedOrigins = ['https://superset.example.com']; + try { + const ws = newClient({ origin: 'https://evil.example' }); + + const status = await new Promise(resolve => { + ws.on('unexpected-response', (_req, res) => resolve(res.statusCode)); + ws.on('error', () => {}); + }); + + expect(status).toBe(403); + } finally { + server.opts.allowedOrigins = []; + } +}); + +test('an upgrade while the subscriber is unhealthy is answered with 503', async () => { + server.markSubscriberUnhealthy('test drop'); + const ws = newClient(); + + const status = await new Promise(resolve => { + ws.on('unexpected-response', (_req, res) => resolve(res.statusCode)); + ws.on('error', () => {}); + }); + + expect(status).toBe(503); +}); + +test('unauthenticated sockets count toward maxTotalConnections', async () => { + server.opts.maxTotalConnections = 1; + try { + await open(); + expect(server.pendingConnectionCount()).toBe(1); + + const second = await open(); + + expect(await second.closed).toMatchObject({ code: 1013 }); + } finally { + server.opts.maxTotalConnections = 0; + } +}); + +test('maxPendingConnections caps unauthenticated sockets separately', async () => { + server.opts.maxPendingConnections = 1; + try { + const first = await open(); + const second = await open(); + expect(await second.closed).toMatchObject({ code: 1013 }); + + // Authenticating releases the reservation for the next connection. + first.ws.send(authFrame(signTicket())); + await vi.waitFor(() => expect(first.frames).toHaveLength(1)); + expect(server.pendingConnectionCount()).toBe(0); + const third = await open(); + third.ws.send(authFrame(signTicket())); + await vi.waitFor(() => expect(third.frames).toHaveLength(1)); + } finally { + server.opts.maxPendingConnections = 0; + } +}); + +test('a pending socket releases its reservation when it closes', async () => { + const { ws, closed } = await open(); + expect(server.pendingConnectionCount()).toBe(1); + + ws.close(); + await closed; + + expect(server.pendingConnectionCount()).toBe(0); +}); + +test('auth is refused when the subscriber dropped after the upgrade', async () => { + const { ws, closed } = await open(); + server.markSubscriberUnhealthy('test drop'); + + ws.send(authFrame(signTicket())); + + expect(await closed).toMatchObject({ code: 1012 }); + expect(server.channels).toEqual({}); + expect(server.sockets).toEqual({}); +}); diff --git a/superset-websocket/spec/index.test.ts b/superset-websocket/spec/index.test.ts index db376996638f..611940017883 100644 --- a/superset-websocket/spec/index.test.ts +++ b/superset-websocket/spec/index.test.ts @@ -76,23 +76,17 @@ const signRealtimeToken = ( secret = config.jwtSecret, ) => jwt.sign(realtimeClaims(overrides), secret); -/** - * Build an upgrade-shaped HTTP request, optionally carrying the JWT cookie and - * an `Origin` header. - */ +/** Build an upgrade-shaped HTTP request, optionally carrying an `Origin` header. */ const makeRequest = ({ - token, origin, url = 'http://localhost', }: { - token?: string; origin?: string; url?: string; } = {}): http.IncomingMessage => { const request = new http.IncomingMessage(new net.Socket()); request.method = 'GET'; request.url = url; - if (token) request.headers.cookie = `${config.jwtCookieName}=${token}`; if (origin) request.headers.origin = origin; return request; }; @@ -125,9 +119,23 @@ const trackSocket = ( ...overrides, }); -/** Run the full accept path for `token`: verify at upgrade, then connect. */ -const connect = (ws: WebSocket, token: string) => - server.wsConnection(ws, server.readSocketIdentity(makeRequest({ token }))); +/** Deliver a raw inbound frame to the `message` handler `wsConnection` registered. */ +const receive = (ws: WebSocket, frame: unknown) => { + // Automocked instances share one `on` mock, so take the latest registration. + const call = (ws.on as unknown as Mock).mock.calls + .filter(([event]) => event === 'message') + .at(-1); + expect(call).toBeDefined(); + (call![1] as (data: Buffer) => void)( + Buffer.from(typeof frame === 'string' ? frame : JSON.stringify(frame)), + ); +}; + +/** Accept a connection, then authenticate it with an `auth` frame for `token`. */ +const connect = (ws: WebSocket, token: string, tabId?: string) => { + server.wsConnection(ws); + receive(ws, { type: 'auth', token, tab_id: tabId }); +}; describe('server', () => { let statsdIncrementMock: Mock; @@ -475,11 +483,11 @@ describe('server', () => { }); }); - describe('readSocketIdentity', () => { + describe('readTokenIdentity', () => { test('extracts the identity a valid token proves', () => { const token = signRealtimeToken(); - expect(server.readSocketIdentity(makeRequest({ token }))).toEqual({ + expect(server.readTokenIdentity(token)).toEqual({ channel: channelId, principalType: 'user', subject: '5', @@ -490,15 +498,7 @@ describe('server', () => { test('accepts a token signed with the previous secret', () => { const token = signRealtimeToken({}, config.previousJwtSecret); - expect(server.readSocketIdentity(makeRequest({ token })).channel).toEqual( - channelId, - ); - }); - - test('rejects a request with no JWT cookie', () => { - expect(() => server.readSocketIdentity(makeRequest())).toThrow( - 'JWT not present', - ); + expect(server.readTokenIdentity(token).channel).toEqual(channelId); }); test.each([ @@ -532,9 +532,7 @@ describe('server', () => { ) => { const token = signRealtimeToken(overrides, secret); - expect(() => - server.readSocketIdentity(makeRequest({ token })), - ).toThrow(); + expect(() => server.readTokenIdentity(token)).toThrow(); }, ); }); @@ -599,11 +597,14 @@ describe('server', () => { // per-tab channel. A message targeting user:5:tabA reaches only tab A. const wsA = new wsMock('localhost'); const sendA = vi.spyOn(wsA, 'send'); - server.wsConnection(wsA, makeIdentity('user:5'), 'tabA'); + connect(wsA, signRealtimeToken(), 'tabA'); const wsB = new wsMock('localhost'); const sendB = vi.spyOn(wsB, 'send'); - server.wsConnection(wsB, makeIdentity('user:5'), 'tabB'); + connect(wsB, signRealtimeToken(), 'tabB'); + // Discard the `realtime.auth` replies; only routed messages matter here. + sendA.mockClear(); + sendB.mockClear(); server.routeRedisMessage('realtime', taskStatus(['user:5:tabA'])); @@ -666,7 +667,8 @@ describe('server', () => { // receive a broadcast exactly once (dedup by socket, not by channel). const ws = new wsMock('localhost'); const send = vi.spyOn(ws, 'send'); - server.wsConnection(ws, makeIdentity('user:5'), 'tabA'); + connect(ws, signRealtimeToken(), 'tabA'); + send.mockClear(); server.routeRedisMessage( 'realtime', @@ -865,7 +867,7 @@ describe('server', () => { }); test('a tab_id also binds the socket to its per-tab channel', async () => { - server.wsConnection(ws, makeIdentity(channelId), 'tabA'); + connect(ws, signRealtimeToken(), 'tabA'); // One socket id, indexed under both the principal and the per-tab channel. const principalSockets = server.channels[channelId].sockets; @@ -877,6 +879,40 @@ describe('server', () => { expect(server.sockets[socketId].tabChannel).toEqual(`${channelId}:tabA`); }); + test('a connection that never authenticates is closed after the auth window', () => { + vi.useFakeTimers(); + try { + const closeSpy = vi.spyOn(ws, 'close'); + server.wsConnection(ws); + + vi.advanceTimersByTime(server.AUTH_TIMEOUT_MS - 1); + expect(closeSpy).not.toHaveBeenCalled(); + vi.advanceTimersByTime(1); + + expect(closeSpy).toHaveBeenCalledWith( + server.AUTH_CLOSE_CODE, + 'authentication timeout', + ); + expect(server.sockets).toEqual({}); + } finally { + vi.useRealTimers(); + } + }); + + test('an authenticated connection is not closed by the auth window', () => { + vi.useFakeTimers(); + try { + const closeSpy = vi.spyOn(ws, 'close'); + connect(ws, signRealtimeToken()); + + vi.advanceTimersByTime(server.AUTH_TIMEOUT_MS * 2); + + expect(closeSpy).not.toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + test('unsolicited pong payload cannot pollute Object.prototype', async () => { connect(ws, signRealtimeToken()); @@ -1021,11 +1057,13 @@ describe('server', () => { describe('httpUpgrade', () => { let socket: net.Socket; let socketDestroySpy: Mock; + let socketWriteSpy: Mock; let wssUpgradeSpy: Mock; beforeEach(() => { socket = new net.Socket(); socketDestroySpy = vi.spyOn(socket, 'destroy'); + socketWriteSpy = vi.spyOn(socket, 'write'); wssUpgradeSpy = vi.spyOn(server.wss, 'handleUpgrade'); }); @@ -1033,104 +1071,33 @@ describe('server', () => { wssUpgradeSpy.mockRestore(); }); - test('invalid JWT', async () => { - const token = signRealtimeToken({}, 'invalid secret'); - - server.httpUpgrade(makeRequest({ token }), socket, Buffer.alloc(5)); - - expect(socketDestroySpy).toHaveBeenCalled(); - expect(wssUpgradeSpy).not.toHaveBeenCalled(); - // rejected upgrades are counted for auditability - expect(statsdIncrementMock).toHaveBeenCalledWith('ws_upgrade_rejected'); - }); - - test('valid JWT, no channel', async () => { - const token = signRealtimeToken({ channel: undefined }); - - server.httpUpgrade(makeRequest({ token }), socket, Buffer.alloc(5)); - - expect(socketDestroySpy).toHaveBeenCalled(); - expect(wssUpgradeSpy).not.toHaveBeenCalled(); - }); - - test('valid upgrade', async () => { - const token = signRealtimeToken(); - - server.httpUpgrade(makeRequest({ token }), socket, Buffer.alloc(5)); + test('accepts the upgrade without inspecting any credential', () => { + server.httpUpgrade(makeRequest(), socket, Buffer.alloc(5)); expect(socketDestroySpy).not.toHaveBeenCalled(); expect(wssUpgradeSpy).toHaveBeenCalled(); }); - test('valid upgrade with previous JWT secret', async () => { - const token = signRealtimeToken({}, config.previousJwtSecret); - - server.httpUpgrade(makeRequest({ token }), socket, Buffer.alloc(5)); - - expect(socketDestroySpy).not.toHaveBeenCalled(); - expect(wssUpgradeSpy).toHaveBeenCalled(); - }); - - test('an accepted upgrade tracks the socket without re-verifying the JWT', async () => { - const token = signRealtimeToken(); + test('an accepted upgrade registers nothing until it authenticates', () => { const ws = new wsMock('localhost'); - // Hand the upgrade callback a socket, as the real handshake would, but - // strip the cookie from the request first: the identity must already have - // been carried over from the single upgrade-time verification. - const request = makeRequest({ token }); wssUpgradeSpy.mockImplementation( (_request, _socket, _head, callback: (ws: WebSocket) => void) => { - request.headers.cookie = ''; callback(ws); }, ); - server.httpUpgrade(request, socket, Buffer.alloc(5)); - - expect(server.channels[channelId].sockets).toHaveLength(1); - }); - - test('an accepted upgrade binds the per-tab channel from ?tab_id', async () => { - const token = signRealtimeToken(); - const ws = new wsMock('localhost'); - const request = makeRequest({ - token, - url: 'http://localhost/?tab_id=tabA', - }); - wssUpgradeSpy.mockImplementation( - (_request, _socket, _head, callback: (ws: WebSocket) => void) => { - callback(ws); - }, - ); - - server.httpUpgrade(request, socket, Buffer.alloc(5)); - - expect(server.channels[channelId].sockets).toHaveLength(1); - expect(server.channels[`${channelId}:tabA`].sockets).toEqual( - server.channels[channelId].sockets, - ); - }); - - test('an accepted upgrade ignores an invalid ?tab_id', async () => { - // A tab id outside the allowed length/charset must not become a channel - // key; the socket falls back to principal-grain (no per-tab channel). - const token = signRealtimeToken(); - const ws = new wsMock('localhost'); - const badTabId = 'a'.repeat(65); - const request = makeRequest({ - token, - url: `http://localhost/?tab_id=${badTabId}`, - }); - wssUpgradeSpy.mockImplementation( - (_request, _socket, _head, callback: (ws: WebSocket) => void) => { - callback(ws); - }, + server.httpUpgrade( + makeRequest({ url: 'http://localhost/?tab_id=tabA&token=ignored' }), + socket, + Buffer.alloc(5), ); - server.httpUpgrade(request, socket, Buffer.alloc(5)); - + expect(server.channels).toEqual({}); + expect(server.sockets).toEqual({}); + receive(ws, { type: 'auth', token: signRealtimeToken() }); expect(server.channels[channelId].sockets).toHaveLength(1); - expect(server.channels[`${channelId}:${badTabId}`]).toBeUndefined(); + // The tab id is only honored from the auth frame, never the URL. + expect(server.channels[`${channelId}:tabA`]).toBeUndefined(); }); describe('origin validation', () => { @@ -1140,23 +1107,22 @@ describe('server', () => { test('rejects upgrade from a disallowed origin', () => { server.opts.allowedOrigins = ['https://superset.example.com']; - const token = signRealtimeToken(); - server.httpUpgrade( - makeRequest({ token, origin: 'https://evil.example' }), + makeRequest({ origin: 'https://evil.example' }), socket, Buffer.alloc(5), ); + expect(socketWriteSpy).toHaveBeenCalledWith( + expect.stringContaining('403 Forbidden'), + ); expect(socketDestroySpy).toHaveBeenCalled(); expect(wssUpgradeSpy).not.toHaveBeenCalled(); }); test('rejects upgrade with no origin when an allowlist is set', () => { server.opts.allowedOrigins = ['https://superset.example.com']; - const token = signRealtimeToken(); - - server.httpUpgrade(makeRequest({ token }), socket, Buffer.alloc(5)); + server.httpUpgrade(makeRequest(), socket, Buffer.alloc(5)); expect(socketDestroySpy).toHaveBeenCalled(); expect(wssUpgradeSpy).not.toHaveBeenCalled(); @@ -1164,10 +1130,8 @@ describe('server', () => { test('allows upgrade from an allowed origin', () => { server.opts.allowedOrigins = ['https://superset.example.com']; - const token = signRealtimeToken(); - server.httpUpgrade( - makeRequest({ token, origin: 'https://superset.example.com' }), + makeRequest({ origin: 'https://superset.example.com' }), socket, Buffer.alloc(5), ); @@ -1210,11 +1174,7 @@ describe('server', () => { const writeSpy = vi.spyOn(socket, 'write'); const upgradeSpy = vi.spyOn(server.wss, 'handleUpgrade'); - server.httpUpgrade( - makeRequest({ token: signRealtimeToken() }), - socket, - Buffer.alloc(5), - ); + server.httpUpgrade(makeRequest(), socket, Buffer.alloc(5)); expect(writeSpy).toHaveBeenCalledWith(expect.stringContaining('503')); expect(destroySpy).toHaveBeenCalled(); @@ -1231,11 +1191,7 @@ describe('server', () => { const destroySpy = vi.spyOn(socket, 'destroy'); const upgradeSpy = vi.spyOn(server.wss, 'handleUpgrade'); - server.httpUpgrade( - makeRequest({ token: signRealtimeToken() }), - socket, - Buffer.alloc(5), - ); + server.httpUpgrade(makeRequest(), socket, Buffer.alloc(5)); expect(destroySpy).not.toHaveBeenCalled(); expect(upgradeSpy).toHaveBeenCalled(); @@ -1283,11 +1239,13 @@ describe('server', () => { let ws: WebSocket; let pingSpy: Mock; let terminateSpy: Mock; + let closeSpy: Mock; beforeEach(() => { ws = new wsMock('localhost'); pingSpy = vi.spyOn(ws, 'ping'); terminateSpy = vi.spyOn(ws, 'terminate'); + closeSpy = vi.spyOn(ws, 'close'); }); test('active sockets', () => { @@ -1301,7 +1259,7 @@ describe('server', () => { expect(Object.keys(server.sockets).length).toBe(1); }); - test('sockets with expired JWTs are terminated', () => { + test('sockets with expired JWTs are closed with a reauthentication code', () => { vi.spyOn(ws, 'readyState', 'get').mockReturnValue(WebSocket.OPEN); trackSocket(channelId, ws, { identity: makeIdentity(channelId, Date.now() - 1), @@ -1310,7 +1268,8 @@ describe('server', () => { server.checkSockets(); expect(pingSpy).not.toHaveBeenCalled(); - expect(terminateSpy).toHaveBeenCalled(); + expect(closeSpy).toHaveBeenCalledWith(4401, 'token expired'); + expect(terminateSpy).not.toHaveBeenCalled(); expect(Object.keys(server.sockets).length).toBe(0); expect(statsdIncrementMock).toHaveBeenCalledWith( 'ws_token_expired_disconnect', diff --git a/superset-websocket/src/config.ts b/superset-websocket/src/config.ts index c6c59618ca30..8d1246be4217 100644 --- a/superset-websocket/src/config.ts +++ b/superset-websocket/src/config.ts @@ -68,7 +68,6 @@ type ConfigType = { jwtAlgorithms: string[]; jwtSecret: string; previousJwtSecret: string; - jwtCookieName: string; realtimeChannelPrefix: string; allowedOrigins: string[]; socketResponseTimeoutMs: number; @@ -77,6 +76,7 @@ type ConfigType = { maxSocketBufferBytes: number; maxConnectionsPerChannel: number; maxTotalConnections: number; + maxPendingConnections: number; }; function defaultConfig(): ConfigType { @@ -88,7 +88,6 @@ function defaultConfig(): ConfigType { jwtAlgorithms: ['HS256'], jwtSecret: '', previousJwtSecret: '', - jwtCookieName: 'superset-ws-token', realtimeChannelPrefix: '', allowedOrigins: [], socketResponseTimeoutMs: 60 * 1000, @@ -100,6 +99,7 @@ function defaultConfig(): ConfigType { // 0 disables the limit (unlimited); set a positive value to opt in. maxConnectionsPerChannel: 0, maxTotalConnections: 0, + maxPendingConnections: 0, statsd: { host: '127.0.0.1', port: 8125, @@ -210,7 +210,6 @@ function applyEnvOverrides(config: ConfigType): ConfigType { LOG_FILENAME: val => (config.logFilename = val), JWT_SECRET: val => (config.jwtSecret = val), PREVIOUS_JWT_SECRET: val => (config.previousJwtSecret = val), - JWT_COOKIE_NAME: val => (config.jwtCookieName = val), REALTIME_CHANNEL_PREFIX: val => (config.realtimeChannelPrefix = val), ALLOWED_ORIGINS: val => (config.allowedOrigins = toStringArray(val)), SOCKET_RESPONSE_TIMEOUT_MS: val => @@ -243,6 +242,11 @@ function applyEnvOverrides(config: ConfigType): ConfigType { val, config.maxTotalConnections, )), + MAX_PENDING_CONNECTIONS: val => + (config.maxPendingConnections = toNonNegativeNumber( + val, + config.maxPendingConnections, + )), REDIS_HOST: val => (config.redis.host = val), REDIS_PORT: val => (config.redis.port = toNonNegativeNumber(val, config.redis.port)), diff --git a/superset-websocket/src/index.ts b/superset-websocket/src/index.ts index 289e68cc0356..59e6887f5d86 100644 --- a/superset-websocket/src/index.ts +++ b/superset-websocket/src/index.ts @@ -26,7 +26,6 @@ import jwt, { type Algorithm, type JwtPayload as JsonWebTokenPayload, } from 'jsonwebtoken'; -import { parseCookie } from 'cookie'; import { Redis, RedisOptions } from 'ioredis'; import StatsD from 'hot-shots'; @@ -99,7 +98,7 @@ export interface SocketInstance { ws: WebSocket; channel: string; // Additional per-tab channel the socket is also registered under, when the - // browser advertised a `tab_id` at connect (`:`). + // browser advertised a `tab_id` in an `auth` frame (`:`). tabChannel?: string; identity: SocketIdentity; pongTs: number; @@ -242,9 +241,14 @@ redisSubscriber.on('error', (err: Error) => { logger.error(`Redis connection error: ${err.message}`); }); const httpServer = http.createServer(); +// Upper bound on an inbound frame. The only frame the server accepts is a small +// `auth` message, so anything larger is rejected by the protocol layer. +const MAX_INBOUND_PAYLOAD_BYTES = 16 * 1024; + export const wss = new WebSocketServer({ noServer: true, clientTracking: false, + maxPayload: MAX_INBOUND_PAYLOAD_BYTES, }); const SOCKET_ACTIVE_STATES: number[] = [WebSocket.OPEN, WebSocket.CONNECTING]; @@ -281,6 +285,22 @@ export let sockets: Record = {}; // connection limit has been reached (1013 = "Try Again Later"). const CONNECTION_LIMIT_CLOSE_CODE = 1013; +// Application close codes of the in-band authentication protocol. A client that +// receives AUTH_CLOSE_CODE reconnects immediately with a fresh ticket (the +// connection never authenticated, its ticket was invalid, or the ticket expired +// without a refresh). FORBIDDEN_CLOSE_CODE means a refresh tried to switch the +// socket to another principal. +export const AUTH_CLOSE_CODE = 4401; +export const FORBIDDEN_CLOSE_CODE = 4403; + +// A new connection is `pending_auth`: it is registered on no channel and receives +// no message until a valid `auth` frame arrives. It is closed with +// AUTH_CLOSE_CODE if none does within this window. +export const AUTH_TIMEOUT_MS = 10000; + +// Topic of the reply to an accepted `auth` frame. +const AUTH_REPLY_TOPIC = 'realtime.auth'; + // The Redis Pub/Sub subscriber is the server's only source of realtime messages. // If it drops, a browser socket can stay open while the server silently misses // messages, so subscriber health gates the transport: while unhealthy, new @@ -348,6 +368,12 @@ export const isSocketActive = (socketId: string): boolean => { export const activeSocketCount = (): number => Object.keys(sockets).filter(isSocketActive).length; +// Connections that have upgraded but not yet authenticated. They hold a server +// socket for up to AUTH_TIMEOUT_MS, so they count toward `maxTotalConnections` +// and are bounded by `maxPendingConnections`. +let pendingConnections = 0; +export const pendingConnectionCount = (): number => pendingConnections; + /** * Counts the active sockets currently registered on the given channel. */ @@ -369,7 +395,10 @@ export const activeChannelSocketCount = (channel: string): number => export const connectionLimitReason = (channel: string): string | null => { const { maxTotalConnections, maxConnectionsPerChannel } = opts; - if (maxTotalConnections > 0 && activeSocketCount() >= maxTotalConnections) { + if ( + maxTotalConnections > 0 && + activeSocketCount() + pendingConnections >= maxTotalConnections + ) { return `total connection limit (${maxTotalConnections}) reached`; } @@ -644,7 +673,7 @@ redisSubscriber.on('ready', () => { }); /** - * Verify and parse a realtime JWT cookie from an HTTP request. + * Verify and parse a realtime JWT ticket. */ const verifyRealtimeJwt = (token: string): RealtimeJwtPayload => { let lastError: unknown; @@ -670,17 +699,11 @@ const verifyRealtimeJwt = (token: string): RealtimeJwtPayload => { }; /** - * Verify a realtime JWT cookie on an HTTP request and extract the socket - * identity it proves. Throws when the token is absent, unverifiable, or its - * claims do not describe a routable principal. + * Verify a realtime JWT ticket and extract the socket identity it proves. Throws + * when the token is unverifiable, or its claims do not describe a routable + * principal. */ -export const readSocketIdentity = ( - request: http.IncomingMessage, -): SocketIdentity => { - const cookies = parseCookie(request.headers.cookie || ''); - const token = cookies[opts.jwtCookieName]; - - if (!token) throw new Error('JWT not present'); +export const readTokenIdentity = (token: string): SocketIdentity => { const jwtPayload = verifyRealtimeJwt(token); const channelId = jwtPayload.channel; const subject = jwtPayload.sub; @@ -715,60 +738,233 @@ export const readSocketIdentity = ( }; }; +/** An inbound `auth` frame: `{type: 'auth', token, tab_id?}`. */ +interface AuthFrame { + token: string; + tabId?: string; +} + /** - * WebSocket `connection` event handler. `identity` is the one already verified - * during the HTTP upgrade (see `httpUpgrade`), so each connection's JWT is - * verified exactly once. + * Parse a raw inbound frame into an `auth` frame, or `null` when it is anything + * else (malformed JSON, another message type, or a missing token). An invalid or + * absent `tab_id` is dropped rather than failing the frame, so the socket falls + * back to principal-grain delivery. */ -export const wsConnection = ( - ws: WebSocket, - identity: SocketIdentity, - tabId?: string, -) => { - const { channel } = identity; - - // Refuse the connection if a configured connection limit has been reached, - // before tracking it against the internal registries. The cap is keyed on the - // principal channel (which counts all of a principal's tabs); the per-tab - // channel is a delivery refinement and is not separately limit-checked. - const limitReason = connectionLimitReason(channel); - if (limitReason) { - statsd.increment('ws_connection_rejected'); - logger.warn(`Refusing connection on channel ${channel}: ${limitReason}`); - ws.close(CONNECTION_LIMIT_CLOSE_CODE, limitReason); - return; +const parseAuthFrame = (raw: WebSocket.RawData): AuthFrame | null => { + let parsed: unknown; + try { + parsed = JSON.parse(raw.toString()); + } catch { + return null; } + if (!parsed || typeof parsed !== 'object') return null; + const { + type, + token, + tab_id: tabId, + } = parsed as { + type?: unknown; + token?: unknown; + tab_id?: unknown; + }; + if (type !== 'auth' || typeof token !== 'string' || token.length === 0) { + return null; + } + return { + token, + tabId: typeof tabId === 'string' && isValidTabId(tabId) ? tabId : undefined, + }; +}; - // When the browser advertised a tab id at connect, also register the socket - // under a per-tab channel derived from its authorized principal channel - // (`:`). The tab id is client-supplied but always prefixed by - // the authorized principal channel, so it can never address another - // principal's sockets. Task-status can then target one tab, while the - // principal channel still reaches all of a principal's tabs. - const tabChannel = tabId ? `${channel}:${tabId}` : undefined; - - const socketInstance: SocketInstance = { - ws, - channel, - tabChannel, - identity, - pongTs: Date.now(), +// Sent through the bounded send path: a client replaying auth frames without +// reading replies hits the same buffer cap as any other outbound message. +const sendAuthOk = (socketId: string, identity: SocketIdentity): void => { + const message: OutboundMessage = { + topic: AUTH_REPLY_TOPIC, + payload: { + status: 'ok', + expires_in: Math.max( + 0, + Math.floor((identity.tokenExpiresAtMs - Date.now()) / 1000), + ), + }, }; + sendToSocket(socketId, JSON.stringify(message)); +}; + +/** + * Removes a socket id from a channel's socket list, dropping the channel when it + * becomes empty. + */ +const untrackFromChannel = (channel: string, socketId: string): void => { + const entry = channels[channel]; + if (!entry) return; + entry.sockets = entry.sockets.filter(id => id !== socketId); + if (entry.sockets.length === 0) delete channels[channel]; +}; - // Register once (mints a single socket id), then also index that same id under - // the per-tab channel - never call trackClient twice (it would mint a second - // id and a duplicate registry entry). - const socketId = trackClient(channel, socketInstance); +/** + * Binds an already tracked socket to the per-tab channel derived from its + * authorized principal channel (`:`), replacing a previous + * binding. The tab id is client-supplied but always prefixed by the authorized + * principal channel, so it can never address another principal's sockets. + * Task-status can then target one tab, while the principal channel still reaches + * all of a principal's tabs. + */ +const bindTabChannel = ( + socketId: string, + socketInstance: SocketInstance, + tabId: string | undefined, +): void => { + const tabChannel = tabId ? `${socketInstance.channel}:${tabId}` : undefined; + if (tabChannel === socketInstance.tabChannel) return; + if (socketInstance.tabChannel) { + untrackFromChannel(socketInstance.tabChannel, socketId); + } if (tabChannel) { (channels[tabChannel] ??= { sockets: [] }).sockets.push(socketId); } - logger.debug( - `socket ${socketId} connected on channel ${channel}` + - (tabChannel ? ` (tab channel ${tabChannel})` : ''), - ); + socketInstance.tabChannel = tabChannel; +}; + +/** + * WebSocket `connection` event handler. The connection starts `pending_auth`: + * it is tracked in no registry and routed no message until its first `auth` frame + * carries a valid ticket, at which point it is registered on its principal channel + * (and tab channel). Later `auth` frames refresh the ticket on the open socket. + */ +export const wsConnection = (ws: WebSocket) => { + // Set once the first `auth` frame is accepted; `undefined` while pending_auth. + let socketId: string | undefined; + + // Reserve capacity at upgrade time so unauthenticated sockets cannot bypass + // the configured caps while they wait out the auth window. + const { maxPendingConnections } = opts; + const limitReason = connectionLimitReason(''); + if ( + limitReason || + (maxPendingConnections > 0 && pendingConnections >= maxPendingConnections) + ) { + statsd.increment('ws_connection_rejected'); + logger.warn( + `Refusing connection: ${limitReason ?? 'pre-auth connection limit reached'}`, + ); + ws.close(CONNECTION_LIMIT_CLOSE_CODE, limitReason ?? 'too many pending'); + return; + } + pendingConnections += 1; + let pendingReleased = false; + const releasePending = () => { + if (pendingReleased) return; + pendingReleased = true; + pendingConnections -= 1; + }; + + const authTimer = setTimeout(() => { + if (socketId === undefined) { + statsd.increment('ws_auth_timeout'); + releasePending(); + ws.close(AUTH_CLOSE_CODE, 'authentication timeout'); + } + }, AUTH_TIMEOUT_MS); + ws.on('close', () => { + clearTimeout(authTimer); + releasePending(); + }); + + const authenticate = (frame: AuthFrame, identity: SocketIdentity) => { + const { channel } = identity; + + // The socket stops being pending here, so it no longer counts against the + // pending reservation when the real limits are evaluated below. + releasePending(); + + // The subscriber may have dropped between the upgrade and this frame; a + // socket registered now would miss messages and no recovery would close it. + if (!subscriberHealthy) { + statsd.increment('ws_upgrade_rejected'); + ws.close(SUBSCRIBER_UNHEALTHY_CLOSE_CODE, 'realtime backend unavailable'); + return; + } + + // Refuse the connection if a configured connection limit has been reached, + // before tracking it against the internal registries. The cap is keyed on the + // principal channel (which counts all of a principal's tabs); the per-tab + // channel is a delivery refinement and is not separately limit-checked. + const limitReason = connectionLimitReason(channel); + if (limitReason) { + statsd.increment('ws_connection_rejected'); + logger.warn(`Refusing connection on channel ${channel}: ${limitReason}`); + ws.close(CONNECTION_LIMIT_CLOSE_CODE, limitReason); + return; + } + + const socketInstance: SocketInstance = { + ws, + channel, + identity, + pongTs: Date.now(), + }; + // Register once (mints a single socket id), then also index that same id + // under the per-tab channel - never call trackClient twice (it would mint a + // second id and a duplicate registry entry). + socketId = trackClient(channel, socketInstance); + bindTabChannel(socketId, socketInstance, frame.tabId); + clearTimeout(authTimer); + logger.debug( + `socket ${socketId} authenticated on channel ${channel}` + + (socketInstance.tabChannel + ? ` (tab channel ${socketInstance.tabChannel})` + : ''), + ); + sendAuthOk(socketId, identity); + }; + + const refresh = ( + id: string, + socketInstance: SocketInstance, + frame: AuthFrame, + identity: SocketIdentity, + ) => { + if (identity.channel !== socketInstance.channel) { + statsd.increment('ws_auth_principal_switch_rejected'); + logger.warn( + `Closing socket ${id}: refresh named a different principal channel`, + ); + ws.close(FORBIDDEN_CLOSE_CODE, 'principal mismatch'); + return; + } + socketInstance.identity = identity; + if (frame.tabId !== undefined) { + bindTabChannel(id, socketInstance, frame.tabId); + } + sendAuthOk(id, identity); + }; + + ws.on('message', (raw: WebSocket.RawData) => { + const frame = parseAuthFrame(raw); + if (!frame) return; + let identity: SocketIdentity; + try { + identity = readTokenIdentity(frame.token); + } catch (err) { + // Record why (never the token itself) so rejected attempts are auditable. + statsd.increment('ws_auth_rejected'); + logger.warn(`Rejected WebSocket auth frame: ${(err as Error).message}`); + ws.close(AUTH_CLOSE_CODE, 'invalid token'); + return; + } + if (socketId === undefined) { + authenticate(frame, identity); + return; + } + const socketInstance = sockets[socketId]; + if (!socketInstance) return; + refresh(socketId, socketInstance, frame, identity); + }); ws.on('pong', function pong(data: Buffer) { - const socketId = data.toString(); + const pongId = data.toString(); // `sockets` is a plain object, so an unsolicited pong carrying an // inherited key ('__proto__', 'constructor', 'hasOwnProperty', ...) as // its payload would otherwise resolve through the prototype chain @@ -776,11 +972,11 @@ export const wsConnection = ( // `pongTs` onto Object.prototype (tripped over by the for...in loops in // checkSockets/cleanChannel on every GC pass). Guarding with an // own-property check rejects every such key in one place. - if (!Object.prototype.hasOwnProperty.call(sockets, socketId)) { - logger.warn(`pong received for nonexistent socket ${socketId}`); + if (!Object.prototype.hasOwnProperty.call(sockets, pongId)) { + logger.warn(`pong received for nonexistent socket ${pongId}`); return; } - sockets[socketId].pongTs = Date.now(); + sockets[pongId].pongTs = Date.now(); }); }; @@ -846,7 +1042,19 @@ export const isOriginAllowed = (request: http.IncomingMessage): boolean => { }; /** - * HTTP `upgrade` event handler, called via httpServer + * Writes a bare HTTP status line to a socket that is not being upgraded and + * closes it, so the client sees why the handshake failed rather than a dropped + * connection. + */ +const rejectUpgrade = (socket: net.Socket, status: string): void => { + socket.write(`HTTP/1.1 ${status}\r\nConnection: close\r\n\r\n`); + socket.destroy(); +}; + +/** + * HTTP `upgrade` event handler, called via httpServer. The handshake carries no + * credential: it only checks the `Origin` allowlist and subscriber health, and + * the client authenticates in-band with its first frame (see `wsConnection`). */ export const httpUpgrade = ( request: http.IncomingMessage, @@ -859,7 +1067,7 @@ export const httpUpgrade = ( request.headers.origin || '(none)' }`, ); - socket.destroy(); + rejectUpgrade(socket, '403 Forbidden'); return; } @@ -870,45 +1078,13 @@ export const httpUpgrade = ( if (!subscriberHealthy) { statsd.increment('ws_upgrade_rejected'); logger.warn('Rejecting WebSocket upgrade: Redis subscriber unhealthy'); - socket.write('HTTP/1.1 503 Service Unavailable\r\n\r\n'); - socket.destroy(); - return; - } - - let identity: SocketIdentity; - try { - identity = readSocketIdentity(request); - } catch (err) { - // Token invalid/absent: do not establish a WebSocket connection. Record a - // structured warning (with the request's remote address) so rejected - // upgrade attempts are auditable, without logging the token itself. - statsd.increment('ws_upgrade_rejected'); - logger.warn( - `Rejected WebSocket upgrade from ${request.socket.remoteAddress ?? 'unknown'}: ` + - `${(err as Error).message}`, - ); - socket.destroy(); + rejectUpgrade(socket, '503 Service Unavailable'); return; } - // Optional per-tab id advertised on the connect URL (`?tab_id=`), used to - // also bind the socket to a per-tab channel for targeted task-status delivery. - // Parsed from the request URL (mirrors httpRequest); absent/blank/invalid -> - // no per-tab channel (principal-grain delivery, as before). - let tabId: string | undefined; - try { - const parsed = new URL( - request.url ?? '', - `http://${request.headers.host ?? 'localhost'}`, - ).searchParams.get('tab_id'); - if (parsed && isValidTabId(parsed)) tabId = parsed; - } catch { - // Malformed URL: ignore the tab hint and fall back to principal-grain. - } - wss.handleUpgrade(request, socket, head, function cb(ws: WebSocket) { wss.emit('connection', ws, request); - wsConnection(ws, identity, tabId); + wsConnection(ws); }); }; @@ -930,10 +1106,10 @@ export const checkSockets = () => { if (now >= socketInstance.identity.tokenExpiresAtMs) { logger.debug( - `terminating socket with expired token: ${socketId}, channel: ${socketInstance.channel}`, + `closing socket with expired token: ${socketId}, channel: ${socketInstance.channel}`, ); statsd.increment('ws_token_expired_disconnect'); - socketInstance.ws.terminate(); + socketInstance.ws.close(AUTH_CLOSE_CODE, 'token expired'); isActive = false; } else if (timeout >= opts.socketResponseTimeoutMs) { logger.debug( @@ -1017,6 +1193,7 @@ export const resetState = () => { channels = {}; sockets = {}; messageBound = false; + pendingConnections = 0; // A reset server is a clean, running (subscribed) server, so tests start from a // healthy transport; a test drives the unhealthy edge with markSubscriberUnhealthy. subscriberHealthy = true; diff --git a/superset-websocket/utils/client-ws-app/public/javascripts/app.js b/superset-websocket/utils/client-ws-app/public/javascripts/app.js index 46590b4bd543..bace67e24255 100644 --- a/superset-websocket/utils/client-ws-app/public/javascripts/app.js +++ b/superset-websocket/utils/client-ws-app/public/javascripts/app.js @@ -24,28 +24,23 @@ function ts() { return new Date().getTime(); } -const cookieName = document.getElementById('cookie').innerHTML; const tokenData = document.getElementById('tokens').innerHTML; const tokens = JSON.parse(tokenData); function connect() { if (socketCount >= tokens.length) return; - // using https://github.com/js-cookie/js-cookie - // oxlint-disable-next-line no-undef - Cookies.set(cookieName, tokens[socketCount], { path: '' }); - // Create WebSocket connection. let url = `ws://127.0.0.1:8080?last_id=${ts()}`; const socket = new WebSocket(url); // Connection opened socket.addEventListener('open', function () { + // The first frame authenticates the socket with its ticket. + socket.send(JSON.stringify({ type: 'auth', token: tokens[socketCount] })); socketCount++; document.getElementById('socket-count').innerHTML = socketCount; connect(); - - socket.send('Hello Server!'); }); // Listen for messages diff --git a/superset-websocket/utils/client-ws-app/routes/index.js b/superset-websocket/utils/client-ws-app/routes/index.js index 2ab6858418af..8f15a2411fe8 100644 --- a/superset-websocket/utils/client-ws-app/routes/index.js +++ b/superset-websocket/utils/client-ws-app/routes/index.js @@ -47,7 +47,6 @@ router.get('/', function (req, res) { res.render('index', { tokens: JSON.stringify(tokens), - c: config.jwtCookieName, }); }); diff --git a/superset-websocket/utils/client-ws-app/views/index.pug b/superset-websocket/utils/client-ws-app/views/index.pug index 2322bec5805c..e0f25769e3e1 100644 --- a/superset-websocket/utils/client-ws-app/views/index.pug +++ b/superset-websocket/utils/client-ws-app/views/index.pug @@ -20,7 +20,6 @@ extends layout block content h1 Websocket test app #tokens(style="display:none")= tokens - #cookie(style="display:none")= c div Sockets connected: span#socket-count 0 diff --git a/superset-websocket/utils/client-ws-app/views/layout.pug b/superset-websocket/utils/client-ws-app/views/layout.pug index a2a1be19c51a..5a046fd8060d 100644 --- a/superset-websocket/utils/client-ws-app/views/layout.pug +++ b/superset-websocket/utils/client-ws-app/views/layout.pug @@ -20,7 +20,6 @@ html head title= title link(rel='stylesheet', href='/stylesheets/style.css') - script(src="https://cdn.jsdelivr.net/npm/js-cookie@rc/dist/js.cookie.min.js") body block content diff --git a/superset/config.py b/superset/config.py index efba67f3f816..3631b4719872 100644 --- a/superset/config.py +++ b/superset/config.py @@ -3358,24 +3358,23 @@ def EMAIL_HEADER_MUTATOR( # pylint: disable=invalid-name,unused-argument # noq # a principal's tabs); a task type may narrow them to a per-tab channel so only # the tab watching a task is notified. The route never reaches the browser, # which dispatches on `topic`. -# The JWT authenticates the socket connection and binds it to its principal -# channel; a browser may also advertise a tab id on the connect URL to bind a -# per-tab channel (derived from the authorized principal channel). Set a strong +# A JWT ticket, fetched from GET /api/v1/realtime/token/, authenticates the socket +# and binds it to its principal channel. The browser sends it as the first frame +# after connecting, together with an optional tab id that binds a per-tab channel +# (derived from the authorized principal channel), and sends a fresh one on the +# open socket before the current one expires. No credential rides the handshake +# or the URL. Set a strong # random WEBSOCKET_JWT_SECRET (>= 32 bytes) in production. # The websocket server can be configured with a previous validation secret -# during rotations; the Flask app always mints new cookies with the current key. -# The websocket server validates the signed token at connection time and -# terminates sockets after JWT expiry; the Flask app re-mints the cookie inside a -# sliding window before expiry (on any request), and the browser proactively -# refreshes and reconnects, so an active realtime surface stays connected. Post- -# mint permission revocation is bounded by this lifetime plus the ping interval. +# during rotations; the Flask app always mints new tickets with the current key. +# The websocket server validates the signed ticket in-band and closes sockets +# (code 4401) when their ticket expires; the browser re-authenticates on the open +# socket at 75% of the ticket lifetime, so an active realtime surface stays +# connected without reconnecting. Post-mint permission revocation is bounded by +# this lifetime plus the ping interval. WEBSOCKET_ENABLE = False WEBSOCKET_URL = "ws://127.0.0.1:8080/" WEBSOCKET_JWT_SECRET = CHANGE_ME_WEBSOCKET_JWT_SECRET -WEBSOCKET_JWT_COOKIE_NAME = "superset-ws-token" # noqa: S105 -WEBSOCKET_JWT_COOKIE_SECURE = False -WEBSOCKET_JWT_COOKIE_SAMESITE: None | (Literal["None", "Lax", "Strict"]) = None -WEBSOCKET_JWT_COOKIE_DOMAIN = None WEBSOCKET_JWT_EXPIRATION_SECONDS = int(timedelta(minutes=15).total_seconds()) # Prefix for the realtime pub/sub channel (default ""). Redis pub/sub is not diff --git a/superset/initialization/__init__.py b/superset/initialization/__init__.py index 3280ea87e826..a9b844ecfccd 100644 --- a/superset/initialization/__init__.py +++ b/superset/initialization/__init__.py @@ -332,6 +332,11 @@ def init_views(self) -> None: appbuilder.add_api(SubjectRestApi) appbuilder.add_api(LogRestApi) + if self.config.get("WEBSOCKET_ENABLE"): + from superset.websocket.api import RealtimeRestApi + + appbuilder.add_api(RealtimeRestApi) + if feature_flag_manager.is_feature_enabled("ENABLE_EXTENSIONS"): from superset.extensions.api import ExtensionsRestApi from superset.extensions.storage.api import ExtensionStorageRestApi @@ -1120,7 +1125,6 @@ def init_app_in_ctx(self) -> None: self.configure_stats_manager() self.configure_semantic_cache() self.configure_task_manager() - self.configure_websocket() # Hook that provides administrators a handle on the Flask APP # after initialization @@ -1803,13 +1807,6 @@ def configure_task_manager(self) -> None: TaskManager.init_app(self.superset_app) - def configure_websocket(self) -> None: - """Mint the websocket channel-token cookie when the transport is enabled.""" - if self.config.get("WEBSOCKET_ENABLE"): - from superset.websocket.channel import register_ws_channel_cookie - - register_ws_channel_cookie(self.superset_app) - def register_blueprints(self) -> None: # Register custom blueprints from config for bp in self.config["BLUEPRINTS"]: diff --git a/superset/views/base.py b/superset/views/base.py index 1de41c6c553c..08021ca10d38 100644 --- a/superset/views/base.py +++ b/superset/views/base.py @@ -100,7 +100,6 @@ "GLOBAL_ASYNC_QUERIES_DEFAULT", "WEBSOCKET_ENABLE", "WEBSOCKET_URL", - "WEBSOCKET_JWT_EXPIRATION_SECONDS", "SQL_VALIDATORS_BY_ENGINE", "SQLALCHEMY_DOCS_URL", "SQLALCHEMY_DISPLAY_TEXT", diff --git a/superset/websocket/api.py b/superset/websocket/api.py new file mode 100644 index 000000000000..da4a3e2661c8 --- /dev/null +++ b/superset/websocket/api.py @@ -0,0 +1,82 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""REST API issuing connection tickets for the realtime websocket transport.""" + +from __future__ import annotations + +from flask import current_app, Response +from flask_appbuilder.api import expose, protect, safe +from flask_appbuilder.security.decorators import permission_name + +from superset.extensions import event_logger +from superset.views.base_api import BaseSupersetApi, statsd_metrics +from superset.websocket.channel import get_realtime_principal, mint_channel_token +from superset.websocket.permissions import REALTIME_NOTIFICATION_RESOURCE + + +class RealtimeRestApi(BaseSupersetApi): + """Issues short-lived tickets the browser presents to the websocket server.""" + + resource_name = "realtime" + class_permission_name = REALTIME_NOTIFICATION_RESOURCE + allow_browser_login = True + openapi_spec_tag = "Realtime" + + @expose("/token/", methods=("GET",)) + @event_logger.log_this + @protect() + @safe + @statsd_metrics + @permission_name("read") + def token(self) -> Response: + """Get a websocket connection ticket for the current principal. + --- + get: + summary: Get a websocket connection ticket + description: >- + Returns a short-lived JWT the client sends as the first frame of a + websocket connection, and again on the open socket to refresh it + before it expires. + responses: + 200: + description: A connection ticket and its lifetime in seconds + content: + application/json: + schema: + type: object + properties: + token: + type: string + expires_in: + type: integer + 401: + $ref: '#/components/responses/401' + 403: + $ref: '#/components/responses/403' + """ + principal = get_realtime_principal() + if principal is None: + return self.response_401() + response = self.response( + 200, + token=mint_channel_token(principal), + expires_in=current_app.config["WEBSOCKET_JWT_EXPIRATION_SECONDS"], + ) + # The ticket is a bearer credential; a guest authenticates by header, not + # cookie, so a shared cache must never store or replay it. + response.headers["Cache-Control"] = "no-store, private" + return response diff --git a/superset/websocket/channel.py b/superset/websocket/channel.py index a202c5fc202a..20ae27169ee7 100644 --- a/superset/websocket/channel.py +++ b/superset/websocket/channel.py @@ -16,9 +16,11 @@ # under the License. """Per-principal websocket routing identity and connection tokens. -The realtime transport requires a JWT cookie minted by Superset after the -request principal passes the ``can_read Realtime`` gate. The JWT carries the -principal identity and a deterministic routing key - ``user:`` for a +The realtime transport requires a short-lived JWT ticket minted by Superset +(``GET /api/v1/realtime/token/``) after the request principal passes the +``can_read Realtime`` gate. The client presents it in-band as the first frame of +the socket and again, on the open socket, to refresh it before expiry. The JWT +carries the principal identity and a deterministic routing key - ``user:`` for a logged-in user, a stable HMAC for an embedded guest - so the ``superset-websocket`` server can validate the socket and bind it to that routing key. @@ -37,13 +39,11 @@ from typing import Literal, TypedDict import jwt -from flask import Flask, request, Response from superset import security_manager from superset.tasks.guest import get_current_guest_subscriber_key from superset.utils.core import get_user_id from superset.websocket.permissions import ( - can_access_realtime_notifications, REALTIME_NOTIFICATION_JWT_AUDIENCE, REALTIME_NOTIFICATION_JWT_ISSUER, ) @@ -79,7 +79,7 @@ def get_realtime_principal() -> RealtimePrincipal | None: """Return JWT principal claims for the current websocket-eligible principal. ``None`` for a request with no realtime identity (anonymous, or a guest whose - token yields no subscriber key), in which case no cookie is minted. + token yields no subscriber key), in which case no ticket is minted. """ if security_manager.get_current_guest_user_if_guest(): guest_channel = channel_id_for(None, get_current_guest_subscriber_key()) @@ -104,7 +104,7 @@ def get_realtime_principal() -> RealtimePrincipal | None: def mint_channel_token(principal: RealtimePrincipal) -> str: - """Sign a JWT binding a websocket connection to ``principal``. + """Sign a JWT ticket binding a websocket connection to ``principal``. The ``superset-websocket`` server verifies this with the same secret, checks the identity claims, then uses ``channel`` as the routing key for targeted @@ -124,111 +124,3 @@ def mint_channel_token(principal: RealtimePrincipal) -> str: return jwt.encode( payload, current_app.config["WEBSOCKET_JWT_SECRET"], algorithm="HS256" ) - - -def _valid_payload_channel(payload: dict[str, object]) -> str | None: - channel = payload.get("channel") - subject = payload.get("sub") - principal_type = payload.get("principal_type") - - if not isinstance(channel, str) or not channel: - return None - if not isinstance(subject, str) or not subject: - return None - if principal_type not in ("user", "guest"): - return None - - if principal_type == "user" and channel != f"user:{subject}": - return None - if principal_type == "guest" and channel != subject: - return None - return channel - - -def _decode_cookie(token: str) -> dict[str, object] | None: - """Return the claims of a valid, unexpired channel token, else None.""" - from flask import current_app - - try: - payload = jwt.decode( - token, - current_app.config["WEBSOCKET_JWT_SECRET"], - algorithms=["HS256"], - audience=REALTIME_NOTIFICATION_JWT_AUDIENCE, - issuer=REALTIME_NOTIFICATION_JWT_ISSUER, - ) - except jwt.InvalidTokenError: - return None - return payload if isinstance(payload, dict) else None - - -def _should_remint(existing: str | None, channel_id: str) -> bool: - """Whether the channel cookie must be (re)minted for ``channel_id``. - - Re-mints when the cookie is missing, invalid/expired, bound to a different - principal's channel, or inside a sliding refresh window before expiry (less - than half its lifetime remaining). The sliding window keeps an active - surface's socket alive indefinitely: any request in the second half of the - token's life issues a fresh cookie, so the socket reconnects before the - server terminates it at expiry (rather than only re-minting after expiry). - """ - from flask import current_app - - if not existing: - return True - payload = _decode_cookie(existing) - if payload is None: - return True - if _valid_payload_channel(payload) != channel_id: - return True - exp = payload.get("exp") - if not isinstance(exp, (int, float)): - return True - lifetime = current_app.config["WEBSOCKET_JWT_EXPIRATION_SECONDS"] - remaining = exp - datetime.now(tz=timezone.utc).timestamp() - return remaining < lifetime / 2 - - -def register_ws_channel_cookie(app: Flask) -> None: - """Keep the websocket channel-token cookie in sync with the request principal. - - Sets an ``httponly`` JWT cookie (which the ``superset-websocket`` server - verifies) bound to the current principal's channel. Re-mints the cookie - whenever the bound channel no longer matches the current principal — e.g. - after a logout/login in the same browser — or when it is inside the sliding - refresh window before expiry, and deletes it when there is no principal - (anonymous / logged out), so a stale token can never bind a new session to a - previous user's channel and an active surface never silently loses realtime. - """ - cookie_name = app.config["WEBSOCKET_JWT_COOKIE_NAME"] - cookie_domain = app.config["WEBSOCKET_JWT_COOKIE_DOMAIN"] - - @app.after_request - def set_ws_channel_cookie(response: Response) -> Response: - principal = ( - get_realtime_principal() if can_access_realtime_notifications() else None - ) - existing = request.cookies.get(cookie_name) - - if principal is None: - # No channel for this principal (anonymous / just logged out): drop any - # stale cookie so it can't be reused by a later session or after the - # realtime notification permission is removed. - if existing: - response.delete_cookie(cookie_name, domain=cookie_domain) - return response - - channel_id = principal["channel"] - if not _should_remint(existing, channel_id): - return response - - response.set_cookie( - cookie_name, - value=mint_channel_token(principal), - httponly=True, - secure=app.config["WEBSOCKET_JWT_COOKIE_SECURE"], - domain=cookie_domain, - samesite=app.config["WEBSOCKET_JWT_COOKIE_SAMESITE"], - max_age=app.config["WEBSOCKET_JWT_EXPIRATION_SECONDS"], - ) - return response diff --git a/tests/unit_tests/websocket/test_api.py b/tests/unit_tests/websocket/test_api.py new file mode 100644 index 000000000000..f4ab0d5a02f6 --- /dev/null +++ b/tests/unit_tests/websocket/test_api.py @@ -0,0 +1,90 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from typing import Any + +import jwt +import pytest +from pytest_mock import MockerFixture + +from superset.websocket.permissions import ( + REALTIME_NOTIFICATION_JWT_AUDIENCE, + REALTIME_NOTIFICATION_JWT_ISSUER, +) + +SECRET = "x" * 40 + +WEBSOCKET_APP_CONFIG = { + "WEBSOCKET_ENABLE": True, + "WEBSOCKET_JWT_SECRET": SECRET, + "WEBSOCKET_JWT_EXPIRATION_SECONDS": 120, +} + + +@pytest.mark.parametrize("app", [WEBSOCKET_APP_CONFIG], indirect=True) +def test_token_endpoint_requires_a_session(client: Any) -> None: + assert client.get("/api/v1/realtime/token/").status_code == 401 + + +@pytest.mark.parametrize("app", [WEBSOCKET_APP_CONFIG], indirect=True) +def test_token_endpoint_returns_a_ticket_for_the_principal( + client: Any, full_api_access: None, mocker: MockerFixture +) -> None: + mocker.patch( + "superset.websocket.api.get_realtime_principal", + return_value={"channel": "user:7", "principal_type": "user", "sub": "7"}, + ) + + response = client.get("/api/v1/realtime/token/") + + assert response.status_code == 200 + assert response.headers["Cache-Control"] == "no-store, private" + body = response.json + assert body["expires_in"] == 120 + claims = jwt.decode( + body["token"], + SECRET, + algorithms=["HS256"], + audience=REALTIME_NOTIFICATION_JWT_AUDIENCE, + issuer=REALTIME_NOTIFICATION_JWT_ISSUER, + ) + assert claims["channel"] == "user:7" + assert claims["exp"] - claims["iat"] == 120 + + +@pytest.mark.parametrize("app", [WEBSOCKET_APP_CONFIG], indirect=True) +def test_token_endpoint_rejects_a_principal_without_a_channel( + client: Any, full_api_access: None, mocker: MockerFixture +) -> None: + mocker.patch("superset.websocket.api.get_realtime_principal", return_value=None) + + assert client.get("/api/v1/realtime/token/").status_code == 401 + + +@pytest.mark.parametrize("app", [WEBSOCKET_APP_CONFIG], indirect=True) +def test_token_endpoint_requires_the_realtime_permission( + client: Any, mocker: MockerFixture +) -> None: + from superset import security_manager + + mocker.patch( + "flask_appbuilder.security.decorators.verify_jwt_in_request", + return_value=True, + ) + mocker.patch.object(security_manager, "is_item_public", return_value=False) + mocker.patch.object(security_manager, "has_access", return_value=False) + + assert client.get("/api/v1/realtime/token/").status_code == 403 diff --git a/tests/unit_tests/websocket/test_channel.py b/tests/unit_tests/websocket/test_channel.py index 74f0c8ad8096..4ca9a4d9fd53 100644 --- a/tests/unit_tests/websocket/test_channel.py +++ b/tests/unit_tests/websocket/test_channel.py @@ -14,8 +14,6 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from datetime import datetime, timedelta, timezone - import jwt from pytest_mock import MockerFixture @@ -127,193 +125,3 @@ def test_realtime_principal_none_for_guest_without_key( mocker.patch.object(channel, "get_current_guest_subscriber_key", return_value=None) assert channel.get_realtime_principal() is None - - -def _make_ws_app(): - """A minimal Flask app carrying the WEBSOCKET_* config the hook reads.""" - from flask import Flask, jsonify - - from superset.websocket.channel import register_ws_channel_cookie - - app = Flask(__name__) - app.config.update( - WEBSOCKET_ENABLE=True, - WEBSOCKET_JWT_SECRET="x" * 40, - WEBSOCKET_JWT_COOKIE_NAME="superset-ws-token", - WEBSOCKET_JWT_COOKIE_SECURE=False, - WEBSOCKET_JWT_COOKIE_SAMESITE=None, - WEBSOCKET_JWT_COOKIE_DOMAIN=None, - WEBSOCKET_JWT_EXPIRATION_SECONDS=3600, - ) - register_ws_channel_cookie(app) - - @app.route("/_ws_probe") - def _ws_probe(): # pragma: no cover - trivial - return jsonify(ok=True) - - return app - - -def _ws_set_cookies(response): - return [c for c in response.headers.getlist("Set-Cookie") if "ws-token" in c] - - -def test_cookie_reminted_when_principal_changes(mocker) -> None: - from superset.websocket import channel - - client = _make_ws_app().test_client() - mocker.patch.object(channel, "can_access_realtime_notifications", return_value=True) - mocker.patch.object( - channel.security_manager, "get_current_guest_user_if_guest", return_value=None - ) - - mocker.patch.object(channel, "get_user_id", return_value=1) - assert _ws_set_cookies(client.get("/_ws_probe")) # minted for user:1 - - # A different principal on the same client must re-mint (not keep user:1). - mocker.patch.object(channel, "get_user_id", return_value=2) - assert _ws_set_cookies(client.get("/_ws_probe")) - - -def test_cookie_signed_with_old_secret_is_reminted_with_current_secret(mocker) -> None: - from superset.websocket import channel - from superset.websocket.permissions import ( - REALTIME_NOTIFICATION_JWT_AUDIENCE, - REALTIME_NOTIFICATION_JWT_ISSUER, - ) - - client = _make_ws_app().test_client() - mocker.patch.object(channel, "can_access_realtime_notifications", return_value=True) - mocker.patch.object( - channel.security_manager, "get_current_guest_user_if_guest", return_value=None - ) - mocker.patch.object(channel, "get_user_id", return_value=1) - - old_key_token = jwt.encode( - { - "channel": "user:1", - "principal_type": "user", - "sub": "1", - "aud": REALTIME_NOTIFICATION_JWT_AUDIENCE, - "iss": REALTIME_NOTIFICATION_JWT_ISSUER, - "exp": datetime.now(tz=timezone.utc) + timedelta(hours=1), - }, - "y" * 40, - algorithm="HS256", - ) - client.set_cookie("superset-ws-token", old_key_token) - - assert _ws_set_cookies(client.get("/_ws_probe")) - - -def _mint_ws_token( - exp_delta_seconds: int, - secret: str = "x" * 40, - channel_id: str = "user:1", -) -> str: - from superset.websocket.permissions import ( - REALTIME_NOTIFICATION_JWT_AUDIENCE, - REALTIME_NOTIFICATION_JWT_ISSUER, - ) - - return jwt.encode( - { - "channel": channel_id, - "principal_type": "user", - "sub": channel_id.split(":")[1], - "aud": REALTIME_NOTIFICATION_JWT_AUDIENCE, - "iss": REALTIME_NOTIFICATION_JWT_ISSUER, - "exp": datetime.now(tz=timezone.utc) + timedelta(seconds=exp_delta_seconds), - }, - secret, - algorithm="HS256", - ) - - -def test_cookie_not_reminted_when_fresh(mocker) -> None: - from superset.websocket import channel - - client = _make_ws_app().test_client() - mocker.patch.object(channel, "can_access_realtime_notifications", return_value=True) - mocker.patch.object( - channel.security_manager, "get_current_guest_user_if_guest", return_value=None - ) - mocker.patch.object(channel, "get_user_id", return_value=1) - - # Full lifetime remaining (> half of 3600s) → the hook leaves it alone, so a - # Set-Cookie header is not written on every request. - client.set_cookie("superset-ws-token", _mint_ws_token(3600)) - assert not _ws_set_cookies(client.get("/_ws_probe")) - - -def test_cookie_reminted_inside_refresh_window(mocker) -> None: - from superset.websocket import channel - - client = _make_ws_app().test_client() - mocker.patch.object(channel, "can_access_realtime_notifications", return_value=True) - mocker.patch.object( - channel.security_manager, "get_current_guest_user_if_guest", return_value=None - ) - mocker.patch.object(channel, "get_user_id", return_value=1) - - # Only 60s remaining (< half of 3600s) → sliding-window re-mint keeps an active - # surface's socket alive past the original expiry (Finding 6). - client.set_cookie("superset-ws-token", _mint_ws_token(60)) - assert _ws_set_cookies(client.get("/_ws_probe")) - - -def test_cookie_cleared_for_anonymous(mocker) -> None: - from superset.websocket import channel - - client = _make_ws_app().test_client() - mocker.patch.object(channel, "can_access_realtime_notifications", return_value=True) - mocker.patch.object( - channel.security_manager, "get_current_guest_user_if_guest", return_value=None - ) - mocker.patch.object(channel, "get_user_id", return_value=None) - - client.set_cookie("superset-ws-token", "stale") - resp = client.get("/_ws_probe") - cleared = [ - c - for c in resp.headers.getlist("Set-Cookie") - if "superset-ws-token" in c and ("Expires" in c or "Max-Age=0" in c) - ] - assert cleared, "stale cookie must be cleared when there is no principal" - - -def test_cookie_not_minted_without_realtime_permission(mocker) -> None: - from superset.websocket import channel - - client = _make_ws_app().test_client() - mocker.patch.object( - channel, "can_access_realtime_notifications", return_value=False - ) - mocker.patch.object( - channel.security_manager, "get_current_guest_user_if_guest", return_value=None - ) - mocker.patch.object(channel, "get_user_id", return_value=1) - - assert not _ws_set_cookies(client.get("/_ws_probe")) - - -def test_cookie_cleared_when_realtime_permission_removed(mocker) -> None: - from superset.websocket import channel - - client = _make_ws_app().test_client() - mocker.patch.object( - channel, "can_access_realtime_notifications", return_value=False - ) - mocker.patch.object( - channel.security_manager, "get_current_guest_user_if_guest", return_value=None - ) - mocker.patch.object(channel, "get_user_id", return_value=1) - - client.set_cookie("superset-ws-token", "stale") - resp = client.get("/_ws_probe") - cleared = [ - c - for c in resp.headers.getlist("Set-Cookie") - if "superset-ws-token" in c and ("Expires" in c or "Max-Age=0" in c) - ] - assert cleared, "stale cookie must be cleared without realtime permission"