Resolve IMDS region from placement/region endpoint

.changes/next-release/bugfix-IMDS-region-local-zone.json

@@ -0,0 +1,5 @@
+{
+  "type": "bugfix",
+  "category": "``IMDS``",
+  "description": "Resolve the region from the IMDS ``placement/region`` endpoint instead of deriving it from the availability zone. The previous heuristic of stripping the last character of the availability zone produced an invalid region for Local Zones and Wavelength Zones (e.g. ``us-east-2-sbn-1a`` resolved to ``us-east-2-sbn-1`` instead of ``us-east-2``)."
+}

awscli/utils.py

@@ -163,7 +163,8 @@ def _create_fetcher(self):
 
 
 class InstanceMetadataRegionFetcher(IMDSFetcher):
-    _URL_PATH = 'latest/meta-data/placement/availability-zone/'
+    _URL_PATH = 'latest/meta-data/placement/region/'
+    _AZ_URL_PATH = 'latest/meta-data/placement/availability-zone/'
 
     def retrieve_region(self):
         """Get the current region from the instance metadata service.
@@ -199,14 +200,33 @@ def retrieve_region(self):
 
     def _get_region(self):
         token = self._fetch_metadata_token()
-        response = self._get_request(
-            url_path=self._URL_PATH,
-            retry_func=self._default_retry,
-            token=token,
-        )
-        availability_zone = response.text
-        region = availability_zone[:-1]
-        return region
+        try:
+            response = self._get_request(
+                url_path=self._URL_PATH,
+                retry_func=self._default_retry,
+                token=token,
+            )
+            return response.text
+        except (self._RETRIES_EXCEEDED_ERROR_CLS, BadIMDSRequestError):
+            # The placement/region endpoint may be unavailable on older or
+            # third-party IMDS implementations. Fall back to deriving the
+            # region from the availability zone. Note this heuristic is wrong
+            # for Local Zones and Wavelength Zones (e.g. the AZ
+            # ``us-east-2-sbn-1a`` belongs to region ``us-east-2``, not
+            # ``us-east-2-sbn-1``), which is why placement/region is preferred.
+            logger.debug(
+                "Failed to retrieve region from IMDS placement/region "
+                "endpoint, falling back to deriving it from the "
+                "availability zone."
+            )
+            response = self._get_request(
+                url_path=self._AZ_URL_PATH,
+                retry_func=self._default_retry,
+                token=token,
+            )
+            availability_zone = response.text
+            region = availability_zone[:-1]
+            return region
 
 
 def split_on_commas(value):

tests/functional/test_clidriver.py

@@ -66,9 +66,8 @@ def test_imds_region_is_used_as_fallback_wo_v2_support(self):
         # First response should be from the IMDS server for security token
         # if server supports IMDSv1 only there will be no response for token
         self.add_response(None)
-        # Then another response from the IMDS server for an availability
-        # zone.
-        self.add_response(b'us-mars-2a')
+        # Then another response from the IMDS server for the region.
+        self.add_response(b'us-mars-2')
         # Once a region is fetched form the IMDS server we need to mock an
         # XML response from ec2 so that the CLI driver doesn't throw an error
         # during parsing.
@@ -84,8 +83,25 @@ def test_imds_region_is_used_as_fallback_with_v2_support(self):
         # First response should be from the IMDS server for security token
         # if server supports IMDSv2 it'll return token
         self.add_response(b'token')
-        # Then another response from the IMDS server for an availability
-        # zone.
+        # Then another response from the IMDS server for the region.
+        self.add_response(b'us-mars-2')
+        # Once a region is fetched form the IMDS server we need to mock an
+        # XML response from ec2 so that the CLI driver doesn't throw an error
+        # during parsing.
+        self.add_response(b'<?xml version="1.0" ?><foo><bar>text</bar></foo>')
+        capture = RegionCapture()
+        self.session.register('before-send.ec2.*', capture)
+        self.driver.main(['ec2', 'describe-instances'])
+        self.assertEqual(capture.region, 'us-mars-2')
+
+    def test_imds_region_falls_back_to_availability_zone(self):
+        # Remove region override from the environment variables.
+        self.environ.pop('AWS_DEFAULT_REGION', 0)
+        # First response should be from the IMDS server for security token.
+        self.add_response(b'token')
+        # If the placement/region endpoint is unavailable, the region
+        # fetcher falls back to the availability-zone endpoint.
+        self.add_response(b'', status_code=404)
         self.add_response(b'us-mars-2a')
         # Once a region is fetched form the IMDS server we need to mock an
         # XML response from ec2 so that the CLI driver doesn't throw an error

tests/unit/test_utils.py

@@ -399,7 +399,7 @@ def setUp(self):
         self._send = self._urllib3_patch.start()
         self._imds_responses = []
         self._send.side_effect = self.get_imds_response
-        self._region = 'us-mars-1a'
+        self._region = 'us-mars-1'
         self.environ = {}
         self.environ_patch = mock.patch('os.environ', self.environ)
         self.environ_patch.start()
@@ -463,6 +463,38 @@ def test_disabling_env_var_not_true(self):
         expected_result = 'us-mars-1'
         self.assertEqual(result, expected_result)
 
+    def test_uses_placement_region_endpoint(self):
+        self.add_imds_token_response()
+        self.add_get_region_imds_response()
+        result = InstanceMetadataRegionFetcher().retrieve_region()
+        self.assertEqual(result, 'us-mars-1')
+        # The region must come from the placement/region endpoint, not the
+        # availability-zone endpoint.
+        region_request = self._send.call_args[0][0]
+        self.assertTrue(region_request.url.endswith('placement/region/'))
+
+    def test_local_zone_resolves_to_parent_region(self):
+        # For a Local Zone, placement/region returns the parent region
+        # directly. Deriving it from the AZ (e.g. stripping the last char of
+        # ``us-east-2-sbn-1a``) would yield the invalid ``us-east-2-sbn-1``.
+        self.add_imds_token_response()
+        self.add_get_region_imds_response(region='us-east-2')
+        result = InstanceMetadataRegionFetcher().retrieve_region()
+        self.assertEqual(result, 'us-east-2')
+
+    def test_falls_back_to_availability_zone(self):
+        # If placement/region is unavailable (older or third-party IMDS),
+        # fall back to deriving the region from the availability zone.
+        self.add_imds_token_response()
+        self.add_imds_response(status_code=404, body=b'')
+        self.add_get_region_imds_response(region='us-mars-1a')
+        result = InstanceMetadataRegionFetcher(num_attempts=1).retrieve_region()
+        self.assertEqual(result, 'us-mars-1')
+        # First the region endpoint is tried, then the AZ endpoint.
+        urls = [call[0][0].url for call in self._send.call_args_list]
+        self.assertTrue(urls[-2].endswith('placement/region/'))
+        self.assertTrue(urls[-1].endswith('placement/availability-zone/'))
+
     def test_includes_user_agent_header(self):
         user_agent = 'my-user-agent'
         self.add_imds_token_response()
@@ -536,6 +568,9 @@ def test_empty_response_is_retried(self):
 
     def test_exhaust_retries_on_region_request(self):
         self.add_imds_token_response()
+        # The region endpoint fails, and so does the availability-zone
+        # fallback, so no region can be resolved.
+        self.add_imds_response(status_code=400, body=b'')
         self.add_imds_response(status_code=400, body=b'')
         result = InstanceMetadataRegionFetcher(
             num_attempts=1