Drop region validation, add URL normalization and configurable timeout

Author: AndrewAsseily

awscli/customizations/sso/resolve.py

@@ -40,10 +40,21 @@
     ),
 )
 
-_ALL_PARTITIONS = ('aws', 'aws-cn', 'aws-us-gov')
-
 _MAX_REDIRECTS = 1
 
+_DEFAULT_PORTS = {'https': 443, 'http': 80}
+
+_DEFAULT_TIMEOUT = 10
+
+
+def _normalize_url(url):
+    parsed = urllib.parse.urlparse(url)
+    netloc = parsed.hostname or ''
+    if parsed.port and parsed.port != _DEFAULT_PORTS.get(parsed.scheme):
+        netloc = f'{netloc}:{parsed.port}'
+    path = parsed.path.rstrip('/')
+    return urllib.parse.urlunparse((parsed.scheme, netloc, path, '', '', ''))
+
 
 def is_aws_owned_domain(hostname):
     hostname = hostname.lower().rstrip('.')
@@ -66,38 +77,25 @@ def _extract_region_from_hostname(hostname):
     return None
 
 
-def _validate_region(region, session):
-    available = set()
-    for partition in _ALL_PARTITIONS:
-        available.update(
-            session.get_available_regions('sso-oidc', partition_name=partition)
-        )
-    if region not in available:
-        raise ConfigurationError(
-            f"Region '{region}' parsed from the resolved start URL is not "
-            f"a known AWS region. Verify the start URL is correct."
-        )
-
-
-def _follow_redirect(url):
+def _follow_redirect(url, timeout=_DEFAULT_TIMEOUT):
     class _NoRedirectHandler(urllib.request.HTTPRedirectHandler):
         def redirect_request(self, req, fp, code, msg, headers, newurl):
             return None
 
     opener = urllib.request.build_opener(_NoRedirectHandler)
     redirect_codes = (301, 302, 303, 307, 308)
 
-    for _attempt in range(_MAX_REDIRECTS):
+    for _ in range(_MAX_REDIRECTS):
         try:
             req = urllib.request.Request(url, method='HEAD')
-            resp = opener.open(req, timeout=10)
+            resp = opener.open(req, timeout=timeout)
             resp.close()
             return url
         except urllib.error.HTTPError as e:
             if e.code in (405, 501):
                 try:
                     req = urllib.request.Request(url, method='GET')
-                    resp = opener.open(req, timeout=10)
+                    resp = opener.open(req, timeout=timeout)
                     resp.close()
                     return url
                 except urllib.error.HTTPError as e2:
@@ -135,7 +133,9 @@ def redirect_request(self, req, fp, code, msg, headers, newurl):
     return url
 
 
-def resolve_start_url(start_url, session, configured_region=None):
+def resolve_start_url(
+    start_url, session, configured_region=None, timeout=None
+):
     parsed = urllib.parse.urlparse(start_url)
 
     if parsed.scheme != 'https':
@@ -158,7 +158,10 @@ def resolve_start_url(start_url, session, configured_region=None):
     LOG.debug(
         "Start URL '%s' is not AWS-owned, following redirects", start_url
     )
-    resolved_url = _follow_redirect(start_url)
+    effective_timeout = timeout if timeout is not None else _DEFAULT_TIMEOUT
+    resolved_url = _normalize_url(
+        _follow_redirect(start_url, timeout=effective_timeout)
+    )
 
     resolved_hostname = urllib.parse.urlparse(resolved_url).hostname
     if not resolved_hostname or not is_aws_owned_domain(resolved_hostname):
@@ -174,14 +177,13 @@ def resolve_start_url(start_url, session, configured_region=None):
 
     region = _extract_region_from_hostname(resolved_hostname)
 
-    if region:
-        _validate_region(region, session)
-    elif configured_region:
-        region = configured_region
-    else:
-        raise ConfigurationError(
-            f"Cannot determine region from start URL '{start_url}'. "
-            f"Please provide sso_region in your configuration."
-        )
+    if not region:
+        if configured_region:
+            region = configured_region
+        else:
+            raise ConfigurationError(
+                f"Cannot determine region from start URL '{start_url}'. "
+                f"Please provide sso_region in your configuration."
+            )
 
     return resolved_url, region

tests/unit/customizations/sso/test_resolve.py

@@ -18,6 +18,7 @@
 from awscli.customizations.sso.resolve import (
     _extract_region_from_hostname,
     _follow_redirect,
+    _normalize_url,
     is_aws_owned_domain,
     resolve_start_url,
 )
@@ -87,19 +88,25 @@ def test_returns_none_for_region_less_hostnames(self, hostname):
         assert _extract_region_from_hostname(hostname) is None
 
 
+class TestNormalizeUrl:
+    @pytest.mark.parametrize(
+        'url, expected',
+        [
+            ('https://example.com:443/', 'https://example.com'),
+            ('https://example.com:443/path/', 'https://example.com/path'),
+            ('https://example.com/', 'https://example.com'),
+            ('https://example.com:8443/path', 'https://example.com:8443/path'),
+            ('http://example.com:80/', 'http://example.com'),
+            ('https://example.com', 'https://example.com'),
+        ],
+    )
+    def test_normalize_url(self, url, expected):
+        assert _normalize_url(url) == expected
+
+
 class TestResolveStartUrl:
-    def _mock_session(self, regions=None):
-        if regions is None:
-            regions = [
-                'us-east-1',
-                'us-west-2',
-                'eu-west-1',
-                'us-gov-west-1',
-                'cn-north-1',
-            ]
-        session = mock.Mock()
-        session.get_available_regions.return_value = regions
-        return session
+    def _mock_session(self):
+        return mock.Mock()
 
     def test_aws_owned_url_returns_configured_region(self):
         session = self._mock_session()
@@ -139,6 +146,7 @@ def test_aws_owned_url_uses_configured_region_not_hostname(self):
             session=session,
             configured_region='us-east-1',
         )
+        assert resolved_url == 'https://ssoins-abc.portal.us-west-2.app.aws'
         assert region == 'us-east-1'
 
     def test_awsapps_url_requires_configured_region(self):
@@ -172,96 +180,39 @@ def test_invalid_url_raises_error(self):
         with pytest.raises(ConfigurationError, match='Invalid sso_start_url'):
             resolve_start_url('https://', session=session)
 
-    def test_vanity_url_invalid_region_raises_error(self):
-        session = self._mock_session(regions=['us-east-1', 'us-west-2'])
-        with mock.patch(
-            'awscli.customizations.sso.resolve._follow_redirect'
-        ) as mock_follow:
-            mock_follow.return_value = (
-                'https://ssoins-abc.portal.fake-region-1.app.aws'
-            )
-            with pytest.raises(
-                ConfigurationError, match='not a known AWS region'
-            ):
-                resolve_start_url(
-                    'https://aws.mycompany.com',
-                    session=session,
-                )
+    def test_vanity_url_follows_redirect(self):
+        session = self._mock_session()
+        redirect_url = 'https://ssoins-abc.portal.us-east-1.app.aws:443/'
+        normalized_url = 'https://ssoins-abc.portal.us-east-1.app.aws'
 
-    def test_vanity_url_invalid_region_does_not_fall_back_to_configured(self):
-        session = self._mock_session(regions=['us-east-1', 'us-west-2'])
-        with mock.patch(
-            'awscli.customizations.sso.resolve._follow_redirect'
-        ) as mock_follow:
-            mock_follow.return_value = (
-                'https://ssoins-abc.portal.fake-region-1.app.aws'
-            )
-            with pytest.raises(
-                ConfigurationError, match='not a known AWS region'
-            ):
-                resolve_start_url(
-                    'https://aws.mycompany.com',
-                    session=session,
-                    configured_region='us-east-1',
-                )
-
-    def test_vanity_url_govcloud_region_accepted(self):
-        session = mock.Mock()
-        session.get_available_regions.side_effect = (
-            lambda service, partition_name='aws': {
-                'aws': ['us-east-1', 'us-west-2'],
-                'aws-us-gov': ['us-gov-west-1'],
-                'aws-cn': ['cn-north-1'],
-            }.get(partition_name, [])
-        )
         with mock.patch(
             'awscli.customizations.sso.resolve._follow_redirect'
         ) as mock_follow:
-            mock_follow.return_value = (
-                'https://ssoins-abc.portal.us-gov-west-1.app.aws'
-            )
+            mock_follow.return_value = redirect_url
             resolved_url, region = resolve_start_url(
                 'https://aws.mycompany.com',
                 session=session,
             )
-            assert region == 'us-gov-west-1'
-
-    def test_vanity_url_china_region_accepted(self):
-        session = mock.Mock()
-        session.get_available_regions.side_effect = (
-            lambda service, partition_name='aws': {
-                'aws': ['us-east-1', 'us-west-2'],
-                'aws-us-gov': ['us-gov-west-1'],
-                'aws-cn': ['cn-north-1'],
-            }.get(partition_name, [])
-        )
-        with mock.patch(
-            'awscli.customizations.sso.resolve._follow_redirect'
-        ) as mock_follow:
-            mock_follow.return_value = (
-                'https://ssoins-abc.portal.cn-north-1.app.aws'
-            )
-            resolved_url, region = resolve_start_url(
-                'https://aws.mycompany.com',
-                session=session,
+            assert resolved_url == normalized_url
+            assert region == 'us-east-1'
+            mock_follow.assert_called_once_with(
+                'https://aws.mycompany.com', timeout=10
             )
-            assert region == 'cn-north-1'
 
-    def test_vanity_url_follows_redirect(self):
+    def test_vanity_url_uses_parsed_region_not_configured(self):
         session = self._mock_session()
-        redirect_url = 'https://ssoins-abc.portal.us-east-1.app.aws:443/'
-
         with mock.patch(
             'awscli.customizations.sso.resolve._follow_redirect'
         ) as mock_follow:
-            mock_follow.return_value = redirect_url
+            mock_follow.return_value = (
+                'https://ssoins-abc.portal.us-west-2.app.aws'
+            )
             resolved_url, region = resolve_start_url(
                 'https://aws.mycompany.com',
                 session=session,
+                configured_region='eu-west-1',
             )
-            assert resolved_url == redirect_url
-            assert region == 'us-east-1'
-            mock_follow.assert_called_once_with('https://aws.mycompany.com')
+            assert region == 'us-west-2'
 
     def test_vanity_url_resolves_to_non_aws_domain_raises_error(self):
         session = self._mock_session()
@@ -283,6 +234,47 @@ def test_vanity_url_resolves_to_http_raises_error(self):
             with pytest.raises(ConfigurationError, match='must use https'):
                 resolve_start_url('https://aws.mycompany.com', session=session)
 
+    def test_vanity_url_region_less_requires_configured_region(self):
+        session = self._mock_session()
+        with mock.patch(
+            'awscli.customizations.sso.resolve._follow_redirect'
+        ) as mock_follow:
+            mock_follow.return_value = 'https://d-abc123.awsapps.com/start'
+            with pytest.raises(
+                ConfigurationError, match='Cannot determine region'
+            ):
+                resolve_start_url('https://aws.mycompany.com', session=session)
+
+    def test_vanity_url_region_less_uses_configured_region(self):
+        session = self._mock_session()
+        with mock.patch(
+            'awscli.customizations.sso.resolve._follow_redirect'
+        ) as mock_follow:
+            mock_follow.return_value = 'https://d-abc123.awsapps.com/start'
+            resolved_url, region = resolve_start_url(
+                'https://aws.mycompany.com',
+                session=session,
+                configured_region='us-east-1',
+            )
+            assert region == 'us-east-1'
+
+    def test_custom_timeout_passed_to_follow_redirect(self):
+        session = self._mock_session()
+        with mock.patch(
+            'awscli.customizations.sso.resolve._follow_redirect'
+        ) as mock_follow:
+            mock_follow.return_value = (
+                'https://ssoins-abc.portal.us-east-1.app.aws'
+            )
+            resolve_start_url(
+                'https://aws.mycompany.com',
+                session=session,
+                timeout=5,
+            )
+            mock_follow.assert_called_once_with(
+                'https://aws.mycompany.com', timeout=5
+            )
+
 
 class TestFollowRedirect:
     def _make_http_error(self, code, headers=None):