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,47 @@ 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_policy_error(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_policy_error(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 +565,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:
@@ -692,6 +734,26 @@ def __handle_and_report_ext_handler_errors(ext_handler_i, error, report_op, mess
             add_event(name=name, version=handler_version, op=report_op, is_success=False, log_event=True,
                       message=message)
 
+    @staticmethod
+    def __handle_and_report_policy_error(ext_handler_i, error, report_op, message, report=True, extension=None):
+        # TODO: Consider merging this function with __handle_and_report_ext_handler_errors() above.
+
+        # Set handler status for all extensions
+        ext_handler_i.set_handler_status(message=message, code=error.code)
+
+        # Create status file for only extensions with settings. Since extensions are not processed in the case of
+        # policy-related failures, no extension status file is created. For CRP to report status, we need to create the
+        # file with failure on behalf of the extension. This should be done for both multi-config and single-config extensions.
+        if extension is not None:
+            ext_handler_i.create_status_file_if_not_exist(extension, status=ExtensionStatusValue.error, code=error.code,
+                                                          operation=report_op, message=message)
+
+        if report:
+            name = ext_handler_i.get_extension_full_name(extension)
+            handler_version = ext_handler_i.ext_handler.version
+            add_event(name=name, version=handler_version, op=report_op, is_success=False, log_event=True,
+                      message=message)
+
     def handle_enable(self, ext_handler_i, extension):
         """
              1- Ensure the handler is installed

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,15 @@ def __init__(self, msg, inner=None):
         super(InvalidPolicyError, self).__init__(msg, inner)
 
 
+class ExtensionPolicyError(ExtensionError):
+    """
+    Error raised during agent extension policy enforcement.
+    """
+    def __init__(self, msg, code, inner=None):
+        msg = "Extension 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 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()