feat: read the stacklet:rls_access claim for the RLS db role (ENG-9644)

redash/authentication/__init__.py

@@ -2,6 +2,7 @@
 import hmac
 import json
 import logging
+import re
 import time
 from datetime import timedelta
 from urllib.parse import quote_plus, urlsplit, urlunsplit
@@ -193,9 +194,8 @@ def jwt_token_load_user_from_request(request):
 
     if not valid_token:
         return None
-
-    if payload.get("stacklet:db_role") == "limited_visibility":
-        raise Unauthorized("Unable to determine SSO identity")
+
+    db_role = _db_role_from_claims(payload)
 
     # it might actually be a username or something, but it doesn't actually matter
     email = identity
@@ -220,14 +220,45 @@ def jwt_token_load_user_from_request(request):
                 user_groups.discard("admin")
             user.update_group_assignments(user_groups)
 
-    db_role = payload.get("stacklet:db_role") or None  # force None instead of empty string, JIC
     if db_role != user.db_role:
         user.db_role = db_role
         models.db.session.commit()
 
     return user
 
 
+# `stacklet:rls_access` is either one of the reserved values below, or the
+# stem the platform derives for a scoped user's RLS identity: up to 20
+# sanitised characters, then `_`, then 6 hex characters of a hash. The
+# reserved values can't match the stem pattern.
+RLS_ACCESS_UNRESTRICTED = "unrestricted"
+RLS_ACCESS_NO_VISIBILITY = "no_visibility"
+RLS_STEM_RE = re.compile(r"[a-z0-9_]+_[0-9a-f]{6}")
+
+
+def _db_role_from_claims(payload):
+    """Return the database role to run queries as, or None for the data source user.
+
+    Reads `stacklet:rls_access`, falling back to the legacy `stacklet:db_role`
+    only for tokens minted before the platform sent it. Raises Unauthorized
+    for users who must not see any data, and for malformed values.
+    """
+    if "stacklet:rls_access" not in payload:
+        db_role = payload.get("stacklet:db_role") or None
+        if db_role in ("limited_visibility", RLS_ACCESS_NO_VISIBILITY):
+            raise Unauthorized("Unable to determine SSO identity")
+        return db_role
+
+    access = payload["stacklet:rls_access"]
+    if access == RLS_ACCESS_UNRESTRICTED:
+        return None
+    if access == RLS_ACCESS_NO_VISIBILITY:
+        raise Unauthorized("Unable to determine SSO identity")
+    if isinstance(access, str) and RLS_STEM_RE.fullmatch(access):
+        return "sso_" + access
+    raise Unauthorized("Invalid stacklet:rls_access claim")
+
+
 def log_user_logged_in(app, user):
     event = {
         "org_id": user.org_id,

tests/test_authentication.py

@@ -9,6 +9,7 @@
 import requests
 from flask import request
 from mock import Mock, patch
+from sqlalchemy.orm import Session
 from sqlalchemy.orm.exc import NoResultFound
 
 from redash import models, settings
@@ -517,6 +518,79 @@ def test_jwt_rejects_an_audience_not_configured(self):
         response = self.get_request("/data_sources", org=self.factory.org, headers={self.token_name: token_data})
         self.assertEqual(response.status_code, 401)
 
+    def _get_with_claims(self, user, **claims):
+        issued_at_timestamp = time.time()
+        data = {
+            "aud": self.auth_audience,
+            "email": user.email,
+            "exp": issued_at_timestamp + 60,
+            "iat": issued_at_timestamp,
+            "iss": self.auth_issuer,
+            **claims,
+        }
+        with open(self.rsa_private_key) as keyfile:
+            sign_key = keyfile.read().strip()
+        token_data = jwt.encode(data, sign_key, algorithm="RS256")
+        return self.get_request("/data_sources", org=self.factory.org, headers={self.token_name: token_data})
+
+    def _stored_db_role(self, user):
+        # a separate session releases its connection on exit, or the next
+        # test's drop_all blocks on the lock left by the read
+        with Session(models.db.engine) as session:
+            return session.get(models.User, user.id).db_role
+
+    def test_jwt_rls_access_sets_sso_db_role(self):
+        user = self.factory.create_user()
+        response = self._get_with_claims(user, **{"stacklet:rls_access": "jane_doe_example_com_a1b2c3"})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(self._stored_db_role(user), "sso_jane_doe_example_com_a1b2c3")
+
+    def test_jwt_rls_access_takes_precedence_over_db_role(self):
+        user = self.factory.create_user()
+        response = self._get_with_claims(
+            user,
+            **{"stacklet:rls_access": "jane_a1b2c3", "stacklet:db_role": "sso_other_d4e5f6"},
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(self._stored_db_role(user), "sso_jane_a1b2c3")
+
+    def test_jwt_rls_access_unrestricted_clears_db_role(self):
+        user = self.factory.create_user(db_role="sso_jane_a1b2c3")
+        response = self._get_with_claims(
+            user,
+            **{"stacklet:rls_access": "unrestricted", "stacklet:db_role": "sso_jane_a1b2c3"},
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertIsNone(self._stored_db_role(user))
+
+    def test_jwt_rls_access_no_visibility_is_denied(self):
+        user = self.factory.create_user()
+        response = self._get_with_claims(user, **{"stacklet:rls_access": "no_visibility"})
+        self.assertEqual(response.status_code, 401)
+
+    def test_jwt_rls_access_malformed_is_denied(self):
+        user = self.factory.create_user()
+        for access in ("", "jane", "Jane_a1b2c3", "jane_a1b2c", "jane_a1b2c3\n", "sso_jane_a1b2cz", None, 123):
+            with self.subTest(access=access):
+                response = self._get_with_claims(user, **{"stacklet:rls_access": access})
+                self.assertEqual(response.status_code, 401)
+
+    def test_jwt_legacy_db_role_without_rls_access(self):
+        user = self.factory.create_user()
+        response = self._get_with_claims(user, **{"stacklet:db_role": "sso_jane_a1b2c3"})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(self._stored_db_role(user), "sso_jane_a1b2c3")
+
+    def test_jwt_legacy_limited_visibility_db_role_is_denied(self):
+        user = self.factory.create_user()
+        response = self._get_with_claims(user, **{"stacklet:db_role": "limited_visibility"})
+        self.assertEqual(response.status_code, 401)
+
+    def test_jwt_legacy_no_visibility_db_role_is_denied(self):
+        user = self.factory.create_user()
+        response = self._get_with_claims(user, **{"stacklet:db_role": "no_visibility"})
+        self.assertEqual(response.status_code, 401)
+
     @patch.object(requests, "get")
     def test_jwk_decode(self, mock_get):
         with open(self.rsa_public_key, "rb") as keyfile: