From b594a8dcc0ec1a3bd8a07ab5828ead62ee14b015 Mon Sep 17 00:00:00 2001
From: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com>
Date: Mon, 6 May 2024 03:37:59 -0400
Subject: [PATCH] Fix custom actions in security manager `has_access` (#39421)

---
 airflow/www/security_manager.py    |  2 +-
 tests/www/test_security_manager.py | 36 ++++++++++++++++++++++++------
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/airflow/www/security_manager.py b/airflow/www/security_manager.py
index 88c2caa32a3d4..316daace7d617 100644
--- 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
index 81a05e5fd063e..2051ccef095e2 100644
--- 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,6 +92,7 @@ class TestAirflowSecurityManagerV2:
                 False,
             ),
             (
+                ACTION_CAN_READ,
                 RESOURCE_BROWSE_MENU,
                 {
                     "is_authorized_dag": False,
@@ -96,11 +100,29 @@ class TestAirflowSecurityManagerV2:
                 },
                 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()
@@ -108,7 +130,7 @@ def test_has_access(
             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: