feat: filter query results by user db_role (ENG-2477,ENG-2476,ENG-2478) (#47)

If the current user has a `db_role`, they should only see query results
that they have generated, so that they don't see results which contain
info about resources they don't have permission to view.

feat: use per-user db role for query exec (ENG-2474) (#48)

* feat: use per-user db role for query exec (ENG-2474)

Login with per-user PG database role, if available, to ensure that RLS
policies are applied to user queries.

* Reject login from unknown SSO users

* Use pre-filtered query for QueryResults rather than session-level `set role`

* Add docstring with non-obvious context

Author: johnsca

redash/authentication/__init__.py

@@ -191,6 +191,9 @@ 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")
 
     # it might actually be a username or something, but it doesn't actually matter
     email = identity

redash/models/__init__.py

@@ -7,6 +7,7 @@
 import pytz
 from sqlalchemy import UniqueConstraint, and_, cast, distinct, func, or_
 from sqlalchemy.dialects.postgresql import ARRAY, DOUBLE_PRECISION, JSONB
+from flask_login import current_user
 from sqlalchemy.event import listens_for
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import (
@@ -36,6 +37,7 @@
     gfk_type,
     key_type,
     primary_key,
+    BaseQuery,
 )
 from redash.models.changes import Change, ChangeTrackingMixin  # noqa
 from redash.models.mixins import BelongsToOrgMixin, TimestampMixin
@@ -346,8 +348,11 @@ def unused(cls, days=7):
         )
 
     @classmethod
-    def get_latest(cls, data_source, query, max_age=0):
-        query_hash = gen_query_hash(query)
+    def get_latest(cls, data_source, query, max_age=0, is_hash=False):
+        if is_hash:
+            query_hash = query
+        else:
+            query_hash = gen_query_hash(query)
 
         if max_age == -1 and settings.QUERY_RESULTS_EXPIRED_TTL_ENABLED:
             max_age = settings.QUERY_RESULTS_EXPIRED_TTL
@@ -391,6 +396,35 @@ def groups(self):
         return self.data_source.groups
 
 
+@listens_for(BaseQuery, "before_compile", retval=True)
+def prefilter_query_results(query):
+    """
+    Ensure that a user with a db_role defined can only see QueryResults that
+    they themselves created.
+
+    This is to ensure that they don't see results that might include resources
+    from accounts that shouldn't be visibile to them. Ideally, this would use
+    `set role` and the RLS policy that is applied to the table, but without a
+    "post-query" type event, that's not really feasible.
+
+    The RLS policy on the redash.query_results table still applies to the
+    arbitrary query that the user executes, so that they can't issue a query
+    directly against that table and get around this check.
+    """
+    for desc in query.column_descriptions:
+        if desc['type'] is QueryResult:
+            db_role = getattr(current_user, "db_role", None)
+            if not db_role:
+                continue
+            limit = query._limit
+            offset = query._offset
+            query = query.limit(None).offset(None)
+            query.offset(None)
+            query = query.filter(desc['entity'].db_role == db_role)
+            query = query.limit(limit).offset(offset)
+    return query
+
+
 def should_schedule_next(previous_iteration, now, interval, time=None, day_of_week=None, failures=0):
     # if previous_iteration is None, it means the query has never been run before
     # so we should schedule it immediately

redash/query_runner/pg.py

@@ -1,3 +1,4 @@
+import hashlib
 import logging
 import os
 import select
@@ -20,7 +21,7 @@
     JobTimeoutException,
     register,
 )
-
+from redash import settings
 from redash.stacklet.auth import inject_iam_auth
 
 logger = logging.getLogger(__name__)
@@ -251,12 +252,29 @@ def _get_tables(self, schema):
 
         return list(schema.values())
 
+    def _gen_role_pass(self, role_name: str) -> str:
+        """
+        Generate a password for a given role using the datasource secret and role name.
+        """
+        secret = settings.DATASOURCE_SECRET_KEY
+        return hashlib.sha256(f"{secret}:{role_name}".encode("utf-8")).hexdigest()
+
     @inject_iam_auth
-    def _get_connection(self):
+    def _get_connection(self, user):
+        if getattr(user, "db_role", None):
+            auth_config = dict(
+                user=user.db_role,
+                password=self._gen_role_pass(user.db_role),
+            )
+        else:
+            auth_config = dict(
+                user=self.configuration.get("user"),
+                password=self.configuration.get("password"),
+            )
+        logger.info(f"Connecting to datasource as {auth_config['user']}")
         self.ssl_config = _get_ssl_config(self.configuration)
         connection = psycopg2.connect(
-            user=self.configuration.get("user"),
-            password=self.configuration.get("password"),
+            **auth_config,
             host=self.configuration.get("host"),
             port=self.configuration.get("port"),
             dbname=self.configuration.get("dbname"),
@@ -267,7 +285,7 @@ def _get_connection(self):
         return connection
 
     def run_query(self, query, user):
-        connection = self._get_connection()
+        connection = self._get_connection(user)
         _wait(connection, timeout=10)
 
         cursor = connection.cursor()

redash/serializers/__init__.py

@@ -137,6 +137,18 @@ def serialize_query(
     if with_visualizations:
         d["visualizations"] = [serialize_visualization(vis, with_query=False) for vis in query.visualizations]
 
+    if getattr(current_user, "db_role", None):
+        # Override the latest_query_data_id for users with a db_role because
+        # they may not actually be able to see that one due to their db_role
+        # and may have one specific to them instead.
+        latest_result = models.QueryResult.get_latest(
+            data_source=query.data_source,
+            query=query.query_hash,
+            max_age=-1,
+            is_hash=True,
+        )
+        d["latest_query_data_id"] = latest_result and latest_result.id or None
+
     return d