Block extensions disallowed by policy

azurelinuxagent/ga/exthandlers.py

@@ -39,6 +39,7 @@
 from azurelinuxagent.common.agent_supported_feature import get_agent_supported_features_list_for_extensions, \
     SupportedFeatureNames, get_supported_feature_by_name, get_agent_supported_features_list_for_crp
 from azurelinuxagent.ga.cgroupconfigurator import CGroupConfigurator
+from azurelinuxagent.ga.policy.policy_engine import ExtensionPolicyEngine, ExtensionPolicyError
 from azurelinuxagent.common.datacontract import get_properties, set_properties
 from azurelinuxagent.common.errorstate import ErrorState
 from azurelinuxagent.common.event import add_event, elapsed_milliseconds, WALAEventOperation, \
@@ -482,10 +483,49 @@ def handle_ext_handlers(self, goal_state_id):
         depends_on_err_msg = None
         extensions_enabled = conf.get_extensions_enabled()
 
+        # Instantiate policy engine, and use same engine to handle all extension handlers.
+        # If an error is thrown during policy engine initialization, we block all extensions and report the error via handler/extension status for
+        # each extension.
+        policy_error = None
+        try:
+            policy_engine = ExtensionPolicyEngine()
+        except Exception as ex:
+            policy_error = ex
+
         for extension, ext_handler in all_extensions:
 
             handler_i = ExtHandlerInstance(ext_handler, self.protocol, extension=extension)
 
+            # Invoke policy engine to determine if extension is allowed. If not, block extension and report error on
+            # behalf of the extension.
+            policy_err_map = {
+                ExtensionRequestedState.Enabled: ('enable', ExtensionErrorCodes.PluginEnableProcessingFailed),
+                # TODO: CRP does not currently have a terminal error code for uninstall. Once CRP adds
+                # an error code for uninstall or for policy, use this code instead of PluginDisableProcessingFailed
+                # Note that currently, CRP waits for 90 minutes to time out for a failed uninstall operation, instead of
+                # failing fast.
+                ExtensionRequestedState.Uninstall: ('uninstall', ExtensionErrorCodes.PluginDisableProcessingFailed),
+                ExtensionRequestedState.Disabled: ('disable', ExtensionErrorCodes.PluginDisableProcessingFailed),
+            }
+            policy_op, policy_err_code = policy_err_map.get(ext_handler.state)
+            if policy_error is not None:
+                err = ExtensionPolicyError(msg="", inner=policy_error, code=policy_err_code)
+                self.__handle_and_report_ext_handler_errors(handler_i, err,
+                                                            report_op=handler_i.operation,
+                                                            message=ustr(err), extension=extension, report=True)
+                continue
+
+            extension_allowed = policy_engine.should_allow_extension(ext_handler.name)
+            if not extension_allowed:
+                msg = "failed to {0} extension '{1}' because extension is not specified in allowlist. To {0}, " \
+                      "add extension to the allowed list in the policy file ('{2}').".format(policy_op,
+                                                                                             ext_handler.name,
+                                                                                             conf.get_policy_file_path())
+                err = ExtensionPolicyError(msg, code=policy_err_code)
+                self.__handle_and_report_ext_handler_errors(handler_i, err,
+                                                            report_op=handler_i.operation,
+                                                            message=ustr(err), extension=extension, report=True)
+
             # In case of extensions disabled, we skip processing extensions. But CRP is still waiting for some status
             # back for the skipped extensions. In order to propagate the status back to CRP, we will report status back
             # here with an error message.
@@ -527,7 +567,11 @@ def handle_ext_handlers(self, goal_state_id):
                 continue
 
             # Process extensions and get if it was successfully executed or not
-            extension_success = self.handle_ext_handler(handler_i, extension, goal_state_id)
+            # If extension was blocked by policy, treat the extension as failed and do not process the handler.
+            if not extension_allowed:
+                extension_success = False
+            else:
+                extension_success = self.handle_ext_handler(handler_i, extension, goal_state_id)
 
             dep_level = self.__get_dependency_level((extension, ext_handler))
             if 0 <= dep_level < max_dep_level:

azurelinuxagent/ga/policy/policy_engine.py

@@ -22,7 +22,7 @@
 from azurelinuxagent.common import logger
 from azurelinuxagent.common.event import WALAEventOperation, add_event
 from azurelinuxagent.common import conf
-from azurelinuxagent.common.exception import AgentError
+from azurelinuxagent.common.exception import AgentError, ExtensionError
 from azurelinuxagent.common.protocol.extensions_goal_state_from_vm_settings import _CaseFoldedDict
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
 
@@ -36,12 +36,6 @@
 _MAX_SUPPORTED_POLICY_VERSION = "0.1.0"
 
 
-class PolicyError(AgentError):
-    """
-    Error raised during agent policy enforcement.
-    """
-
-
 class InvalidPolicyError(AgentError):
     """
     Error raised if user-provided policy is invalid.
@@ -51,6 +45,16 @@ def __init__(self, msg, inner=None):
         super(InvalidPolicyError, self).__init__(msg, inner)
 
 
+class ExtensionPolicyError(ExtensionError):
+    """
+    Error raised during agent extension policy enforcement.
+    """
+    # TODO: when CRP adds terminal error code for policy-related extension failures, set that as the default code.
+    def __init__(self, msg, inner=None, code=-1):
+        msg = "Extension is disallowed by agent policy and will not be processed: {0}".format(msg)
+        super(ExtensionPolicyError, self).__init__(msg, inner, code)
+
+
 class _PolicyEngine(object):
     """
     Implements base policy engine API.

tests/ga/test_extension.py

@@ -427,7 +427,7 @@ def test_migration_ignores_tree_remove_errors(self, shutil_mock):  # pylint: dis
 class TestExtensionBase(AgentTestCase):
     def _assert_handler_status(self, report_vm_status, expected_status,
                                expected_ext_count, version,
-                               expected_handler_name="OSTCExtensions.ExampleHandlerLinux", expected_msg=None):
+                               expected_handler_name="OSTCExtensions.ExampleHandlerLinux", expected_msg=None, expected_code=None):
         self.assertTrue(report_vm_status.called)
         args, kw = report_vm_status.call_args  # pylint: disable=unused-variable
         vm_status = args[0]
@@ -443,6 +443,9 @@ def _assert_handler_status(self, report_vm_status, expected_status,
         if expected_msg is not None:
             self.assertIn(expected_msg, handler_status.message)
 
+        if expected_code is not None:
+            self.assertEqual(expected_code, handler_status.code)
+
 
 # Deprecated. New tests should be added to the TestExtension class
 @patch('time.sleep', side_effect=lambda _: mock_sleep(0.001))
@@ -3507,5 +3510,144 @@ def test_report_msg_if_handler_manifest_contains_invalid_values(self):
                 self.assertIn("'supportsMultipleExtensions' has a non-boolean value", kw_messages[2]['message'])
 
 
+class TestExtensionPolicy(TestExtensionBase):
+    def setUp(self):
+        AgentTestCase.setUp(self)
+        self.policy_path = os.path.join(self.tmp_dir, "waagent_policy.json")
+
+        # Patch attributes to enable policy feature
+        self.patch_policy_path = patch('azurelinuxagent.common.conf.get_policy_file_path',
+                                       return_value=str(self.policy_path))
+        self.patch_policy_path.start()
+        self.patch_conf_flag = patch('azurelinuxagent.ga.policy.policy_engine.conf.get_extension_policy_enabled',
+                                     return_value=True)
+        self.patch_conf_flag.start()
+        self.maxDiff = None     # When long error messages don't match, display the entire diff.
+
+    def tearDown(self):
+        patch.stopall()
+        AgentTestCase.tearDown(self)
+
+    def _create_policy_file(self, policy):
+        with open(self.policy_path, mode='w') as policy_file:
+            if isinstance(policy, dict):
+                json.dump(policy, policy_file, indent=4)
+            else:
+                policy_file.write(policy)
+            policy_file.flush()
+
+    def _test_policy_failure(self, policy, op, expected_status_code, expected_handler_status,
+                             expected_status_msg=None):
+
+        with mock_wire_protocol(wire_protocol_data.DATA_FILE) as protocol:
+            if op == ExtensionRequestedState.Uninstall:
+                protocol.mock_wire_data.set_incarnation(2)
+                protocol.mock_wire_data.set_extensions_config_state(ExtensionRequestedState.Uninstall)
+                protocol.client.update_goal_state()
+            protocol.aggregate_status = None
+            protocol.report_vm_status = MagicMock()
+            exthandlers_handler = get_exthandlers_handler(protocol)
+
+            self._create_policy_file(policy)
+            exthandlers_handler.run()
+            exthandlers_handler.report_ext_handlers_status()
+
+            report_vm_status = protocol.report_vm_status
+            self.assertTrue(report_vm_status.called)
+            self._assert_handler_status(report_vm_status, expected_handler_status, 0, "1.0.0", 'OSTCExtensions.ExampleHandlerLinux',
+                                        expected_msg=expected_status_msg, expected_code=expected_status_code)
+
+    def test_should_fail_enable_if_extension_disallowed(self):
+        policy = \
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+               }
+            }
+        expected_msg = "failed to enable extension 'OSTCExtensions.ExampleHandlerLinux' because extension is not specified in allowlist."
+        self._test_policy_failure(policy=policy, op=ExtensionRequestedState.Enabled, expected_status_code=ExtensionErrorCodes.PluginEnableProcessingFailed,
+                          expected_handler_status='NotReady', expected_status_msg=expected_msg)
+
+    def test_should_fail_enable_for_invalid_policy(self):
+        policy = \
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": "False"
+                }
+            }
+        expected_msg = "attribute 'extensionPolicies.allowListedExtensionsOnly'; must be 'boolean'"
+        self._test_policy_failure(policy=policy, op=ExtensionRequestedState.Enabled, expected_status_code=ExtensionErrorCodes.PluginEnableProcessingFailed,
+                          expected_handler_status='NotReady', expected_status_msg=expected_msg)
+
+    def test_should_fail_extension_if_error_thrown_during_policy_engine_init(self):
+        policy = \
+            {
+                "policyVersion": "0.1.0"
+            }
+        with patch('azurelinuxagent.ga.policy.policy_engine.ExtensionPolicyEngine.__init__',
+                   side_effect=Exception("mock exception")):
+            expected_msg = "Extension is disallowed by agent policy and will not be processed: \nInner error: mock exception"
+            self._test_policy_failure(policy=policy, op=ExtensionRequestedState.Enabled,
+                                      expected_status_code=ExtensionErrorCodes.PluginEnableProcessingFailed,
+                                      expected_handler_status='NotReady', expected_status_msg=expected_msg)
+
+    def test_should_fail_uninstall_if_extension_disallowed(self):
+        policy = \
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "signatureRequired": False,
+                    "extensions": {}
+                },
+            }
+        expected_msg = "failed to uninstall extension 'OSTCExtensions.ExampleHandlerLinux' because extension is not specified in allowlist."
+        self._test_policy_failure(policy=policy, op=ExtensionRequestedState.Uninstall, expected_status_code=ExtensionErrorCodes.PluginDisableProcessingFailed,
+                                  expected_handler_status='NotReady', expected_status_msg=expected_msg)
+
+    def test_should_fail_enable_if_dependent_extension_disallowed(self):
+        self._create_policy_file({
+            "policyVersion": "0.1.0",
+            "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "extensions": {
+                        "OSTCExtensions.ExampleHandlerLinux": {}
+                    }
+                }
+            })
+        with mock_wire_protocol(wire_protocol_data.DATA_FILE_EXT_SEQUENCING) as protocol:
+            protocol.aggregate_status = None
+            protocol.report_vm_status = MagicMock()
+            exthandlers_handler = get_exthandlers_handler(protocol)
+            dep_ext_level_2 = extension_emulator(name="OSTCExtensions.ExampleHandlerLinux")
+            dep_ext_level_1 = extension_emulator(name="OSTCExtensions.OtherExampleHandlerLinux")
+
+            exthandlers_handler.run()
+            exthandlers_handler.report_ext_handlers_status()
+
+            # OtherExampleHandlerLinux should be disallowed by policy, ExampleHandlerLinux should be skipped because
+            # dependent extension failed
+            self._assert_handler_status(protocol.report_vm_status, "NotReady", 0, "1.0.0",
+                                        expected_handler_name="OSTCExtensions.OtherExampleHandlerLinux",
+                                        expected_msg=("failed to enable extension 'OSTCExtensions.OtherExampleHandlerLinux' "
+                                                      "because extension is not specified in allowlist."))
+
+            self._assert_handler_status(protocol.report_vm_status, "NotReady", 0, "1.0.0",
+                                        expected_handler_name="OSTCExtensions.ExampleHandlerLinux",
+                                        expected_msg="Skipping processing of extensions since execution of dependent "
+                                                     "extension OSTCExtensions.OtherExampleHandlerLinux failed")
+
+            # check handler list and dependency levels
+            self.assertTrue(exthandlers_handler.ext_handlers is not None)
+            self.assertTrue(exthandlers_handler.ext_handlers is not None)
+            self.assertEqual(len(exthandlers_handler.ext_handlers), 2)
+            self.assertEqual(1, next(handler for handler in exthandlers_handler.ext_handlers if
+                                     handler.name == dep_ext_level_1.name).settings[0].dependencyLevel)
+            self.assertEqual(2, next(handler for handler in exthandlers_handler.ext_handlers if
+                                     handler.name == dep_ext_level_2.name).settings[0].dependencyLevel)
+
+
 if __name__ == '__main__':
     unittest.main()

tests/ga/test_multi_config_extension.py

@@ -630,6 +630,70 @@ def test_it_should_handle_and_report_enable_errors_properly(self):
                 }
                 self._assert_extension_status(sc_handler, expected_extensions)
 
+    def test_it_should_handle_and_report_disallowed_extensions_properly(self):
+        """If multiconfig extension is disallowed by policy, all instances should be blocked."""
+        policy_path = os.path.join(self.tmp_dir, "waagent_policy.json")
+        patch('azurelinuxagent.common.conf.get_policy_file_path', return_value=str(policy_path)).start()
+        patch('azurelinuxagent.ga.policy.policy_engine.conf.get_extension_policy_enabled',
+              return_value=True).start()
+        policy = \
+            {
+                "policyVersion": "0.0.1",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "signatureRequired": True,
+                    "extensions": {
+                        "Microsoft.Powershell.ExampleExtension": {}
+                    }
+                }
+            }
+        with open(policy_path, mode='w') as policy_file:
+            json.dump(policy, policy_file, indent=4)
+            policy_file.flush()
+        self.test_data['ext_conf'] = os.path.join(self._MULTI_CONFIG_TEST_DATA,
+                                                  "ext_conf_multi_config_no_dependencies.xml")
+        with self._setup_test_env() as (exthandlers_handler, protocol, no_of_extensions):
+            disallowed_mc_1 = extension_emulator(name="OSTCExtensions.ExampleHandlerLinux.firstExtension",
+                                                 supports_multiple_extensions=True)
+            disallowed_mc_2 = extension_emulator(name="OSTCExtensions.ExampleHandlerLinux.secondExtension",
+                                                 supports_multiple_extensions=True)
+            disallowed_mc_3 = extension_emulator(name="OSTCExtensions.ExampleHandlerLinux.thirdExtension",
+                                                 supports_multiple_extensions=True)
+            allowed_ext = extension_emulator(name="Microsoft.Powershell.ExampleExtension")
+            with enable_invocations(disallowed_mc_1, disallowed_mc_2, disallowed_mc_3,
+                                    allowed_ext) as invocation_record:
+                exthandlers_handler.run()
+                exthandlers_handler.report_ext_handlers_status()
+                self.assertEqual(no_of_extensions,
+                                 len(protocol.aggregate_status['aggregateStatus']['handlerAggregateStatus']),
+                                 "incorrect extensions reported")
+
+                # We should only enable the allowed extension, no instances of the multiconfig extension should be enabled
+                invocation_record.compare(
+                    (allowed_ext, ExtensionCommandNames.INSTALL),
+                    (allowed_ext, ExtensionCommandNames.ENABLE)
+                )
+
+                mc_handlers = self._assert_and_get_handler_status(aggregate_status=protocol.aggregate_status,
+                                                                  handler_name="OSTCExtensions.ExampleHandlerLinux",
+                                                                  expected_count=3, status="NotReady")
+                msg = "failed to enable extension 'OSTCExtensions.ExampleHandlerLinux' because extension is not specified in allowlist"
+                expected_extensions = {
+                    "firstExtension": {"status": ExtensionStatusValue.error, "seq_no": 1, "message": msg},
+                    "secondExtension": {"status": ExtensionStatusValue.error, "seq_no": 2, "message": msg},
+                    "thirdExtension": {"status": ExtensionStatusValue.error, "seq_no": 3, "message": msg},
+                }
+                self._assert_extension_status(mc_handlers, expected_extensions, multi_config=True)
+
+                sc_handler = self._assert_and_get_handler_status(aggregate_status=protocol.aggregate_status,
+                                                                 handler_name="Microsoft.Powershell.ExampleExtension",
+                                                                 status="Ready", message=None)
+                expected_extensions = {
+                    "Microsoft.Powershell.ExampleExtension": {"status": ExtensionStatusValue.success, "seq_no": 9,
+                                                              "message": None}
+                }
+                self._assert_extension_status(sc_handler, expected_extensions)
+
     def test_it_should_cleanup_extension_state_on_disable(self):
 
         def __assert_state_file(handler_name, handler_version, extensions, state, not_present=None):

tests_e2e/orchestrator/runbook.yml

@@ -55,6 +55,8 @@ variable:
       - recover_network_interface
       - cgroup_v2_disabled
       - log_collector
+      - ext_policy
+      - ext_policy_with_dependencies
 
   #
   # Additional arguments pass to the test suites