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

7 files changed, 193 insertions, 89 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
diff --git a/cloudinit/features.py b/cloudinit/features.py
index e455213d..c44fa29e 100644
--- a/cloudinit/features.py
+++ b/cloudinit/features.py
@@ -26,6 +26,17 @@ After the 20.2 release, we instead raise an exception.
 This flag can be removed after Focal is no longer supported
 """
 
+
+ALLOW_EC2_MIRRORS_ON_NON_AWS_INSTANCE_TYPES = False
+"""
+When configuring apt mirrors, old behavior is to allow
+the use of ec2 mirrors if the datasource availability_zone format
+matches one of the possible aws ec2 regions. After the 20.2 release, we
+no longer publish ec2 region mirror urls on non-AWS cloud platforms.
+Besides feature_overrides.py, users can override this by providing
+#cloud-config apt directives.
+"""
+
 try:
     # pylint: disable=wildcard-import
     from cloudinit.feature_overrides import *  # noqa
diff --git a/tests/unittests/test_distros/test_generic.py b/tests/unittests/test_distros/test_generic.py
index 6b6c1566..44607489 100644
--- a/tests/unittests/test_distros/test_generic.py
+++ b/tests/unittests/test_distros/test_generic.py
@@ -6,6 +6,7 @@ from cloudinit import util
 from cloudinit.tests import helpers
 
 import os
+import pytest
 import shutil
 import tempfile
 from unittest import mock
@@ -37,24 +38,6 @@ gapmi = distros._get_arch_package_mirror_info
 
 class TestGenericDistro(helpers.FilesystemMockingTestCase):
 
-    def return_first(self, mlist):
-        if not mlist:
-            return None
-        return mlist[0]
-
-    def return_second(self, mlist):
-        if not mlist:
-            return None
-        return mlist[1]
-
-    def return_none(self, _mlist):
-        return None
-
-    def return_last(self, mlist):
-        if not mlist:
-            return None
-        return(mlist[-1])
-
     def setUp(self):
         super(TestGenericDistro, self).setUp()
         # Make a temp directoy for tests to use.
@@ -145,61 +128,6 @@ class TestGenericDistro(helpers.FilesystemMockingTestCase):
         arch_mirrors = gapmi(package_mirrors, arch="amd64")
         self.assertEqual(package_mirrors[0], arch_mirrors)
 
-    def test_get_package_mirror_info_az_ec2(self):
-        arch_mirrors = gapmi(package_mirrors, arch="amd64")
-        data_source_mock = mock.Mock(availability_zone="us-east-1a")
-
-        results = gpmi(arch_mirrors, data_source=data_source_mock,
-                       mirror_filter=self.return_first)
-        self.assertEqual(results,
-                         {'primary': 'http://us-east-1.ec2/',
-                          'security': 'http://security-mirror1-intel'})
-
-        results = gpmi(arch_mirrors, data_source=data_source_mock,
-                       mirror_filter=self.return_second)
-        self.assertEqual(results,
-                         {'primary': 'http://us-east-1a.clouds/',
-                          'security': 'http://security-mirror2-intel'})
-
-        results = gpmi(arch_mirrors, data_source=data_source_mock,
-                       mirror_filter=self.return_none)
-        self.assertEqual(results, package_mirrors[0]['failsafe'])
-
-    def test_get_package_mirror_info_az_non_ec2(self):
-        arch_mirrors = gapmi(package_mirrors, arch="amd64")
-        data_source_mock = mock.Mock(availability_zone="nova.cloudvendor")
-
-        results = gpmi(arch_mirrors, data_source=data_source_mock,
-                       mirror_filter=self.return_first)
-        self.assertEqual(results,
-                         {'primary': 'http://nova.cloudvendor.clouds/',
-                          'security': 'http://security-mirror1-intel'})
-
-        results = gpmi(arch_mirrors, data_source=data_source_mock,
-                       mirror_filter=self.return_last)
-        self.assertEqual(results,
-                         {'primary': 'http://nova.cloudvendor.clouds/',
-                          'security': 'http://security-mirror2-intel'})
-
-    def test_get_package_mirror_info_none(self):
-        arch_mirrors = gapmi(package_mirrors, arch="amd64")
-        data_source_mock = mock.Mock(availability_zone=None)
-
-        # because both search entries here replacement based on
-        # availability-zone, the filter will be called with an empty list and
-        # failsafe should be taken.
-        results = gpmi(arch_mirrors, data_source=data_source_mock,
-                       mirror_filter=self.return_first)
-        self.assertEqual(results,
-                         {'primary': 'http://fs-primary-intel',
-                          'security': 'http://security-mirror1-intel'})
-
-        results = gpmi(arch_mirrors, data_source=data_source_mock,
-                       mirror_filter=self.return_last)
-        self.assertEqual(results,
-                         {'primary': 'http://fs-primary-intel',
-                          'security': 'http://security-mirror2-intel'})
-
     def test_systemd_in_use(self):
         cls = distros.fetch("ubuntu")
         d = cls("ubuntu", {}, None)
@@ -259,4 +187,116 @@ class TestGenericDistro(helpers.FilesystemMockingTestCase):
             ["pw", "usermod", "myuser", "-p", "01-Jan-1970"])
 
 
+class TestGetPackageMirrors:
+
+    def return_first(self, mlist):
+        if not mlist:
+            return None
+        return mlist[0]
+
+    def return_second(self, mlist):
+        if not mlist:
+            return None
+
+        return mlist[1] if len(mlist) > 1 else None
+
+    def return_none(self, _mlist):
+        return None
+
+    def return_last(self, mlist):
+        if not mlist:
+            return None
+        return(mlist[-1])
+
+    @pytest.mark.parametrize(
+        "allow_ec2_mirror, platform_type, mirrors",
+        [
+            (True, "ec2", [
+                {'primary': 'http://us-east-1.ec2/',
+                 'security': 'http://security-mirror1-intel'},
+                {'primary': 'http://us-east-1a.clouds/',
+                 'security': 'http://security-mirror2-intel'}
+            ]),
+            (True, "other", [
+                {'primary': 'http://us-east-1.ec2/',
+                 'security': 'http://security-mirror1-intel'},
+                {'primary': 'http://us-east-1a.clouds/',
+                 'security': 'http://security-mirror2-intel'}
+            ]),
+            (False, "ec2", [
+                {'primary': 'http://us-east-1.ec2/',
+                 'security': 'http://security-mirror1-intel'},
+                {'primary': 'http://us-east-1a.clouds/',
+                 'security': 'http://security-mirror2-intel'}
+            ]),
+            (False, "other", [
+                {'primary': 'http://us-east-1a.clouds/',
+                 'security': 'http://security-mirror1-intel'},
+                {'primary': 'http://fs-primary-intel',
+                 'security': 'http://security-mirror2-intel'}
+            ])
+        ])
+    def test_get_package_mirror_info_az_ec2(self,
+                                            allow_ec2_mirror,
+                                            platform_type,
+                                            mirrors):
+        flag_path = "cloudinit.distros." \
+                    "ALLOW_EC2_MIRRORS_ON_NON_AWS_INSTANCE_TYPES"
+        with mock.patch(flag_path, allow_ec2_mirror):
+            arch_mirrors = gapmi(package_mirrors, arch="amd64")
+            data_source_mock = mock.Mock(
+                availability_zone="us-east-1a",
+                platform_type=platform_type)
+
+            results = gpmi(arch_mirrors, data_source=data_source_mock,
+                           mirror_filter=self.return_first)
+            assert(results == mirrors[0])
+
+            results = gpmi(arch_mirrors, data_source=data_source_mock,
+                           mirror_filter=self.return_second)
+            assert(results == mirrors[1])
+
+            results = gpmi(arch_mirrors, data_source=data_source_mock,
+                           mirror_filter=self.return_none)
+            assert(results == package_mirrors[0]['failsafe'])
+
+    def test_get_package_mirror_info_az_non_ec2(self):
+        arch_mirrors = gapmi(package_mirrors, arch="amd64")
+        data_source_mock = mock.Mock(availability_zone="nova.cloudvendor")
+
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
+                       mirror_filter=self.return_first)
+        assert(results == {
+            'primary': 'http://nova.cloudvendor.clouds/',
+            'security': 'http://security-mirror1-intel'}
+        )
+
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
+                       mirror_filter=self.return_last)
+        assert(results == {
+            'primary': 'http://nova.cloudvendor.clouds/',
+            'security': 'http://security-mirror2-intel'}
+        )
+
+    def test_get_package_mirror_info_none(self):
+        arch_mirrors = gapmi(package_mirrors, arch="amd64")
+        data_source_mock = mock.Mock(availability_zone=None)
+
+        # because both search entries here replacement based on
+        # availability-zone, the filter will be called with an empty list and
+        # failsafe should be taken.
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
+                       mirror_filter=self.return_first)
+        assert(results == {
+            'primary': 'http://fs-primary-intel',
+            'security': 'http://security-mirror1-intel'}
+        )
+
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
+                       mirror_filter=self.return_last)
+        assert(results == {
+            'primary': 'http://fs-primary-intel',
+            'security': 'http://security-mirror2-intel'}
+        )
+
 # vi: ts=4 expandtab
diff --git a/tests/unittests/test_handler/test_handler_apt_configure_sources_list_v1.py b/tests/unittests/test_handler/test_handler_apt_configure_sources_list_v1.py
index e5382544..369480be 100644
--- a/tests/unittests/test_handler/test_handler_apt_configure_sources_list_v1.py
+++ b/tests/unittests/test_handler/test_handler_apt_configure_sources_list_v1.py
@@ -101,6 +101,7 @@ class TestAptSourceConfigSourceList(t_help.FilesystemMockingTestCase):
             cfg = {'apt_mirror_search': mirror}
         else:
             cfg = {'apt_mirror': mirror}
+
         mycloud = self._get_cloud(distro)
 
         with mock.patch.object(util, 'write_file') as mockwf:
@@ -108,8 +109,9 @@ class TestAptSourceConfigSourceList(t_help.FilesystemMockingTestCase):
                                    return_value="faketmpl") as mocklf:
                 with mock.patch.object(os.path, 'isfile',
                                        return_value=True) as mockisfile:
-                    with mock.patch.object(templater, 'render_string',
-                                           return_value="fake") as mockrnd:
+                    with mock.patch.object(
+                        templater, 'render_string',
+                            return_value='fake') as mockrnd:
                         with mock.patch.object(util, 'rename'):
                             cc_apt_configure.handle("test", cfg, mycloud,
                                                     LOG, None)
diff --git a/tests/unittests/test_handler/test_handler_apt_source_v1.py b/tests/unittests/test_handler/test_handler_apt_source_v1.py
index f2349157..367971cb 100644
--- a/tests/unittests/test_handler/test_handler_apt_source_v1.py
+++ b/tests/unittests/test_handler/test_handler_apt_source_v1.py
@@ -43,10 +43,17 @@ class FakeDistro(object):
         return
 
 
+class FakeDatasource:
+    """Fake Datasource helper object"""
+    def __init__(self):
+        self.region = 'region'
+
+
 class FakeCloud(object):
     """Fake Cloud helper object"""
     def __init__(self):
         self.distro = FakeDistro()
+        self.datasource = FakeDatasource()
 
 
 class TestAptSourceConfig(TestCase):
diff --git a/tests/unittests/test_handler/test_handler_apt_source_v3.py b/tests/unittests/test_handler/test_handler_apt_source_v3.py
index 220100e2..ac847238 100644
--- a/tests/unittests/test_handler/test_handler_apt_source_v3.py
+++ b/tests/unittests/test_handler/test_handler_apt_source_v3.py
@@ -49,6 +49,18 @@ MOCK_LSB_RELEASE_DATA = {
     'release': '18.04', 'codename': 'bionic'}
 
 
+class FakeDatasource:
+    """Fake Datasource helper object"""
+    def __init__(self):
+        self.region = 'region'
+
+
+class FakeCloud:
+    """Fake Cloud helper object"""
+    def __init__(self):
+        self.datasource = FakeDatasource()
+
+
 class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
     """TestAptSourceConfig
     Main Class to test apt configs
@@ -471,7 +483,7 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
         fromfn = ("%s/%s_%s" % (pre, archive, post))
         tofn = ("%s/test.ubuntu.com_%s" % (pre, post))
 
-        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, arch)
+        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, FakeCloud(), arch)
 
         self.assertEqual(mirrors['MIRROR'],
                          "http://test.ubuntu.com/%s/" % component)
@@ -559,7 +571,8 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
                "security": [{'arches': ["default"],
                              "uri": smir}]}
 
-        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, 'amd64')
+        mirrors = cc_apt_configure.find_apt_mirror_info(
+            cfg, FakeCloud(), 'amd64')
 
         self.assertEqual(mirrors['MIRROR'],
                          pmir)
@@ -594,7 +607,7 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
                "security": [{'arches': ["default"], "uri": "nothis-security"},
                             {'arches': [arch], "uri": smir}]}
 
-        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, arch)
+        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, FakeCloud(), arch)
 
         self.assertEqual(mirrors['PRIMARY'], pmir)
         self.assertEqual(mirrors['MIRROR'], pmir)
@@ -613,7 +626,8 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
                             {'arches': ["default"],
                              "uri": smir}]}
 
-        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, 'amd64')
+        mirrors = cc_apt_configure.find_apt_mirror_info(
+            cfg, FakeCloud(), 'amd64')
 
         self.assertEqual(mirrors['MIRROR'],
                          pmir)
@@ -673,7 +687,7 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
 
         with mock.patch.object(cc_apt_configure.util, 'search_for_mirror',
                                side_effect=[pmir, smir]) as mocksearch:
-            mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None,
+            mirrors = cc_apt_configure.find_apt_mirror_info(cfg, FakeCloud(),
                                                             'amd64')
 
         calls = [call(["pfailme", pmir]),
@@ -712,7 +726,8 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
         # should not be called, since primary is specified
         with mock.patch.object(cc_apt_configure.util,
                                'search_for_mirror') as mockse:
-            mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, arch)
+            mirrors = cc_apt_configure.find_apt_mirror_info(
+                cfg, FakeCloud(), arch)
         mockse.assert_not_called()
 
         self.assertEqual(mirrors['MIRROR'],