Download certificates when goal state source is Fast Track

azurelinuxagent/common/protocol/goal_state.py

@@ -56,8 +56,9 @@ class GoalStateProperties(object):
     HostingEnv = 0x2
     SharedConfig = 0x4
     ExtensionsGoalState = 0x8
-    RemoteAccessInfo = 0x10
-    All = RoleConfig | HostingEnv | SharedConfig | ExtensionsGoalState | RemoteAccessInfo
+    Certificates = 0x10
+    RemoteAccessInfo = 0x20
+    All = RoleConfig | HostingEnv | SharedConfig | ExtensionsGoalState | Certificates | RemoteAccessInfo
 
 
 class GoalStateInconsistentError(ProtocolError):
@@ -140,7 +141,7 @@ def extensions_goal_state(self):
 
     @property
     def certs(self):
-        if not self._goal_state_properties & GoalStateProperties.ExtensionsGoalState:
+        if not self._goal_state_properties & GoalStateProperties.Certificates:
             raise ProtocolError("Certificates is not in goal state properties")
         else:
             return self._certs
@@ -289,6 +290,9 @@ def _update(self, force_update):
         # case, to ensure it fetches the new certificate.
         #
         if self._extensions_goal_state.source == GoalStateSource.FastTrack:
+            certs_uri = findtext(xml_doc, "Certificates")
+            if certs_uri is not None:
+                self._download_certificates(certs_uri)
             self._check_certificates()
 
     def _check_certificates(self):
@@ -301,6 +305,20 @@ def _check_certificates(self):
                     message = "Certificate {0} needed by {1} is missing from the goal state".format(settings.certificateThumbprint, extension.name)
                     raise GoalStateInconsistentError(message)
 
+    def _download_certificates(self, certs_uri):
+        xml_text = self._wire_client.fetch_config(certs_uri, self._wire_client.get_header_for_cert())
+        certs = Certificates(xml_text, self.logger)
+        # Log and save the certificates summary (i.e. the thumbprint but not the certificate itself) to the goal state history
+        for c in certs.summary:
+            message = "Downloaded certificate {0}".format(c)
+            self.logger.info(message)
+            add_event(op=WALAEventOperation.GoalState, message=message)
+        if len(certs.warnings) > 0:
+            self.logger.warn(certs.warnings)
+            add_event(op=WALAEventOperation.GoalState, message=certs.warnings)
+        self._history.save_certificates(json.dumps(certs.summary))
+        return certs
+
     def _restore_wire_server_goal_state(self, incarnation, xml_text, xml_doc, vm_settings_support_stopped_error):
         msg = 'The HGAP stopped supporting vmSettings; will fetched the goal state from the WireServer.'
         self.logger.info(msg)
@@ -435,18 +453,8 @@ def _fetch_full_wire_server_goal_state(self, incarnation, xml_doc):
 
             certs = EmptyCertificates()
             certs_uri = findtext(xml_doc, "Certificates")
-            if (GoalStateProperties.ExtensionsGoalState & self._goal_state_properties) and certs_uri is not None:
-                xml_text = self._wire_client.fetch_config(certs_uri, self._wire_client.get_header_for_cert())
-                certs = Certificates(xml_text, self.logger)
-                # Log and save the certificates summary (i.e. the thumbprint but not the certificate itself) to the goal state history
-                for c in certs.summary:
-                    message = "Downloaded certificate {0}".format(c)
-                    self.logger.info(message)
-                    add_event(op=WALAEventOperation.GoalState, message=message)
-                if len(certs.warnings) > 0:
-                    self.logger.warn(certs.warnings)
-                    add_event(op=WALAEventOperation.GoalState, message=certs.warnings)
-                self._history.save_certificates(json.dumps(certs.summary))
+            if (GoalStateProperties.Certificates & self._goal_state_properties) and certs_uri is not None:
+                certs = self._download_certificates(certs_uri)
 
             remote_access = None
             if GoalStateProperties.RemoteAccessInfo & self._goal_state_properties:

tests/protocol/mockwiredata.py

@@ -165,6 +165,7 @@ def __init__(self, data_files=None):
         self.in_vm_artifacts_profile = None
         self.vm_settings = None
         self.etag = None
+        self.prev_etag = None
         self.imds_info = None
 
         self.reload()
@@ -242,9 +243,12 @@ def mock_http_get(self, url, *_, **kwargs):
         elif "/vmSettings" in url:
             if self.vm_settings is None:
                 resp.status = httpclient.NOT_FOUND
+            elif self.call_counts["vm_settings"] > 0 and self.prev_etag == self.etag:
+                resp.status = httpclient.NOT_MODIFIED
             else:
                 content = self.vm_settings
                 response_headers = [('ETag', self.etag)]
+                self.prev_etag = self.etag
             self.call_counts["vm_settings"] += 1
         elif '{0}/metadata/compute'.format(IMDS_ENDPOINT) in url:
             content = json.dumps(self.imds_info.get("compute", "{}"))

tests/protocol/test_extensions_goal_state_from_vm_settings.py

@@ -58,6 +58,7 @@ def test_it_should_parse_requested_version_properly(self):
         data_file = mockwiredata.DATA_FILE_VM_SETTINGS.copy()
         data_file["vm_settings"] = "hostgaplugin/vm_settings-requested_version.json"
         with mock_wire_protocol(data_file) as protocol:
+            protocol.mock_wire_data.set_etag(888)
             goal_state = GoalState(protocol.client)
             families = goal_state.extensions_goal_state.agent_families
             for family in families:

tests/protocol/test_goal_state.py

@@ -28,6 +28,7 @@
 class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
     def test_it_should_use_vm_settings_by_default(self):
         with mock_wire_protocol(mockwiredata.DATA_FILE_VM_SETTINGS) as protocol:
+            protocol.mock_wire_data.set_etag(888)
             extensions_goal_state = GoalState(protocol.client).extensions_goal_state
             self.assertTrue(
                 isinstance(extensions_goal_state, ExtensionsGoalStateFromVmSettings),
@@ -156,11 +157,12 @@ def http_get_handler(url, *_, **__):
             protocol.set_http_handlers(http_get_handler=None)
             goal_state.update()
             self._assert_directory_contents(
-                self._find_history_subdirectory("234-987"), ["VmSettings.json"])
+                self._find_history_subdirectory("234-987"), ["VmSettings.json", "Certificates.json"])
 
     def test_it_should_redact_the_protected_settings_when_saving_to_the_history_directory(self):
         with mock_wire_protocol(mockwiredata.DATA_FILE_VM_SETTINGS) as protocol:
             protocol.mock_wire_data.set_incarnation(888)
+            protocol.mock_wire_data.set_etag(888)
 
             goal_state = GoalState(protocol.client)
 
@@ -173,7 +175,7 @@ def test_it_should_redact_the_protected_settings_when_saving_to_the_history_dire
             if len(protected_settings) == 0:
                 raise Exception("The test goal state does not include any protected settings")
 
-            history_directory = self._find_history_subdirectory("888-1")
+            history_directory = self._find_history_subdirectory("888-888")
             extensions_config_file = os.path.join(history_directory, "ExtensionsConfig.xml")
             vm_settings_file = os.path.join(history_directory, "VmSettings.json")
             for file_name in extensions_config_file, vm_settings_file:
@@ -198,14 +200,14 @@ def test_it_should_save_vm_settings_on_parse_errors(self):
             data_file = mockwiredata.DATA_FILE_VM_SETTINGS.copy()
             data_file["vm_settings"] = invalid_vm_settings_file
             protocol.mock_wire_data = mockwiredata.WireProtocolData(data_file)
-            protocol.mock_wire_data.set_etag(888)
 
             with self.assertRaises(ProtocolError):  # the parsing error will cause an exception
                 _ = GoalState(protocol.client)
 
             # Do an extra call to update the goal state; this should save the vmsettings to the history directory
             # only once (self._find_history_subdirectory asserts 1 single match)
             time.sleep(0.1)  # add a short delay to ensure that a new timestamp would be saved in the history folder
+            protocol.mock_wire_data.set_etag(888)
             with self.assertRaises(ProtocolError):
                 _ = GoalState(protocol.client)
 
@@ -375,13 +377,63 @@ def test_it_should_raise_when_the_tenant_certificate_is_missing(self):
         with mock_wire_protocol(data_file) as protocol:
             data_file["vm_settings"] = "hostgaplugin/vm_settings-missing_cert.json"
             protocol.mock_wire_data.reload()
+            protocol.mock_wire_data.set_etag(888)
 
             with self.assertRaises(GoalStateInconsistentError) as context:
                 _ = GoalState(protocol.client)
 
             expected_message = "Certificate 59A10F50FFE2A0408D3F03FE336C8FD5716CF25C needed by Microsoft.OSTCExtensions.VMAccessForLinux is missing from the goal state"
             self.assertIn(expected_message, str(context.exception))
 
+    def test_it_should_download_certs_on_a_new_fast_track_goal_state(self):
+        data_file = mockwiredata.DATA_FILE_VM_SETTINGS.copy()
+
+        with mock_wire_protocol(data_file) as protocol:
+            goal_state = GoalState(protocol.client)
+
+            cert = "BD447EF71C3ADDF7C837E84D630F3FAC22CCD22F"
+            crt_path = os.path.join(self.tmp_dir, cert + ".crt")
+            prv_path = os.path.join(self.tmp_dir, cert + ".prv")
+
+            # Check that crt and prv files are downloaded after processing goal state
+            self.assertTrue(os.path.isfile(crt_path))
+            self.assertTrue(os.path.isfile(prv_path))
+
+            # Remove .crt file
+            os.remove(crt_path)
+            if os.path.isfile(crt_path):
+                raise Exception("{0}.crt was not removed.".format(cert))
+
+            # Update goal state and check that .crt was downloaded
+            protocol.mock_wire_data.set_etag(888)
+            goal_state.update()
+            self.assertTrue(os.path.isfile(crt_path))
+
+    def test_it_should_download_certs_on_a_new_fabric_goal_state(self):
+        data_file = mockwiredata.DATA_FILE_VM_SETTINGS.copy()
+
+        with mock_wire_protocol(data_file) as protocol:
+            protocol.mock_wire_data.set_vm_settings_source(GoalStateSource.Fabric)
+            goal_state = GoalState(protocol.client)
+
+            cert = "BD447EF71C3ADDF7C837E84D630F3FAC22CCD22F"
+            crt_path = os.path.join(self.tmp_dir, cert + ".crt")
+            prv_path = os.path.join(self.tmp_dir, cert + ".prv")
+
+            # Check that crt and prv files are downloaded after processing goal state
+            self.assertTrue(os.path.isfile(crt_path))
+            self.assertTrue(os.path.isfile(prv_path))
+
+            # Remove .crt file
+            os.remove(crt_path)
+            if os.path.isfile(crt_path):
+                raise Exception("{0}.crt was not removed.".format(cert))
+
+            # Update goal state and check that .crt was downloaded
+            protocol.mock_wire_data.set_incarnation(999)
+            goal_state.update()
+            self.assertTrue(os.path.isfile(crt_path))
+
     def test_it_should_refresh_the_goal_state_when_it_is_inconsistent(self):
         #
         # Some scenarios can produce inconsistent goal states. For example, during hibernation/resume, the Fabric goal state changes (the
@@ -412,7 +464,7 @@ def http_get_handler(url, *_, **__):
             goal_state = GoalState(protocol.client)
 
             self.assertEqual(2, protocol.mock_wire_data.call_counts['goalstate'], "There should have been exactly 2 requests for the goal state (original + refresh)")
-            self.assertEqual(2, http_get_handler.certificate_requests, "There should have been exactly 2 requests for the goal state certificates (original + refresh)")
+            self.assertEqual(4, http_get_handler.certificate_requests, "There should have been exactly 4 requests for the goal state certificates (2x original + 2x refresh)")
 
             thumbprints = [c.thumbprint for c in goal_state.certs.cert_list.certificates]
 

tests/protocol/test_hostplugin.py

@@ -998,7 +998,7 @@ def test_it_should_save_the_timestamp_of_the_most_recent_fast_track_goal_state(s
 
             # A fabric goal state should remove the state file
             protocol.mock_wire_data.set_vm_settings_source(GoalStateSource.Fabric)
-
+            protocol.mock_wire_data.set_etag(888)
             _ = host_ga_plugin.fetch_vm_settings()
 
             self.assertFalse(os.path.exists(state_file), "{0} was not removed by a Fabric goal state".format(state_file))

tests/protocol/test_wire.py

@@ -1101,7 +1101,7 @@ def test_forced_update_should_update_the_goal_state_and_the_host_plugin_when_the
 
     def test_reset_should_init_provided_goal_state_properties(self):
         with mock_wire_protocol(mockwiredata.DATA_FILE) as protocol:
-            protocol.client.reset_goal_state(goal_state_properties=GoalStateProperties.All & ~GoalStateProperties.ExtensionsGoalState)
+            protocol.client.reset_goal_state(goal_state_properties=GoalStateProperties.All & ~GoalStateProperties.Certificates)
 
             with self.assertRaises(ProtocolError) as context:
                 _ = protocol.client.get_certs()