Disable ec2 mirror for non aws instances (#390)

For versions before 20.2, we allowed the use of ec2 mirrors if the datasource availability_zone matches one of the ec2 regions. We are now updating that behavior to allow allow the use of ec2 mirrors on ec2 instances or if the user directly passes an an ec2 mirror url through #cloud-config apt directives.

LP: #1456277

2 files changed, 37 insertions, 8 deletions

diff --git a/cloudinit/distros/__init__.py b/cloudinit/distros/__init__.py
index 89940cf0..2fc91bbc 100755
--- a/cloudinit/distros/__init__.py
+++ b/cloudinit/distros/__init__.py
@@ -28,6 +28,9 @@ from cloudinit import type_utils
 from cloudinit import subp
 from cloudinit import util
 
+from cloudinit.features import \
+    ALLOW_EC2_MIRRORS_ON_NON_AWS_INSTANCE_TYPES
+
 from cloudinit.distros.parsers import hosts
 from .networking import LinuxNetworking
 
@@ -849,7 +852,12 @@ def _get_package_mirror_info(mirror_info, data_source=None,
         # ec2 availability zones are named cc-direction-[0-9][a-d] (us-east-1b)
         # the region is us-east-1. so region = az[0:-1]
         if _EC2_AZ_RE.match(data_source.availability_zone):
-            subst['ec2_region'] = "%s" % data_source.availability_zone[0:-1]
+            ec2_region = data_source.availability_zone[0:-1]
+
+            if ALLOW_EC2_MIRRORS_ON_NON_AWS_INSTANCE_TYPES:
+                subst['ec2_region'] = "%s" % ec2_region
+            elif data_source.platform_type == "ec2":
+                subst['ec2_region'] = "%s" % ec2_region
 
     if data_source and data_source.region:
         subst['region'] = data_source.region
diff --git a/cloudinit/distros/tests/test_init.py b/cloudinit/distros/tests/test_init.py
index 40939133..db534654 100644
--- a/cloudinit/distros/tests/test_init.py
+++ b/cloudinit/distros/tests/test_init.py
@@ -67,6 +67,9 @@ class TestGetPackageMirrorInfo:
         assert {'primary': 'http://other'} == _get_package_mirror_info(
             mirror_info, mirror_filter=lambda x: False)
 
+    @pytest.mark.parametrize('allow_ec2_mirror, platform_type', [
+        (True, 'ec2')
+    ])
     @pytest.mark.parametrize('availability_zone,region,patterns,expected', (
         # Test ec2_region alone
         ('fk-fake-1f', None, ['http://EC2-%(ec2_region)s/ubuntu'],
@@ -120,16 +123,34 @@ class TestGetPackageMirrorInfo:
          ['http://%(region)s/ubuntu'], ['http://fk-fake-1/ubuntu'])
         for invalid_char in INVALID_URL_CHARS
     ))
-    def test_substitution(self, availability_zone, region, patterns, expected):
+    def test_valid_substitution(self,
+                                allow_ec2_mirror,
+                                platform_type,
+                                availability_zone,
+                                region,
+                                patterns,
+                                expected):
         """Test substitution works as expected."""
+        flag_path = "cloudinit.distros." \
+                    "ALLOW_EC2_MIRRORS_ON_NON_AWS_INSTANCE_TYPES"
+
         m_data_source = mock.Mock(
-            availability_zone=availability_zone, region=region
+            availability_zone=availability_zone,
+            region=region,
+            platform_type=platform_type
         )
         mirror_info = {'search': {'primary': patterns}}
 
-        ret = _get_package_mirror_info(
-            mirror_info,
-            data_source=m_data_source,
-            mirror_filter=lambda x: x
-        )
+        with mock.patch(flag_path, allow_ec2_mirror):
+            ret = _get_package_mirror_info(
+                mirror_info,
+                data_source=m_data_source,
+                mirror_filter=lambda x: x
+            )
+        print(allow_ec2_mirror)
+        print(platform_type)
+        print(availability_zone)
+        print(region)
+        print(patterns)
+        print(expected)
         assert {'primary': expected} == ret