diff --git a/sdk/python/feast/offline_server.py b/sdk/python/feast/offline_server.py index e82b5239767..ce8de1c9981 100644 --- a/sdk/python/feast/offline_server.py +++ b/sdk/python/feast/offline_server.py @@ -56,7 +56,10 @@ def _configure_grpc_fips() -> bool: get_offline_store_from_config, ) from feast.permissions.action import AuthzedAction # noqa: E402 -from feast.permissions.security_manager import assert_permissions # noqa: E402 +from feast.permissions.security_manager import ( # noqa: E402 + assert_permissions, + get_security_manager, +) from feast.permissions.server.arrow import ( # noqa: E402 AuthorizationMiddlewareFactory, inject_user_details_decorator, @@ -192,6 +195,17 @@ def _call_api(self, api: str, command: dict, key: str): assert api is not None, "api can not be empty" remove_data = False + # Permission checks in the handlers below read the SecurityManager's permission + # list, which is loaded per project, so it has to follow the project this request + # names -- otherwise a caller reaches another project's resources through + # whichever project the server itself was started from. The SecurityManager keeps + # its own ContextVar, separate from the store's, and `permissions` reads that one, + # so binding the store alone does not scope the check. Bind both: the security + # manager for the policy, the store for the objects the handlers resolve. + project = command.get("project") + sm = get_security_manager() + sm_token = sm.set_current_project(project) if sm is not None else None + project_token = self.store.set_current_project(project) try: if api == OfflineServer.offline_write_batch.__name__: self.offline_write_batch(command, key) @@ -211,6 +225,9 @@ def _call_api(self, api: str, command: dict, key: str): traceback.print_exc() raise e finally: + self.store.reset_current_project(project_token) + if sm is not None and sm_token is not None: + sm.reset_current_project(sm_token) if remove_data: # Get service is consumed, so we clear the corresponding flight and data del self.flights[key] @@ -284,6 +301,12 @@ def do_get(self, context: fl.ServerCallContext, ticket: fl.Ticket): api = command["api"] logger.debug(f"get command is {command}") logger.debug(f"requested api is {api}") + # As in _call_api, and for the same reason, bind both the security manager and + # the store to the project this request names. + project = command.get("project") + sm = get_security_manager() + sm_token = sm.set_current_project(project) if sm is not None else None + project_token = self.store.set_current_project(project) try: if api == OfflineServer.get_historical_features.__name__: table = self.get_historical_features(command, key).to_arrow() @@ -302,6 +325,10 @@ def do_get(self, context: fl.ServerCallContext, ticket: fl.Ticket): logger.exception(e) traceback.print_exc() raise e + finally: + self.store.reset_current_project(project_token) + if sm is not None and sm_token is not None: + sm.reset_current_project(sm_token) # Get service is consumed, so we clear the corresponding flight and data del self.flights[key] diff --git a/sdk/python/tests/unit/test_offline_server.py b/sdk/python/tests/unit/test_offline_server.py index 321b19bb9b7..d970a9d6605 100644 --- a/sdk/python/tests/unit/test_offline_server.py +++ b/sdk/python/tests/unit/test_offline_server.py @@ -1,3 +1,5 @@ +import contextlib +import json import os import subprocess import sys @@ -5,6 +7,8 @@ from unittest.mock import MagicMock, mock_open, patch import assertpy +import pyarrow as pa +import pyarrow.flight as fl import pytest from feast.infra.offline_stores.remote import ( @@ -12,11 +16,18 @@ RemoteOfflineStoreConfig, _create_retrieval_metadata, ) +from feast.infra.registry.base_registry import BaseRegistry from feast.offline_server import ( OfflineServer, _configure_grpc_fips, _is_fips_enabled, ) +from feast.permissions.security_manager import ( + SecurityManager, + get_security_manager, + no_security_manager, + set_security_manager, +) def test_create_retrieval_metadata_with_sql_string(): @@ -208,3 +219,143 @@ def tracking_import(name, *args, **kwargs): assert result.returncode == 0, ( f"Subprocess failed:\nstdout: {result.stdout}\nstderr: {result.stderr}" ) + + +_SERVER_HOME_PROJECT = "the_project_the_server_was_started_from" + + +def _server_for(command): + """Build a mocked OfflineServer plus the flight key for a command.""" + key = ("command_id", json.dumps(command)) + server = MagicMock(spec=OfflineServer) + server.store = MagicMock() + server.store.set_current_project.return_value = "project-token" + server.flights = {key: MagicMock()} + return server, key + + +@contextlib.contextmanager +def _recording_security_manager(seen): + """ + Install a SecurityManager whose registry records the project every permission + lookup resolves to. + + Asserting on `store.set_current_project` alone would not show whether the check is + scoped: the store and the SecurityManager hold separate ContextVars, and + `SecurityManager.permissions` reads its own. Recording the project the registry is + asked for is the thing the fix is actually about. + """ + registry = MagicMock(spec=BaseRegistry) + registry.list_permissions.side_effect = lambda project=None, **kwargs: ( + seen.append(project) or [] + ) + sm = SecurityManager(project=_SERVER_HOME_PROJECT, registry=registry) + set_security_manager(sm) + try: + yield sm + finally: + no_security_manager() + + +def _permission_check(): + """Stand in for the assert_permissions call every real handler makes.""" + get_security_manager().permissions + + +def test_do_get_scopes_the_permission_lookup_to_the_requested_project(): + """ + The permission list is loaded per project, so it has to follow the project the + request names. Otherwise a caller reaches every project the server serves through + whichever project the server itself was started from. + """ + command = {"api": "get_historical_features", "project": "project_b"} + server, key = _server_for(command) + seen = [] + + def handler(*args, **kwargs): + _permission_check() + result = MagicMock() + result.to_arrow.return_value = pa.table({"a": [1]}) + return result + + server.get_historical_features.side_effect = handler + + # do_get is wrapped by inject_user_details_decorator, which returns early when + # the call carries no `auth` middleware. + context = MagicMock() + context.get_middleware.return_value = None + + with _recording_security_manager(seen) as sm: + OfflineServer.do_get( + server, context=context, ticket=fl.Ticket(ticket=str(key).encode()) + ) + sm.permissions # after the dispatcher returns + + assert seen == ["project_b", _SERVER_HOME_PROJECT], ( + "the check inside the handler must resolve to the requested project, and the " + "binding must be gone once the request is done" + ) + server.store.set_current_project.assert_called_once_with("project_b") + server.store.reset_current_project.assert_called_once_with("project-token") + + +def test_call_api_scopes_the_permission_lookup_to_the_requested_project(): + """The put-side dispatcher scopes the permission lookup the same way.""" + command = {"api": "validate_data_source", "project": "project_b"} + server, key = _server_for(command) + seen = [] + server.validate_data_source.side_effect = lambda *a, **k: _permission_check() + + with _recording_security_manager(seen) as sm: + OfflineServer._call_api(server, command["api"], command, key) + sm.permissions + + assert seen == ["project_b", _SERVER_HOME_PROJECT] + server.store.set_current_project.assert_called_once_with("project_b") + server.store.reset_current_project.assert_called_once_with("project-token") + + +def test_call_api_resets_the_project_when_the_handler_raises(): + """A failed request must not leave its project bound for the next one.""" + command = {"api": "validate_data_source", "project": "project_b"} + server, key = _server_for(command) + seen = [] + server.validate_data_source.side_effect = RuntimeError("boom") + + with _recording_security_manager(seen) as sm: + with pytest.raises(RuntimeError): + OfflineServer._call_api(server, command["api"], command, key) + sm.permissions + + assert seen == [_SERVER_HOME_PROJECT] + server.store.reset_current_project.assert_called_once_with("project-token") + + +def test_call_api_without_a_project_falls_back_to_the_servers_own(): + """ + A command that carries no project passes `None`, which the SecurityManager falls + back from to the project it was built with -- the behaviour before this scoping + existed. + """ + command = {"api": "validate_data_source"} + server, key = _server_for(command) + seen = [] + server.validate_data_source.side_effect = lambda *a, **k: _permission_check() + + with _recording_security_manager(seen): + OfflineServer._call_api(server, command["api"], command, key) + + assert seen == [_SERVER_HOME_PROJECT] + server.store.set_current_project.assert_called_once_with(None) + + +def test_dispatchers_work_without_a_security_manager(): + """An unauthenticated deployment has no SecurityManager at all.""" + no_security_manager() + command = {"api": "validate_data_source", "project": "project_b"} + server, key = _server_for(command) + + OfflineServer._call_api(server, command["api"], command, key) + + server.store.set_current_project.assert_called_once_with("project_b") + server.store.reset_current_project.assert_called_once_with("project-token")