Skip to content

Commit

Permalink
Plumb IMDS v1 opt out configuration to region provider
Browse files Browse the repository at this point in the history
The AWS CLI v2 currently has a seperate, but similar, implementation
from upstream botocore. In the future, we should be able to
consolidate the implementations to make porting more streamlined.
  • Loading branch information
kyleknap committed Dec 11, 2023
1 parent 44214bf commit f3c81a6
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 33 deletions.
5 changes: 4 additions & 1 deletion awscli/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,10 @@ def _create_fetcher(self):
'ec2_metadata_service_endpoint'),
'ec2_metadata_service_endpoint_mode': resolve_imds_endpoint_mode(
self._session
)
),
'ec2_metadata_v1_disabled': self._session.get_config_variable(
'ec2_metadata_v1_disabled'
),
}
fetcher = InstanceMetadataRegionFetcher(
timeout=metadata_timeout,
Expand Down
77 changes: 45 additions & 32 deletions tests/unit/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import botocore
import botocore.model
import botocore.session as session
from botocore.exceptions import ConnectionClosedError
from botocore.exceptions import ConnectionClosedError, MetadataRetrievalError
from awscli.clidriver import create_clidriver
from awscli.testutils import unittest, skip_if_windows, mock
from awscli.compat import is_windows
Expand Down Expand Up @@ -316,17 +316,21 @@ def test_not_create_process_if_stream_not_created(self):
self.assertEqual(self.popen.call_count, 0)


class TestInstanceMetadataRegionFetcher(unittest.TestCase):
class BaseIMDSRegionTest(unittest.TestCase):
def setUp(self):
urllib3_session_send = 'botocore.httpsession.URLLib3Session.send'
self._urllib3_patch = mock.patch(urllib3_session_send)
self._send = self._urllib3_patch.start()
self._imds_responses = []
self._send.side_effect = self.get_imds_response
self._region = 'us-mars-1a'
self.environ = {}
self.environ_patch = mock.patch('os.environ', self.environ)
self.environ_patch.start()

def tearDown(self):
self._urllib3_patch.stop()
self.environ_patch.stop()

def add_imds_response(self, body, status_code=200):
response = botocore.awsrequest.AWSResponse(
Expand Down Expand Up @@ -354,6 +358,8 @@ def get_imds_response(self, request):
raise response
return response


class TestInstanceMetadataRegionFetcher(BaseIMDSRegionTest):
def test_disabled_by_environment(self):
env = {'AWS_EC2_METADATA_DISABLED': 'true'}
fetcher = InstanceMetadataRegionFetcher(env=env)
Expand Down Expand Up @@ -454,15 +460,7 @@ def test_exhaust_retries_on_region_request(self):
self.assertEqual(result, None)


class TestIMDSRegionProvider(unittest.TestCase):
def setUp(self):
self.environ = {}
self.environ_patch = mock.patch('os.environ', self.environ)
self.environ_patch.start()

def tearDown(self):
self.environ_patch.stop()

class TestIMDSRegionProvider(BaseIMDSRegionTest):
def assert_does_provide_expected_value(self, fetcher_region=None,
expected_result=None,):
fake_session = mock.Mock(spec=session.Session)
Expand All @@ -484,74 +482,89 @@ def test_does_provide_none(self):
expected_result=None,
)

@mock.patch('botocore.httpsession.URLLib3Session.send')
def test_use_truncated_user_agent(self, send):
def test_use_truncated_user_agent(self):
driver = create_clidriver()
driver.session.user_agent_version = '3.0'
self.add_imds_token_response()
self.add_get_region_imds_response()
provider = IMDSRegionProvider(driver.session)
provider.provide()
args, _ = send.call_args
args, _ = self._send.call_args
self.assertEqual(args[0].headers['User-Agent'], 'aws-cli/3.0')

@mock.patch('botocore.httpsession.URLLib3Session.send')
def test_can_use_ipv6(self, send):
def test_can_use_ipv6(self):
driver = create_clidriver()
driver.session.set_config_variable('imds_use_ipv6', True)
self.add_imds_token_response()
self.add_get_region_imds_response()
provider = IMDSRegionProvider(driver.session)
provider.provide()
args, _ = send.call_args
args, _ = self._send.call_args
self.assertIn('[fd00:ec2::254]', args[0].url)

@mock.patch('botocore.httpsession.URLLib3Session.send')
def test_use_ipv4_by_default(self, send):
def test_use_ipv4_by_default(self):
driver = create_clidriver()
self.add_imds_token_response()
self.add_get_region_imds_response()
provider = IMDSRegionProvider(driver.session)
provider.provide()
args, _ = send.call_args
args, _ = self._send.call_args
self.assertIn('169.254.169.254', args[0].url)

@mock.patch('botocore.httpsession.URLLib3Session.send')
def test_can_set_imds_endpoint_mode_to_ipv4(self, send):
def test_can_set_imds_endpoint_mode_to_ipv4(self):
driver = create_clidriver()
driver.session.set_config_variable(
'ec2_metadata_service_endpoint_mode', 'ipv4')
self.add_imds_token_response()
self.add_get_region_imds_response()
provider = IMDSRegionProvider(driver.session)
provider.provide()
args, _ = send.call_args
args, _ = self._send.call_args
self.assertIn('169.254.169.254', args[0].url)

@mock.patch('botocore.httpsession.URLLib3Session.send')
def test_can_set_imds_endpoint_mode_to_ipv6(self, send):
def test_can_set_imds_endpoint_mode_to_ipv6(self):
driver = create_clidriver()
driver.session.set_config_variable(
'ec2_metadata_service_endpoint_mode', 'ipv6')
self.add_imds_token_response()
self.add_get_region_imds_response()
provider = IMDSRegionProvider(driver.session)
provider.provide()
args, _ = send.call_args
args, _ = self._send.call_args
self.assertIn('[fd00:ec2::254]', args[0].url)

@mock.patch('botocore.httpsession.URLLib3Session.send')
def test_can_set_imds_service_endpoint(self, send):
def test_can_set_imds_service_endpoint(self):
driver = create_clidriver()
driver.session.set_config_variable(
'ec2_metadata_service_endpoint', 'http://myendpoint/')
self.add_imds_token_response()
self.add_get_region_imds_response()
provider = IMDSRegionProvider(driver.session)
provider.provide()
args, _ = send.call_args
args, _ = self._send.call_args
self.assertIn('http://myendpoint/', args[0].url)

@mock.patch('botocore.httpsession.URLLib3Session.send')
def test_imds_service_endpoint_overrides_ipv6_endpoint(self, send):
def test_imds_service_endpoint_overrides_ipv6_endpoint(self):
driver = create_clidriver()
driver.session.set_config_variable(
'ec2_metadata_service_endpoint_mode', 'ipv6')
driver.session.set_config_variable(
'ec2_metadata_service_endpoint', 'http://myendpoint/')
self.add_imds_token_response()
self.add_get_region_imds_response()
provider = IMDSRegionProvider(driver.session)
provider.provide()
args, _ = send.call_args
args, _ = self._send.call_args
self.assertIn('http://myendpoint/', args[0].url)

def test_disable_ec2_metadata_v1(self):
driver = create_clidriver()
driver.session.set_config_variable('ec2_metadata_v1_disabled', 'true')
self.add_imds_response(b'', 404)
provider = IMDSRegionProvider(driver.session)
with self.assertRaises(MetadataRetrievalError):
provider.provide()


class TestLazyPager(unittest.TestCase):
def setUp(self):
Expand Down

0 comments on commit f3c81a6

Please sign in to comment.