diff --git a/redash/authentication/__init__.py b/redash/authentication/__init__.py index 1c41b6420c..4742f7462d 100644 --- a/redash/authentication/__init__.py +++ b/redash/authentication/__init__.py @@ -196,6 +196,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 diff --git a/redash/models/__init__.py b/redash/models/__init__.py index 9ce86b7612..36488603f2 100644 --- a/redash/models/__init__.py +++ b/redash/models/__init__.py @@ -5,6 +5,7 @@ import numbers import pytz +from flask_login import current_user from sqlalchemy import distinct, or_, and_, UniqueConstraint, cast from sqlalchemy.dialects import postgresql from sqlalchemy.event import listens_for @@ -42,7 +43,7 @@ from redash.utils.configuration import ConfigurationContainer from redash.models.parameterized_query import ParameterizedQuery -from .base import db, gfk_type, Column, GFKBase, SearchBaseQuery, key_type, primary_key +from .base import db, gfk_type, Column, GFKBase, BaseQuery, SearchBaseQuery, key_type, primary_key from .changes import ChangeTrackingMixin, Change # noqa from .mixins import BelongsToOrgMixin, TimestampMixin from .organizations import Organization @@ -409,6 +410,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 ): diff --git a/redash/models/users.py b/redash/models/users.py index d01e48df52..0097a13182 100644 --- a/redash/models/users.py +++ b/redash/models/users.py @@ -54,34 +54,12 @@ def update_user_active_at(sender, *args, **kwargs): redis_connection.hset(LAST_ACTIVE_KEY, current_user.id, int(time.time())) -def set_db_role(sender, *args, **kwargs): - """ - Check to see if the current user has a db_role assigned, and, - if so, set it for the session. - """ - logger.info(f"Checking for DB role on user: {current_user}") - db_role = getattr(current_user, "db_role", None) - if db_role: - logger.info(f"Setting session DB role: {db_role}") - db.session.execute("SET ROLE :db_role", {"db_role": db_role}) - - -def reset_db_role(sender, *args, **kwargs): - """ - Reset any DB role set for the current user. - """ - logger.info("Resetting session DB role") - db.session.execute("RESET ROLE") - - def init_app(app): """ A Flask extension to keep user details updates in Redis and sync it periodically to the database (User.details). """ request_started.connect(update_user_active_at, app) - request_started.connect(set_db_role, app) - request_finished.connect(reset_db_role, app) class PermissionsCheckMixin(object): diff --git a/redash/query_runner/pg.py b/redash/query_runner/pg.py index 67608fa223..6f2e1eb9ac 100644 --- a/redash/query_runner/pg.py +++ b/redash/query_runner/pg.py @@ -1,3 +1,4 @@ +import hashlib import os import logging import select @@ -9,6 +10,7 @@ import psycopg2 from psycopg2.extras import Range +from redash import settings from redash.query_runner import * from redash.utils import JSONEncoder, json_dumps, json_loads @@ -238,12 +240,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"), @@ -254,7 +273,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()