From 9938ee273f1cc6cc4dc6d8bd43ceba74ca4faa5a Mon Sep 17 00:00:00 2001 From: Evan Rusackas Date: Thu, 11 Jun 2026 19:37:13 -0700 Subject: [PATCH] feat: terminate active sessions when an account is disabled (#40695) Co-authored-by: Claude Opus 4.8 --- UPDATING.md | 4 + superset/initialization/__init__.py | 17 ++ ...7a1c93e0b21_add_sessions_invalidated_at.py | 177 +++++++++++++++ superset/models/user_attributes.py | 17 +- superset/security/manager.py | 6 + superset/security/session_invalidation.py | 205 +++++++++++++++++ .../security/session_invalidation_tests.py | 79 +++++++ tests/unit_tests/security/audit_log_test.py | 8 +- .../security/test_session_invalidation.py | 214 ++++++++++++++++++ 9 files changed, 724 insertions(+), 3 deletions(-) create mode 100644 superset/migrations/versions/2026-06-02_10-00_f7a1c93e0b21_add_sessions_invalidated_at.py create mode 100644 superset/security/session_invalidation.py create mode 100644 tests/integration_tests/security/session_invalidation_tests.py create mode 100644 tests/unit_tests/security/test_session_invalidation.py diff --git a/UPDATING.md b/UPDATING.md index d0208eb5894..ce75b52b471 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -91,6 +91,10 @@ superset revoke-guest-tokens This change is backward compatible. The feature is off by default, and even when enabled nothing is revoked until an admin explicitly bumps the version: the expected version starts at `0`, and tokens minted before this change (which carry no version claim) are treated as version `0`. No database migration is required. +### Sessions are terminated when an account is disabled + +Disabling a user account (setting `active` to `False`, via the admin UI, REST API, or CLI) now terminates that user's outstanding sessions on their next request, instead of relying on a passive check. This works for both client-side cookie sessions and server-side session stores via a per-user invalidation epoch (`user_attribute.sessions_invalidated_at`, added by a migration). The mechanism is inert for users that were never disabled (NULL epoch), so there is no behavior change for active users. Re-enabling an account and logging in again starts a fresh, valid session. The migration backfills the epoch for accounts that are already disabled at upgrade time, so re-enabling such an account does not revive a session that predates this feature. + ### Dataset import validates catalog against the target connection Importing a dataset now validates the `catalog` field against the target database connection. When the connection has multi-catalog disabled (`allow_multi_catalog` off) and the dataset's catalog is not the connection's default catalog, the import fails instead of silently persisting the non-default catalog. This matches the validation already enforced on the dataset update path and prevents imported datasets from querying an unintended database. diff --git a/superset/initialization/__init__.py b/superset/initialization/__init__.py index c54821ccd98..ec7f7e233a2 100644 --- a/superset/initialization/__init__.py +++ b/superset/initialization/__init__.py @@ -764,6 +764,23 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods gc.collect() return response + @self.superset_app.before_request + def enforce_session_validity() -> Any: + """Force logout of sessions invalidated by a per-user epoch.""" + from superset.security.session_invalidation import ( + enforce_session_validity as _enforce, + ) + + return _enforce() + + # Stamp the per-user invalidation epoch when an account is disabled, + # so outstanding sessions are terminated on their next request. + from superset.security.session_invalidation import ( + register_session_invalidation_events, + ) + + register_session_invalidation_events(appbuilder.sm.user_model) + @self.superset_app.context_processor def get_common_bootstrap_data() -> dict[str, Any]: # Import here to avoid circular imports diff --git a/superset/migrations/versions/2026-06-02_10-00_f7a1c93e0b21_add_sessions_invalidated_at.py b/superset/migrations/versions/2026-06-02_10-00_f7a1c93e0b21_add_sessions_invalidated_at.py new file mode 100644 index 00000000000..7a0a7167136 --- /dev/null +++ b/superset/migrations/versions/2026-06-02_10-00_f7a1c93e0b21_add_sessions_invalidated_at.py @@ -0,0 +1,177 @@ +# 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. +"""add sessions_invalidated_at to user_attribute + +Revision ID: f7a1c93e0b21 +Revises: 31dae2559c05 +Create Date: 2026-06-02 10:00:00.000000 + +""" + +from datetime import datetime, timezone + +import sqlalchemy as sa +from alembic import op + +from superset.migrations.shared.utils import ( + add_columns, + create_index, + drop_columns, + drop_index, +) + +# revision identifiers, used by Alembic. +revision = "f7a1c93e0b21" +down_revision = "31dae2559c05" + +TABLE = "user_attribute" +COLUMN = "sessions_invalidated_at" +INDEX = "ix_user_attribute_sessions_invalidated_at" +UQ = "uq_user_attribute_user_id" + + +MERGE_COLUMNS = ("avatar_url", "welcome_dashboard_id", "sessions_invalidated_at") + + +def upgrade(): + add_columns(TABLE, sa.Column(COLUMN, sa.DateTime(), nullable=True)) + create_index(TABLE, INDEX, [COLUMN]) + # The model treats ``user_attribute`` as one row per user (all readers use + # ``extra_attributes[0]``). Enforce that invariant so the session-invalidation + # upsert is race-safe. Collapse any pre-existing duplicates into the lowest + # ``id`` per user before adding the unique constraint. + # + # The merge/de-dup is driven in Python rather than via correlated subqueries: + # MySQL forbids referencing the UPDATE/DELETE target table in a subquery + # (error 1093), so a portable SQL form is awkward. Reading the duplicate rows + # and resolving the winner in Python works identically on SQLite/MySQL/Postgres. + _dedupe_user_attributes() + # SQLite has no ALTER ... ADD CONSTRAINT, so use batch mode (copy-and-move) + # which all dialects support. + with op.batch_alter_table(TABLE) as batch_op: + batch_op.create_unique_constraint(UQ, ["user_id"]) + # Backfill an epoch for accounts that are already disabled at upgrade time. + # Without this, a user disabled before this lands and later re-enabled would + # have no epoch, so a pre-existing session (which also carries no recorded + # login time) would silently revive. Stamping the epoch now means any such + # outstanding session is treated as invalidated. + _backfill_disabled_users() + + +def _dedupe_user_attributes(): + """Collapse duplicate ``user_attribute`` rows into the lowest ``id`` per user. + + Settings stored only on a higher-``id`` duplicate are merged forward into the + kept row (the kept row's NULL columns take the lowest-``id`` non-NULL sibling + value) so nothing is silently lost, then the redundant rows are deleted. + """ + bind = op.get_bind() + columns = ", ".join(("id", "user_id", *MERGE_COLUMNS)) + rows = bind.execute( + sa.text(f"SELECT {columns} FROM {TABLE} ORDER BY id") # noqa: S608 + ).fetchall() + + by_user: dict[int, list] = {} + for row in rows: + mapping = row._mapping + if mapping["user_id"] is None: + continue + by_user.setdefault(mapping["user_id"], []).append(mapping) + + for user_rows in by_user.values(): + if len(user_rows) < 2: + continue + # Rows are ordered by id, so the first is the keeper. + keeper = user_rows[0] + duplicates = user_rows[1:] + updates = {} + for column in MERGE_COLUMNS: + if keeper[column] is not None: + continue + for dup in duplicates: + if dup[column] is not None: + updates[column] = dup[column] + break + if updates: + assignments = ", ".join(f"{col} = :{col}" for col in updates) + bind.execute( + sa.text( + f"UPDATE {TABLE} SET {assignments} WHERE id = :id" # noqa: S608 + ), + {**updates, "id": keeper["id"]}, + ) + bind.execute( + sa.text(f"DELETE FROM {TABLE} WHERE id = :id"), # noqa: S608 + [{"id": dup["id"]} for dup in duplicates], + ) + + +def _backfill_disabled_users(): + """Stamp the invalidation epoch for users that are already disabled. + + Upserts one ``user_attribute`` row per disabled user (``ab_user.active`` is + false) that has no epoch yet, so re-enabling such an account does not revive + a session that predates this feature. Done in Python to stay portable across + SQLite/MySQL/Postgres; the unique constraint added above guarantees one row + per user, so the insert path cannot create duplicates. + """ + bind = op.get_bind() + now = datetime.now(timezone.utc).replace(tzinfo=None) + + # FAB stores ``active`` as a boolean/tinyint; some legacy rows leave it NULL, + # which is not "explicitly disabled", so only match a false value. + disabled_user_ids = [ + row._mapping["id"] + for row in bind.execute( + sa.text("SELECT id FROM ab_user WHERE active = :inactive"), + {"inactive": False}, + ).fetchall() + ] + if not disabled_user_ids: + return + + existing = { + row._mapping["user_id"] + for row in bind.execute( + sa.text(f"SELECT user_id FROM {TABLE}") # noqa: S608 + ).fetchall() + } + + for user_id in disabled_user_ids: + if user_id in existing: + bind.execute( + sa.text( + f"UPDATE {TABLE} SET {COLUMN} = :now, changed_on = :now " # noqa: S608, E501 + f"WHERE user_id = :user_id AND {COLUMN} IS NULL" + ), + {"now": now, "user_id": user_id}, + ) + else: + bind.execute( + sa.text( + f"INSERT INTO {TABLE} (user_id, {COLUMN}, created_on, changed_on) " # noqa: S608, E501 + "VALUES (:user_id, :now, :now, :now)" + ), + {"now": now, "user_id": user_id}, + ) + + +def downgrade(): + with op.batch_alter_table(TABLE) as batch_op: + batch_op.drop_constraint(UQ, type_="unique") + drop_index(TABLE, INDEX) + drop_columns(TABLE, COLUMN) diff --git a/superset/models/user_attributes.py b/superset/models/user_attributes.py index 512270c89ce..92ac0c9dde1 100644 --- a/superset/models/user_attributes.py +++ b/superset/models/user_attributes.py @@ -16,7 +16,14 @@ # under the License. from flask_appbuilder import Model -from sqlalchemy import Column, ForeignKey, Integer, String +from sqlalchemy import ( + Column, + DateTime, + ForeignKey, + Integer, + String, + UniqueConstraint, +) from sqlalchemy.orm import relationship from superset import security_manager @@ -34,6 +41,9 @@ class UserAttribute(Model, AuditMixinNullable): """ __tablename__ = "user_attribute" + # One attribute row per user; readers rely on ``extra_attributes[0]`` and the + # session-invalidation upsert depends on this for race safety. + __table_args__ = (UniqueConstraint("user_id", name="uq_user_attribute_user_id"),) id = Column(Integer, primary_key=True) user_id = Column(Integer, ForeignKey("ab_user.id")) user = relationship( @@ -42,3 +52,8 @@ class UserAttribute(Model, AuditMixinNullable): welcome_dashboard_id = Column(Integer, ForeignKey("dashboards.id")) welcome_dashboard = relationship("Dashboard") avatar_url = Column(String(100)) + # When set, any session for this user whose login predates this timestamp + # is forced to log out (see ``superset.security.session_invalidation``). + # Stamped when the account is disabled, so outstanding sessions terminate + # regardless of the session backend. NULL means "never invalidated". + sessions_invalidated_at = Column(DateTime, nullable=True, index=True) diff --git a/superset/security/manager.py b/superset/security/manager.py index 96acd3ea461..f2de87e609a 100644 --- a/superset/security/manager.py +++ b/superset/security/manager.py @@ -639,6 +639,12 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods return lm def on_user_login(self, user: Any) -> None: + # pylint: disable=import-outside-toplevel + from superset.security.session_invalidation import stamp_login_time + + # Record the authentication time so outstanding sessions can be + # invalidated when the account is later disabled. + stamp_login_time() _log_audit_event( "UserLoggedIn", {"username": user.username, "user_id": user.id}, diff --git a/superset/security/session_invalidation.py b/superset/security/session_invalidation.py new file mode 100644 index 00000000000..c76d16e386e --- /dev/null +++ b/superset/security/session_invalidation.py @@ -0,0 +1,205 @@ +# 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. +""" +Backend-agnostic session invalidation. + +Outstanding sessions are terminated by comparing the time a session was +authenticated (``session["_login_at"]``, stamped at login) against a per-user +invalidation epoch (``UserAttribute.sessions_invalidated_at``). When a session +predates the user's epoch it is forced to log out on its next request. + +The epoch is stamped whenever an account is *disabled* (``active`` flips to +``False``), via a SQLAlchemy ``after_update`` listener so it fires regardless of +the code path that disabled the user (FAB admin UI, REST API, or CLI). This +works for both client-side cookie sessions and server-side session stores +without enumerating the store by user. A deleted user is already rejected by +Flask-Login's user loader, so deletion needs no epoch. + +The mechanism is inert until an epoch is set: users that were never disabled +(NULL epoch) are never affected, so it is backwards compatible by default. +""" + +import logging +import math +from datetime import datetime, timezone +from typing import Any, Optional + +from flask import flash, session +from flask_babel import gettext as __ +from flask_login import current_user, logout_user +from sqlalchemy import event, inspect +from sqlalchemy.exc import IntegrityError +from werkzeug.wrappers import Response + +logger = logging.getLogger(__name__) + +#: Session key holding the epoch-seconds timestamp of when the session logged in. +SESSION_LOGIN_AT_KEY = "_login_at" + + +def _utcnow() -> datetime: + return datetime.now(timezone.utc) + + +def stamp_login_time() -> None: + """Record the current session's authentication time. Call on login.""" + session[SESSION_LOGIN_AT_KEY] = _utcnow().timestamp() + + +def _as_utc_timestamp(value: datetime) -> float: + """Epoch seconds for ``value``, treating naive datetimes as UTC. + + The ``sessions_invalidated_at`` column is a naive ``DateTime`` storing a UTC + instant; calling ``.timestamp()`` on a naive datetime would otherwise assume + *local* time and skew the comparison by the local UTC offset. + """ + if value.tzinfo is None: + value = value.replace(tzinfo=timezone.utc) + return value.timestamp() + + +def is_session_invalidated( + login_at: Optional[float], invalidated_at: Optional[datetime] +) -> bool: + """ + Return True if a session authenticated at ``login_at`` is invalidated by an + epoch of ``invalidated_at``. + + - No epoch (``invalidated_at is None``) ⇒ never invalidated (the common case + and the reason the feature is inert/backwards compatible by default). + - An epoch with no recorded login time ⇒ invalidated. A session old enough + to predate the feature carries no ``_login_at``; if the user has since + been disabled, fail closed. + - Otherwise the session is invalidated iff it logged in before the epoch. + """ + if invalidated_at is None: + return False + if login_at is None: + return True + return login_at < _as_utc_timestamp(invalidated_at) + + +def _get_user_invalidated_at(user: Any) -> Optional[datetime]: + extra_attributes = getattr(user, "extra_attributes", None) + if not extra_attributes: + return None + return extra_attributes[0].sessions_invalidated_at + + +def enforce_session_validity() -> Optional[Response]: + """ + ``before_request`` hook: force logout of sessions invalidated by the user's + epoch. + + Fails open — any error here logs a warning and allows the request rather + than risk locking everyone out on a bug in the check. + """ + try: + user = current_user + if not user or not getattr(user, "is_authenticated", False): + return None + # Guest (embedded) users are not FAB users and have their own + # revocation mechanism; skip them. + if getattr(user, "is_guest_user", False): + return None + invalidated_at = _get_user_invalidated_at(user) + if invalidated_at is None: + return None + login_at = session.get(SESSION_LOGIN_AT_KEY) + if not is_session_invalidated(login_at, invalidated_at): + return None + + # Clear the authenticated session and let the request continue as + # anonymous: each route's own decorator then responds correctly for its + # type (401 for the REST API, redirect-to-login for HTML views) without + # this hook needing to know the route kind. + logout_user() + session.clear() + flash(__("Your session has ended. Please sign in again."), "warning") + return None + except Exception: # noqa: BLE001 # pylint: disable=broad-except + logger.warning( + "Session-invalidation check failed; allowing request", exc_info=True + ) + return None + + +def invalidate_user_sessions(connection: Any, user_id: int) -> None: + """Stamp the invalidation epoch for ``user_id`` using ``connection``. + + Upserts the user's ``UserAttribute`` row so the mechanism works even for + users that have no attribute row yet. ``user_attribute.user_id`` carries a + unique constraint, so the insert is safe against a concurrent disable of the + same user: the loser's insert raises ``IntegrityError``, which is caught and + retried as an update rather than creating a duplicate row. + """ + # pylint: disable=import-outside-toplevel + from superset.models.user_attributes import UserAttribute + + table = UserAttribute.__table__ + # Round the epoch up to the next whole second. Some backends (e.g. MySQL) + # store ``DATETIME`` columns without sub-second precision and truncate the + # value; a session that logged in earlier in the same wall-clock second + # carries a fractional ``_login_at`` that would otherwise compare as >= the + # truncated epoch and survive invalidation. Ceiling the stamp guarantees it + # strictly exceeds any login time from the same second. + now_epoch = _utcnow().timestamp() + now = datetime.fromtimestamp(math.ceil(now_epoch), timezone.utc).replace( + tzinfo=None + ) + + def _stamp_existing() -> int: + return connection.execute( + table.update() + .where(table.c.user_id == user_id) + .values(sessions_invalidated_at=now, changed_on=now) + ).rowcount + + if _stamp_existing(): + return + + try: + with connection.begin_nested(): + connection.execute( + table.insert().values( + user_id=user_id, + sessions_invalidated_at=now, + created_on=now, + changed_on=now, + ) + ) + except IntegrityError: + # A concurrent disable inserted the row first; stamp it instead. + _stamp_existing() + + +def _stamp_epoch_on_disable(_mapper: Any, connection: Any, target: Any) -> None: + history = inspect(target).attrs.active.history + # Only act when ``active`` actually changed to False — ignore the + # last_login / login_count updates FAB writes on every login. + if not history.has_changes() or target.active: + return + invalidate_user_sessions(connection, target.id) + + +def register_session_invalidation_events(user_model: Any) -> None: + """Register the ``after_update`` listener that stamps the epoch on disable. + + Idempotent: safe to call on every app initialization (e.g. across tests). + """ + if not event.contains(user_model, "after_update", _stamp_epoch_on_disable): + event.listen(user_model, "after_update", _stamp_epoch_on_disable) diff --git a/tests/integration_tests/security/session_invalidation_tests.py b/tests/integration_tests/security/session_invalidation_tests.py new file mode 100644 index 00000000000..fc957ece619 --- /dev/null +++ b/tests/integration_tests/security/session_invalidation_tests.py @@ -0,0 +1,79 @@ +# 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 superset import db +from superset.models.user_attributes import UserAttribute +from tests.integration_tests.base_tests import SupersetTestCase + +USERNAME = "session_invalidation_user" +PASSWORD = "general" # noqa: S105 + + +class TestSessionInvalidation(SupersetTestCase): + def setUp(self) -> None: + self.create_user( + USERNAME, + PASSWORD, + "Admin", + email="session_invalidation_user@fab.org", + ) + db.session.commit() + + def tearDown(self) -> None: + if user := self.get_user(USERNAME): + db.session.query(UserAttribute).filter_by(user_id=user.id).delete() + db.session.delete(user) + db.session.commit() + + def _set_active(self, active: bool) -> None: + # Update via the ORM so the after_update event fires (mirrors the + # FAB admin UI / REST API path that flips ``active``). + user = self.get_user(USERNAME) + user.active = active + db.session.commit() + + def test_disabling_user_forces_logout_of_active_session(self) -> None: + self.login(USERNAME, PASSWORD) + + # Authenticated request works before the account is disabled. + assert self.client.get("/api/v1/me/").status_code == 200 + + # Disabling the account stamps the per-user invalidation epoch... + self._set_active(False) + user = self.get_user(USERNAME) + attribute = ( + db.session.query(UserAttribute).filter_by(user_id=user.id).one_or_none() + ) + assert attribute is not None + assert attribute.sessions_invalidated_at is not None + + # ...so the previously-authenticated session is now forced out. The + # hook clears the session and the protected REST route answers 401. + assert self.client.get("/api/v1/me/").status_code == 401 + + def test_active_user_session_is_unaffected(self) -> None: + """A user who was never disabled (NULL epoch) is never logged out.""" + self.login(USERNAME, PASSWORD) + assert self.client.get("/api/v1/me/").status_code == 200 + # No epoch was ever stamped. + user = self.get_user(USERNAME) + attribute = ( + db.session.query(UserAttribute).filter_by(user_id=user.id).one_or_none() + ) + assert attribute is None or attribute.sessions_invalidated_at is None + # Repeated requests stay authenticated. + assert self.client.get("/api/v1/me/").status_code == 200 diff --git a/tests/unit_tests/security/audit_log_test.py b/tests/unit_tests/security/audit_log_test.py index 276da24a6f9..25967a3a227 100644 --- a/tests/unit_tests/security/audit_log_test.py +++ b/tests/unit_tests/security/audit_log_test.py @@ -206,9 +206,12 @@ def test_group_api_post_delete_logs_event(mock_log: MagicMock) -> None: # --- Login / Logout --- +@patch("superset.security.session_invalidation.stamp_login_time") @patch("superset.security.manager._log_audit_event") -def test_on_user_login_logs_event(mock_log: MagicMock) -> None: - """on_user_login logs a UserLoggedIn event.""" +def test_on_user_login_logs_event( + mock_log: MagicMock, mock_stamp_login_time: MagicMock +) -> None: + """on_user_login logs a UserLoggedIn event and stamps the session.""" sm = SupersetSecurityManager.__new__(SupersetSecurityManager) user = MagicMock(spec=User) user.username = "testuser" @@ -216,6 +219,7 @@ def test_on_user_login_logs_event(mock_log: MagicMock) -> None: sm.on_user_login(user) + mock_stamp_login_time.assert_called_once() mock_log.assert_called_once_with( "UserLoggedIn", {"username": "testuser", "user_id": 7} ) diff --git a/tests/unit_tests/security/test_session_invalidation.py b/tests/unit_tests/security/test_session_invalidation.py new file mode 100644 index 00000000000..85f5e1c70fc --- /dev/null +++ b/tests/unit_tests/security/test_session_invalidation.py @@ -0,0 +1,214 @@ +# 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 datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from sqlalchemy.exc import IntegrityError + +from superset.security.session_invalidation import ( + _as_utc_timestamp, + enforce_session_validity, + invalidate_user_sessions, + is_session_invalidated, + SESSION_LOGIN_AT_KEY, +) + + +def test_no_epoch_is_never_invalidated() -> None: + """A user that was never disabled (NULL epoch) is never invalidated.""" + assert is_session_invalidated(login_at=None, invalidated_at=None) is False + assert is_session_invalidated(login_at=1_000.0, invalidated_at=None) is False + + +def test_epoch_with_no_login_time_fails_closed() -> None: + """A pre-feature session (no _login_at) on a disabled user is invalidated.""" + epoch = datetime.now(timezone.utc) + assert is_session_invalidated(login_at=None, invalidated_at=epoch) is True + + +def test_session_before_epoch_is_invalidated() -> None: + epoch = datetime.now(timezone.utc) + before = (epoch - timedelta(minutes=5)).timestamp() + assert is_session_invalidated(login_at=before, invalidated_at=epoch) is True + + +def test_session_after_epoch_is_valid() -> None: + """A fresh login after a disable+re-enable must not be invalidated.""" + epoch = datetime.now(timezone.utc) + after = (epoch + timedelta(minutes=5)).timestamp() + assert is_session_invalidated(login_at=after, invalidated_at=epoch) is False + + +def test_login_exactly_at_epoch_is_valid() -> None: + epoch = datetime.now(timezone.utc) + assert ( + is_session_invalidated(login_at=epoch.timestamp(), invalidated_at=epoch) + is False + ) + + +def test_naive_epoch_is_treated_as_utc() -> None: + """ + The DB column is a naive UTC ``DateTime``; the comparison must treat it as + UTC, not local time (otherwise it skews by the local offset). + """ + aware = datetime(2026, 6, 2, 12, 0, 0, tzinfo=timezone.utc) + naive = aware.replace(tzinfo=None) + assert _as_utc_timestamp(naive) == aware.timestamp() + + just_before = aware.timestamp() - 1 + just_after = aware.timestamp() + 1 + assert is_session_invalidated(login_at=just_before, invalidated_at=naive) is True + assert is_session_invalidated(login_at=just_after, invalidated_at=naive) is False + + +MODULE = "superset.security.session_invalidation" + + +def _user(*, authenticated: bool = True, guest: bool = False) -> SimpleNamespace: + return SimpleNamespace(is_authenticated=authenticated, is_guest_user=guest) + + +def test_enforce_skips_unauthenticated_user() -> None: + """No authenticated user ⇒ nothing to enforce, request proceeds.""" + with ( + patch(f"{MODULE}.current_user", _user(authenticated=False)), + patch(f"{MODULE}.logout_user") as logout, + ): + assert enforce_session_validity() is None + logout.assert_not_called() + + +def test_enforce_skips_guest_user() -> None: + """Guest (embedded) users have their own revocation path and are skipped.""" + with ( + patch(f"{MODULE}.current_user", _user(guest=True)), + patch(f"{MODULE}._get_user_invalidated_at") as get_epoch, + patch(f"{MODULE}.logout_user") as logout, + ): + assert enforce_session_validity() is None + get_epoch.assert_not_called() + logout.assert_not_called() + + +def test_enforce_no_epoch_leaves_session_alone() -> None: + """A user with no invalidation epoch is never logged out.""" + with ( + patch(f"{MODULE}.current_user", _user()), + patch(f"{MODULE}._get_user_invalidated_at", return_value=None), + patch(f"{MODULE}.logout_user") as logout, + ): + assert enforce_session_validity() is None + logout.assert_not_called() + + +def test_enforce_valid_session_is_not_logged_out() -> None: + """A session that logged in after the epoch stays authenticated.""" + epoch = datetime.now(timezone.utc) + after = (epoch + timedelta(minutes=5)).timestamp() + fake_session = MagicMock() + fake_session.get.return_value = after + with ( + patch(f"{MODULE}.current_user", _user()), + patch(f"{MODULE}._get_user_invalidated_at", return_value=epoch), + patch(f"{MODULE}.session", fake_session), + patch(f"{MODULE}.logout_user") as logout, + ): + assert enforce_session_validity() is None + fake_session.get.assert_called_once_with(SESSION_LOGIN_AT_KEY) + logout.assert_not_called() + + +def test_enforce_invalidated_session_is_logged_out() -> None: + """A session predating the epoch is cleared and flashed a warning.""" + epoch = datetime.now(timezone.utc) + before = (epoch - timedelta(minutes=5)).timestamp() + fake_session = MagicMock() + fake_session.get.return_value = before + with ( + patch(f"{MODULE}.current_user", _user()), + patch(f"{MODULE}._get_user_invalidated_at", return_value=epoch), + patch(f"{MODULE}.session", fake_session), + patch(f"{MODULE}.logout_user") as logout, + patch(f"{MODULE}.flash") as flash, + ): + assert enforce_session_validity() is None + logout.assert_called_once() + fake_session.clear.assert_called_once() + flash.assert_called_once() + + +def test_enforce_fails_open_on_error() -> None: + """Any error in the check logs a warning and allows the request.""" + # A real (non-guest, authenticated) user so the check reaches the epoch + # lookup, which then raises — exercising the fail-open handler. + user = SimpleNamespace(is_authenticated=True, is_guest_user=False) + with ( + patch(f"{MODULE}.current_user", user), + patch(f"{MODULE}._get_user_invalidated_at", side_effect=RuntimeError("boom")), + patch(f"{MODULE}.logout_user") as logout, + patch(f"{MODULE}.logger") as logger, + ): + assert enforce_session_validity() is None + logout.assert_not_called() + logger.warning.assert_called_once() + + +def test_invalidate_updates_existing_row() -> None: + """When a row already exists, the upsert updates it and skips the insert.""" + connection = MagicMock() + connection.execute.return_value.rowcount = 1 + invalidate_user_sessions(connection, user_id=7) + # Single UPDATE, no INSERT / SAVEPOINT. + assert connection.execute.call_count == 1 + connection.begin_nested.assert_not_called() + + +def test_invalidate_inserts_when_missing() -> None: + """When no row exists, the upsert inserts one inside a SAVEPOINT.""" + connection = MagicMock() + # First execute is the UPDATE (rowcount 0); second is the INSERT. + connection.execute.return_value.rowcount = 0 + invalidate_user_sessions(connection, user_id=7) + assert connection.execute.call_count == 2 + connection.begin_nested.assert_called_once() + + +def test_invalidate_retries_as_update_on_race() -> None: + """If a concurrent disable wins the insert, the IntegrityError is caught + and the row is stamped via UPDATE instead of duplicating it.""" + connection = MagicMock() + update_result = SimpleNamespace(rowcount=0) + retry_result = SimpleNamespace(rowcount=1) + + calls: list[str] = [] + + def execute(statement, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + compiled = str(statement).strip().upper() + if compiled.startswith("UPDATE"): + calls.append("update") + # First UPDATE finds nothing; the retry UPDATE succeeds. + return retry_result if len(calls) > 1 else update_result + calls.append("insert") + raise IntegrityError("insert", {}, Exception()) + + connection.execute.side_effect = execute + invalidate_user_sessions(connection, user_id=7) + # update (miss) -> insert (race loses) -> update (retry succeeds) + assert calls == ["update", "insert", "update"]