villebro commented on code in PR #43431: URL: https://github.com/apache/superset/pull/43431#discussion_r3838982380
########## superset/websocket/channel.py: ########## @@ -0,0 +1,103 @@ +# 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. +"""Per-principal websocket channel identity and connection tokens. + +The realtime transport delivers a **per-principal** channel's events only to +that principal's sockets. A channel is derived deterministically from the +authenticated request principal — ``user:<id>`` for a logged-in user, a stable +HMAC for an embedded guest — so the task layer can publish a user's events to +their channel without threading a per-request channel id (unlike the legacy +per-session-random channel). The channel is minted into a JWT cookie that the +``superset-websocket`` server verifies to bind a socket to its channel. + +This is the connection-auth tier. The lossy public list-view pub/sub +(``entity-changes:*``) needs no per-principal channel — see +``TaskManager.publish_entity_change``. +""" + +from __future__ import annotations + +import logging +from datetime import datetime, timedelta, timezone + +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 + +logger = logging.getLogger(__name__) + + +def get_channel_id() -> str | None: + """Return the realtime channel for the current request principal, or ``None``. + + ``user:<id>`` for a logged-in user, ``guest:<hmac>`` for an embedded guest + (reusing the guest identity from ``superset.tasks.guest`` for consistency), + and ``None`` for an anonymous request (no channel, no cookie). + """ + if security_manager.get_current_guest_user_if_guest(): + return get_current_guest_subscriber_key() + if (user_id := get_user_id()) is not None: + return f"user:{user_id}" + return None + + +def mint_channel_token(channel_id: str) -> str: + """Sign a JWT binding a websocket connection to ``channel_id``. + + The ``superset-websocket`` server verifies this with the same secret and + reads the ``channel`` claim to route per-principal events to the socket. + """ + from flask import current_app + + now = datetime.now(tz=timezone.utc) + expiration = current_app.config["WEBSOCKET_JWT_EXPIRATION_SECONDS"] + payload = {"channel": channel_id, "exp": now + timedelta(seconds=expiration)} + return jwt.encode( + payload, current_app.config["WEBSOCKET_JWT_SECRET"], algorithm="HS256" + ) + + +def register_ws_channel_cookie(app: Flask) -> None: + """Set the websocket channel-token cookie on responses for channel-bearing users. + + Mirrors the connection-auth handshake of the legacy async-events transport: + an ``httponly`` JWT cookie the browser sends when opening the socket, which + the ``superset-websocket`` server verifies. Refreshed only when missing so it + is not re-signed on every request. + """ + cookie_name = app.config["WEBSOCKET_JWT_COOKIE_NAME"] + + @app.after_request + def set_ws_channel_cookie(response: Response) -> Response: + channel_id = get_channel_id() + if channel_id is None: + return response + if request.cookies.get(cookie_name): + return response Review Comment: Addressed in 100e208622. `register_ws_channel_cookie` no longer trusts a cookie just because it exists: `_cookie_channel()` decodes/verifies the existing token and the cookie is re-minted whenever its `channel` claim ≠ the current principal (missing, invalid, expired, or a different user after logout→login), and `delete_cookie` is called when there is no principal (anonymous / just logged out — the after_request runs after the session is cleared, so `get_channel_id()` returns `None`). Covered by `test_cookie_reminted_when_principal_changes` and `test_cookie_cleared_for_anonymous`. ########## superset/initialization/__init__.py: ########## @@ -1046,6 +1046,7 @@ def init_app_in_ctx(self) -> None: self.configure_ssh_manager() self.configure_stats_manager() self.configure_task_manager() + self.configure_websocket() Review Comment: Addressed in 100e208622. `register_ws_channel_cookie` no longer trusts a cookie just because it exists: `_cookie_channel()` decodes/verifies the existing token and the cookie is re-minted whenever its `channel` claim ≠ the current principal (missing, invalid, expired, or a different user after logout→login), and `delete_cookie` is called when there is no principal (anonymous / just logged out — the after_request runs after the session is cleared, so `get_channel_id()` returns `None`). Covered by `test_cookie_reminted_when_principal_changes` and `test_cookie_cleared_for_anonymous`. ########## superset/constants.py: ########## @@ -29,6 +29,7 @@ CHANGE_ME_SECRET_KEY = "CHANGE_ME_TO_A_COMPLEX_RANDOM_SECRET" # noqa: S105 CHANGE_ME_GUEST_TOKEN_JWT_SECRET = "test-guest-secret-change-me" # noqa: S105 +CHANGE_ME_WEBSOCKET_JWT_SECRET = "test-ws-secret-change-me" # noqa: S105 Review Comment: Addressed in 100e208622 via a startup guard, `check_websocket_secret()` (initialization/__init__.py, called in the check sequence): when `WEBSOCKET_ENABLED` is on and `WEBSOCKET_JWT_SECRET` is the placeholder or shorter than 32 bytes, it warns and `sys.exit(1)` outside debug/testing — so an enabled production deployment cannot run with the known default. This mirrors the established `check_guest_token_secret` / `check_secret_key` convention (a `CHANGE_ME_*` placeholder plus a startup check), rather than introducing a different pattern. Covered by check_websocket_secret_test.py. ########## superset/config.py: ########## @@ -2931,6 +2932,27 @@ def EMAIL_HEADER_MUTATOR( # pylint: disable=invalid-name,unused-argument # noq # async out per dashboard. GLOBAL_ASYNC_QUERIES_DEFAULT = True +# Realtime websocket transport (the `superset-websocket` server) config. +# When enabled, GTF task changes are pushed to the browser so charts and list +# views update without waiting for the interval poll (which stays as the +# fallback). Requires the superset-websocket server and a Redis coordination +# backend (DISTRIBUTED_COORDINATION_CONFIG). Two channel tiers: +# - a public per-entity-type pub/sub (e.g. entity-changes:task) for lossy +# list-view activity (opaque id + status), and +# - a per-principal channel (user:<id> / guest:<hmac>) for the dashboard +# chart-data path, authenticated by the JWT cookie below. +# The JWT authenticates the socket connection and binds it to its channel; the +# server delivers a per-principal channel's events only to that principal's +# sockets. Set a strong random WEBSOCKET_JWT_SECRET (>= 32 bytes) in production. +WEBSOCKET_ENABLED = False +WEBSOCKET_URL = "ws://127.0.0.1:8080/" +WEBSOCKET_JWT_SECRET = CHANGE_ME_WEBSOCKET_JWT_SECRET Review Comment: Addressed in 100e208622 via a startup guard, `check_websocket_secret()` (initialization/__init__.py, called in the check sequence): when `WEBSOCKET_ENABLED` is on and `WEBSOCKET_JWT_SECRET` is the placeholder or shorter than 32 bytes, it warns and `sys.exit(1)` outside debug/testing — so an enabled production deployment cannot run with the known default. This mirrors the established `check_guest_token_secret` / `check_secret_key` convention (a `CHANGE_ME_*` placeholder plus a startup check), rather than introducing a different pattern. Covered by check_websocket_secret_test.py. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
