azure: pps imds (#1093)

Without UDF support, DS Azure cannot mount the provisioning ISO,
which contains platform metadata necessary to support
pre-provisioning. The required metadata is made available in IMDS 
starting with api version 2021-08-01. This change will leverage IMDS
to obtain the required metadata to support pre-preprovisioning if 
provisioning ISO was not available.

2 files changed, 128 insertions, 29 deletions

diff --git a/cloudinit/sources/DataSourceAzure.py b/cloudinit/sources/DataSourceAzure.py
index 26c407ef..d52c6a7f 100755
--- a/cloudinit/sources/DataSourceAzure.py
+++ b/cloudinit/sources/DataSourceAzure.py
@@ -84,8 +84,8 @@ DEFAULT_PROVISIONING_ISO_DEV = '/dev/sr0'
 IMDS_TIMEOUT_IN_SECONDS = 2
 IMDS_URL = "http://169.254.169.254/metadata"
 IMDS_VER_MIN = "2019-06-01"
-IMDS_VER_WANT = "2021-01-01"
-
+IMDS_VER_WANT = "2021-08-01"
+IMDS_EXTENDED_VER_MIN = "2021-03-01"
 
 # This holds SSH key data including if the source was
 # from IMDS, as well as the SSH key data itself.
@@ -93,7 +93,7 @@ SSHKeys = namedtuple("SSHKeys", ("keys_from_imds", "ssh_keys"))
 
 
 class metadata_type(Enum):
-    compute = "{}/instance".format(IMDS_URL)
+    all = "{}/instance".format(IMDS_URL)
     network = "{}/instance/network".format(IMDS_URL)
     reprovisiondata = "{}/reprovisiondata".format(IMDS_URL)
 
@@ -494,10 +494,21 @@ class DataSourceAzure(sources.DataSource):
             "Found provisioning metadata in %s" % metadata_source,
             logger_func=LOG.debug)
 
-        perform_reprovision = reprovision or self._should_reprovision(ret)
+        imds_md = self.get_imds_data_with_api_fallback(
+            self.fallback_interface,
+            retries=10
+        )
+        if not imds_md and not ovf_is_accessible:
+            msg = 'No OVF or IMDS available'
+            report_diagnostic_event(msg)
+            raise sources.InvalidMetaDataException(msg)
+
+        perform_reprovision = (
+            reprovision or
+            self._should_reprovision(ret, imds_md))
         perform_reprovision_after_nic_attach = (
             reprovision_after_nic_attach or
-            self._should_reprovision_after_nic_attach(ret))
+            self._should_reprovision_after_nic_attach(ret, imds_md))
 
         if perform_reprovision or perform_reprovision_after_nic_attach:
             if util.is_FreeBSD():
@@ -507,15 +518,12 @@ class DataSourceAzure(sources.DataSource):
             if perform_reprovision_after_nic_attach:
                 self._wait_for_all_nics_ready()
             ret = self._reprovision()
+            # fetch metadata again as it has changed after reprovisioning
+            imds_md = self.get_imds_data_with_api_fallback(
+                self.fallback_interface,
+                retries=10
+            )
 
-        imds_md = self.get_imds_data_with_api_fallback(
-            self.fallback_interface,
-            retries=10
-        )
-        if not imds_md and not ovf_is_accessible:
-            msg = 'No OVF or IMDS available'
-            report_diagnostic_event(msg)
-            raise sources.InvalidMetaDataException(msg)
         (md, userdata_raw, cfg, files) = ret
         self.seed = metadata_source
         crawled_data.update({
@@ -691,7 +699,7 @@ class DataSourceAzure(sources.DataSource):
             self,
             fallback_nic,
             retries,
-            md_type=metadata_type.compute,
+            md_type=metadata_type.all,
             exc_cb=retry_on_url_exc,
             infinite=False):
         """
@@ -1407,7 +1415,17 @@ class DataSourceAzure(sources.DataSource):
                 "connectivity issues: %s" % e, logger_func=LOG.warning)
             return False
 
-    def _should_reprovision_after_nic_attach(self, candidate_metadata) -> bool:
+    def _ppstype_from_imds(self, imds_md: dict = None) -> str:
+        try:
+            return imds_md['extended']['compute']['ppsType']
+        except Exception as e:
+            report_diagnostic_event(
+                "Could not retrieve pps configuration from IMDS: %s" %
+                e, logger_func=LOG.debug)
+            return None
+
+    def _should_reprovision_after_nic_attach(
+            self, ovf_md, imds_md=None) -> bool:
         """Whether or not we should wait for nic attach and then poll
         IMDS for reprovisioning data. Also sets a marker file to poll IMDS.
 
@@ -1419,14 +1437,16 @@ class DataSourceAzure(sources.DataSource):
         the ISO, thus cloud-init needs to have a way of knowing that it should
         jump back into the waiting mode in order to retrieve the ovf_env.
 
-        @param candidate_metadata: Metadata obtained from reading ovf-env.
+        @param ovf_md: Metadata obtained from reading ovf-env.
+        @param imds_md: Metadata obtained from IMDS
         @return: Whether to reprovision after waiting for nics to be attached.
         """
-        if not candidate_metadata:
+        if not ovf_md:
             return False
-        (_md, _userdata_raw, cfg, _files) = candidate_metadata
+        (_md, _userdata_raw, cfg, _files) = ovf_md
         path = REPROVISION_NIC_ATTACH_MARKER_FILE
         if (cfg.get('PreprovisionedVMType', None) == "Savable" or
+            self._ppstype_from_imds(imds_md) == "Savable" or
                 os.path.isfile(path)):
             if not os.path.isfile(path):
                 LOG.info("Creating a marker file to wait for nic attach: %s",
@@ -1436,7 +1456,7 @@ class DataSourceAzure(sources.DataSource):
             return True
         return False
 
-    def _should_reprovision(self, ret):
+    def _should_reprovision(self, ovf_md, imds_md=None):
         """Whether or not we should poll IMDS for reprovisioning data.
         Also sets a marker file to poll IMDS.
 
@@ -1447,12 +1467,13 @@ class DataSourceAzure(sources.DataSource):
         However, since the VM reports ready to the Fabric, we will not attach
         the ISO, thus cloud-init needs to have a way of knowing that it should
         jump back into the polling loop in order to retrieve the ovf_env."""
-        if not ret:
+        if not ovf_md:
             return False
-        (_md, _userdata_raw, cfg, _files) = ret
+        (_md, _userdata_raw, cfg, _files) = ovf_md
         path = REPROVISION_MARKER_FILE
         if (cfg.get('PreprovisionedVm') is True or
-                cfg.get('PreprovisionedVMType', None) == 'Running' or
+            cfg.get('PreprovisionedVMType', None) == 'Running' or
+            self._ppstype_from_imds(imds_md) == "Running" or
                 os.path.isfile(path)):
             if not os.path.isfile(path):
                 LOG.info("Creating a marker file to poll imds: %s",
@@ -2239,7 +2260,7 @@ def _generate_network_config_from_fallback_config() -> dict:
 @azure_ds_telemetry_reporter
 def get_metadata_from_imds(fallback_nic,
                            retries,
-                           md_type=metadata_type.compute,
+                           md_type=metadata_type.all,
                            api_version=IMDS_VER_MIN,
                            exc_cb=retry_on_url_exc,
                            infinite=False):
@@ -2280,11 +2301,16 @@ def get_metadata_from_imds(fallback_nic,
 def _get_metadata_from_imds(
         retries,
         exc_cb,
-        md_type=metadata_type.compute,
+        md_type=metadata_type.all,
         api_version=IMDS_VER_MIN,
         infinite=False):
     url = "{}?api-version={}".format(md_type.value, api_version)
     headers = {"Metadata": "true"}
+
+    # support for extended metadata begins with 2021-03-01
+    if api_version >= IMDS_EXTENDED_VER_MIN and md_type == metadata_type.all:
+        url = url + "&extended=true"
+
     try:
         response = readurl(
             url, timeout=IMDS_TIMEOUT_IN_SECONDS, headers=headers,
diff --git a/tests/unittests/test_datasource/test_azure.py b/tests/unittests/test_datasource/test_azure.py
index d7206c72..cbc9665d 100644
--- a/tests/unittests/test_datasource/test_azure.py
+++ b/tests/unittests/test_datasource/test_azure.py
@@ -434,16 +434,16 @@ class TestGetMetadataFromIMDS(HttprettyTestCase):
     @mock.patch(MOCKPATH + 'readurl', autospec=True)
     @mock.patch(MOCKPATH + 'EphemeralDHCPv4')
     @mock.patch(MOCKPATH + 'net.is_up')
-    def test_get_compute_metadata_uses_compute_url(
+    def test_get_metadata_uses_instance_url(
             self, m_net_is_up, m_dhcp, m_readurl):
         """Make sure readurl is called with the correct url when accessing
-        network metadata"""
+        metadata"""
         m_net_is_up.return_value = True
         m_readurl.return_value = url_helper.StringResponse(
             json.dumps(IMDS_NETWORK_METADATA).encode('utf-8'))
 
         dsaz.get_metadata_from_imds(
-            'eth0', retries=3, md_type=dsaz.metadata_type.compute)
+            'eth0', retries=3, md_type=dsaz.metadata_type.all)
         m_readurl.assert_called_with(
             "http://169.254.169.254/metadata/instance?api-version="
             "2019-06-01", exception_cb=mock.ANY,
@@ -472,10 +472,10 @@ class TestGetMetadataFromIMDS(HttprettyTestCase):
     @mock.patch(MOCKPATH + 'readurl', autospec=True)
     @mock.patch(MOCKPATH + 'EphemeralDHCPv4')
     @mock.patch(MOCKPATH + 'net.is_up')
-    def test_get_default_metadata_uses_compute_url(
+    def test_get_default_metadata_uses_instance_url(
             self, m_net_is_up, m_dhcp, m_readurl):
         """Make sure readurl is called with the correct url when accessing
-        network metadata"""
+        metadata"""
         m_net_is_up.return_value = True
         m_readurl.return_value = url_helper.StringResponse(
             json.dumps(IMDS_NETWORK_METADATA).encode('utf-8'))
@@ -489,6 +489,26 @@ class TestGetMetadataFromIMDS(HttprettyTestCase):
             timeout=mock.ANY, infinite=False)
 
     @mock.patch(MOCKPATH + 'readurl', autospec=True)
+    @mock.patch(MOCKPATH + 'EphemeralDHCPv4')
+    @mock.patch(MOCKPATH + 'net.is_up')
+    def test_get_metadata_uses_extended_url(
+            self, m_net_is_up, m_dhcp, m_readurl):
+        """Make sure readurl is called with the correct url when accessing
+        metadata"""
+        m_net_is_up.return_value = True
+        m_readurl.return_value = url_helper.StringResponse(
+            json.dumps(IMDS_NETWORK_METADATA).encode('utf-8'))
+
+        dsaz.get_metadata_from_imds(
+            'eth0', retries=3, md_type=dsaz.metadata_type.all,
+            api_version="2021-08-01")
+        m_readurl.assert_called_with(
+            "http://169.254.169.254/metadata/instance?api-version="
+            "2021-08-01&extended=true", exception_cb=mock.ANY,
+            headers=mock.ANY, retries=mock.ANY,
+            timeout=mock.ANY, infinite=False)
+
+    @mock.patch(MOCKPATH + 'readurl', autospec=True)
     @mock.patch(MOCKPATH + 'EphemeralDHCPv4WithReporting', autospec=True)
     @mock.patch(MOCKPATH + 'net.is_up', autospec=True)
     def test_get_metadata_performs_dhcp_when_network_is_down(
@@ -950,6 +970,43 @@ scbus-1 on xpt0 bus 0
             dsrc.crawl_metadata()
         self.assertEqual(str(cm.exception), error_msg)
 
+    def test_crawl_metadata_call_imds_once_no_reprovision(self):
+        """If reprovisioning, report ready at the end"""
+        ovfenv = construct_valid_ovf_env(
+            platform_settings={"PreprovisionedVm": "False"}
+        )
+
+        data = {
+            'ovfcontent': ovfenv,
+            'sys_cfg': {}
+        }
+        dsrc = self._get_ds(data)
+        dsrc.crawl_metadata()
+        self.assertEqual(1, self.m_get_metadata_from_imds.call_count)
+
+    @mock.patch(
+        'cloudinit.sources.DataSourceAzure.EphemeralDHCPv4WithReporting')
+    @mock.patch('cloudinit.sources.DataSourceAzure.util.write_file')
+    @mock.patch(
+        'cloudinit.sources.DataSourceAzure.DataSourceAzure._report_ready')
+    @mock.patch('cloudinit.sources.DataSourceAzure.DataSourceAzure._poll_imds')
+    def test_crawl_metadata_call_imds_twice_with_reprovision(
+        self, poll_imds_func, m_report_ready, m_write, m_dhcp
+    ):
+        """If reprovisioning, imds metadata will be fetched twice"""
+        ovfenv = construct_valid_ovf_env(
+            platform_settings={"PreprovisionedVm": "True"}
+        )
+
+        data = {
+            'ovfcontent': ovfenv,
+            'sys_cfg': {}
+        }
+        dsrc = self._get_ds(data)
+        poll_imds_func.return_value = ovfenv
+        dsrc.crawl_metadata()
+        self.assertEqual(2, self.m_get_metadata_from_imds.call_count)
+
     @mock.patch(
         'cloudinit.sources.DataSourceAzure.EphemeralDHCPv4WithReporting')
     @mock.patch('cloudinit.sources.DataSourceAzure.util.write_file')
@@ -2638,6 +2695,22 @@ class TestPreprovisioningShouldReprovision(CiTestCase):
         dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
         self.assertFalse(dsa._should_reprovision((None, None, {}, None)))
 
+    @mock.patch(MOCKPATH + 'util.write_file', autospec=True)
+    def test__should_reprovision_uses_imds_md(self, write_file, isfile):
+        """The _should_reprovision method should be able to
+        retrieve the preprovisioning VM type from imds metadata"""
+        isfile.return_value = False
+        dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
+        self.assertTrue(dsa._should_reprovision(
+            (None, None, {}, None),
+            {'extended': {'compute': {'ppsType': 'Running'}}}))
+        self.assertFalse(dsa._should_reprovision(
+            (None, None, {}, None),
+            {}))
+        self.assertFalse(dsa._should_reprovision(
+            (None, None, {}, None),
+            {'extended': {'compute': {"hasCustomData": False}}}))
+
     @mock.patch(MOCKPATH + 'DataSourceAzure._poll_imds')
     def test_reprovision_calls__poll_imds(self, _poll_imds, isfile):
         """_reprovision will poll IMDS."""