Fix custom actions in security manager has_access (#39421)

apache · May 6, 2024 · b594a8d · b594a8d
1 parent 9eac305
commit b594a8d
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 8 deletions.
diff --git a/airflow/www/security_manager.py b/airflow/www/security_manager.py
@@ -337,7 +337,7 @@ def _get_auth_manager_is_authorized_method(self, fab_resource_name: str) -> Call
             # The user is trying to access a page specific to the auth manager
             # (e.g. the user list view in FabAuthManager) or a page defined in a plugin
             return lambda action, resource_pk, user: get_auth_manager().is_authorized_custom_view(
-                method=get_method_from_fab_action_map()[action],
+                method=get_method_from_fab_action_map().get(action, action),
                 resource_name=fab_resource_name,
                 user=user,
             )

diff --git a/tests/www/test_security_manager.py b/tests/www/test_security_manager.py
@@ -50,13 +50,14 @@ def security_manager(app_builder):
 @pytest.mark.db_test
 class TestAirflowSecurityManagerV2:
     @pytest.mark.parametrize(
-        "resource_name, auth_manager_methods, expected",
+        "action_name, resource_name, auth_manager_methods, expected",
         [
-            (RESOURCE_VARIABLE, {"is_authorized_variable": True}, True),
-            (RESOURCE_VARIABLE, {"is_authorized_variable": False}, False),
-            (RESOURCE_DOCS_MENU, {"is_authorized_view": True}, True),
-            (RESOURCE_DOCS_MENU, {"is_authorized_view": False}, False),
+            (ACTION_CAN_READ, RESOURCE_VARIABLE, {"is_authorized_variable": True}, True),
+            (ACTION_CAN_READ, RESOURCE_VARIABLE, {"is_authorized_variable": False}, False),
+            (ACTION_CAN_READ, RESOURCE_DOCS_MENU, {"is_authorized_view": True}, True),
+            (ACTION_CAN_READ, RESOURCE_DOCS_MENU, {"is_authorized_view": False}, False),
             (
+                ACTION_CAN_READ,
                 RESOURCE_ADMIN_MENU,
                 {
                     "is_authorized_view": False,
@@ -69,6 +70,7 @@ class TestAirflowSecurityManagerV2:
                 True,
             ),
             (
+                ACTION_CAN_READ,
                 RESOURCE_ADMIN_MENU,
                 {
                     "is_authorized_view": False,
@@ -81,6 +83,7 @@ class TestAirflowSecurityManagerV2:
                 False,
             ),
             (
+                ACTION_CAN_READ,
                 RESOURCE_BROWSE_MENU,
                 {
                     "is_authorized_dag": False,
@@ -89,26 +92,45 @@ class TestAirflowSecurityManagerV2:
                 False,
             ),
             (
+                ACTION_CAN_READ,
                 RESOURCE_BROWSE_MENU,
                 {
                     "is_authorized_dag": False,
                     "is_authorized_view": True,
                 },
                 True,
             ),
+            (
+                "can_not_a_default_action",
+                "custom_resource",
+                {"is_authorized_custom_view": False},
+                False,
+            ),
+            (
+                "can_not_a_default_action",
+                "custom_resource",
+                {"is_authorized_custom_view": True},
+                True,
+            ),
         ],
     )
     @mock.patch("airflow.www.security_manager.get_auth_manager")
     def test_has_access(
-        self, mock_get_auth_manager, security_manager, resource_name, auth_manager_methods, expected
+        self,
+        mock_get_auth_manager,
+        security_manager,
+        action_name,
+        resource_name,
+        auth_manager_methods,
+        expected,
     ):
         user = Mock()
         auth_manager = Mock()
         for method_name, method_return in auth_manager_methods.items():
             method = Mock(return_value=method_return)
             getattr(auth_manager, method_name).side_effect = method
             mock_get_auth_manager.return_value = auth_manager
-        result = security_manager.has_access(ACTION_CAN_READ, resource_name, user=user)
+        result = security_manager.has_access(action_name, resource_name, user=user)
         assert result == expected
         if len(auth_manager_methods) > 1 and not expected:
             for method_name in auth_manager_methods: