Repository: Azure/WALinuxAgent (python)

Issue:
WALinuxAgent doesn't download all .crt and .prv files from KeyVault
Not sure if this is a bug but I will try to explain as much as I can.

Environment: Virtual Machine Scale Set - Azure (West EU, West US, Japan, Asia,.. all regions impacted)

We are doing deployments of VMSS in azure using ARM template and ansible playbook with some configurations.
Before running the ansible, we are using the following to push certficates from KV:

TEMPLATE
{
    "type": "Microsoft.Compute/virtualMachines",
    "name": "Region1VM",
    ...
    "properties": {
    ...
    "osProfile": {
            "computerName": "Region1VM",
    ...
    "secrets": [
                {
                    "sourceVault": {
                        "id": "[resourceId('Microsoft.KeyVault/vaults', Region1KeyVault)]"
                    },
                    "vaultCertificates": [
                        {
                            "certificateUrl": "[reference(resourceId('Microsoft.KeyVault/vaults/secrets', 'Region1KeyVault', 'SampleCertificateAsSecret')).secretUriWithVersion]",
                            "certificateStore": "My"
                        }
                    ]
                }
            ],
        },
    ...
    }
}

CODE
![image](https://user-images.githubusercontent.com/124268307/216367634-fdc263b6-0e5f-4cd3-9a75-6233dc72fc37.png)

LINK
https://devblogs.microsoft.com/premier-developer/centralized-vm-certificate-deployment-across-multiple-regions-with-arm-templates/#part-2-push-certificate-from-the-regional-key-vault-to-the-virtual-machine

after running this part of the ARM template we do some certificate copies from var/lib/waagent to another location but it fails with the below error:

"could not find or access '/var/lib/waagent/nameexample.prv"

![image](https://user-images.githubusercontent.com/124268307/216368423-265403c7-837f-4dcb-8b9c-1a34bfe0fe91.png)

The problem is, the file that is missing should be downloaded during the push of the certificates from Keyvault but this is not happening and ansible playbook crashes.

If we restart the waagent service, the file "nameexample.prv" will be downloaded and ansible will not crash anymore.

The final lines of the ansible code, will remove this file again from the VM. The next deployment will crash again. 

We have two workarounds here:
FIRST- if we restart the waagent after the crash everything will run as expected
SECOND - if we don't delete the file after the ansible

MAIN PROBLEM - this has been working like this for last 8 months, but now we are getting this errors. 

we don't understand why, the agent doesn't push all the file in KV and we always need to restart the service to do a "complete download" let's say.


 - Distro and Version: RedHat|RHEL|7.3|
 - WALinuxAgent version: 2.9.0.4

**Additional context**
I can give you an agent log from 31st of January 4:30PM issue : 

<img width="1190" alt="image" src="https://user-images.githubusercontent.com/124268307/216373059-724d0fce-704d-408d-a166-db0100be89e2.png">

**Log file attached**

[waagent.log](https://github.com/Azure/WALinuxAgent/files/10570703/waagent.log)


Pull request description:
Download certificates when goal state source is Fast Track
This PR includes a change to download certificates when goal state source is FastTrack

Fixes #2750 

Patch:
diff --git a/azurelinuxagent/common/protocol/goal_state.py b/azurelinuxagent/common/protocol/goal_state.py
index ed96159c..6b2a0c2c 100644
--- a/azurelinuxagent/common/protocol/goal_state.py
+++ b/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):
@@ -96,6 +97,7 @@ class GoalState(object):
             self._hosting_env = None
             self._shared_conf = None
             self._certs = EmptyCertificates()
+            self._certs_uri = None
             self._remote_access = None
 
             self.update(silent=silent)
@@ -140,7 +142,7 @@ class GoalState(object):
 
     @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
@@ -292,6 +294,10 @@ class GoalState(object):
             self._check_certificates()
 
     def _check_certificates(self):
+        # Re-download certificates in case they have been removed from disk since last download
+        if self._goal_state_properties & GoalStateProperties.Certificates and self._certs_uri is not None:
+            self._download_certificates(self._certs_uri)
+        # Check that certificates needed by extensions are in goal state certs.summary
         for extension in self.extensions_goal_state.extensions:
             for settings in extension.settings:
                 if settings.protectedSettings is None:
@@ -301,6 +307,20 @@ class GoalState(object):
                     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 +455,8 @@ class GoalState(object):
 
             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:
@@ -463,6 +473,7 @@ class GoalState(object):
             self._hosting_env = hosting_env
             self._shared_conf = shared_config
             self._certs = certs
+            self._certs_uri = certs_uri
             self._remote_access = remote_access
 
             return extensions_config


Test patch:
diff --git a/tests/protocol/mockwiredata.py b/tests/protocol/mockwiredata.py
index 7ec311af..196ed32d 100644
--- a/tests/protocol/mockwiredata.py
+++ b/tests/protocol/mockwiredata.py
@@ -165,6 +165,7 @@ class WireProtocolData(object):
         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 @@ class WireProtocolData(object):
         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", "{}"))
diff --git a/tests/protocol/test_extensions_goal_state_from_vm_settings.py b/tests/protocol/test_extensions_goal_state_from_vm_settings.py
index fb97a075..1100b05b 100644
--- a/tests/protocol/test_extensions_goal_state_from_vm_settings.py
+++ b/tests/protocol/test_extensions_goal_state_from_vm_settings.py
@@ -58,6 +58,7 @@ class ExtensionsGoalStateFromVmSettingsTestCase(AgentTestCase):
         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:
diff --git a/tests/protocol/test_goal_state.py b/tests/protocol/test_goal_state.py
index 869da68c..61653b2a 100644
--- a/tests/protocol/test_goal_state.py
+++ b/tests/protocol/test_goal_state.py
@@ -28,6 +28,7 @@ from tests.tools import AgentTestCase, patch, load_data
 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 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
             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 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
             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,7 +200,6 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
             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)
@@ -206,6 +207,7 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
             # 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,6 +377,7 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
         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)
@@ -382,6 +385,55 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
             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 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
             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]
 
diff --git a/tests/protocol/test_hostplugin.py b/tests/protocol/test_hostplugin.py
index b85ed757..47e6871b 100644
--- a/tests/protocol/test_hostplugin.py
+++ b/tests/protocol/test_hostplugin.py
@@ -998,7 +998,7 @@ class TestHostPluginVmSettings(HttpRequestPredicates, AgentTestCase):
 
             # 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))
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index bbf018fc..2a36fc29 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -1101,7 +1101,7 @@ class UpdateGoalStateTestCase(HttpRequestPredicates, AgentTestCase):
 
     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()