fix: Scope offline server permission checks to the requested project

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]

sdk/python/tests/unit/test_offline_server.py

@@ -1,22 +1,33 @@
+import contextlib
+import json
 import os
 import subprocess
 import sys
 import textwrap
 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 (
     RemoteOfflineStore,
     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")