Migrate from IMDSv1 to IMDSv2 (#78)

IMDSv2 uses session authentication to retrieve EC2 instance metadata.
https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html

Co-authored-by: Himtanaya Bhadada <himmy@amazon.com>

Author: Himtanaya

aws_embedded_metrics/environment/ec2_environment.py

@@ -24,21 +24,34 @@
 log = logging.getLogger(__name__)
 Config = get_config()
 
+# Documentation for configuring instance metadata can be found here:
+# https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
+TOKEN_ENDPOINT = "http://169.254.169.254/latest/api/token"
+TOKEN_REQUEST_HEADER_KEY = "X-aws-ec2-metadata-token-ttl-seconds"
+TOKEN_REQUEST_HEADER_VALUE = "21600"
 DEFAULT_EC2_METADATA_ENDPOINT = (
     "http://169.254.169.254/latest/dynamic/instance-identity/document"
 )
+METADATA_REQUEST_HEADER_KEY = "X-aws-ec2-metadata-token"
 
 
-async def fetch(  # type: ignore
-    session: aiohttp.ClientSession, url: str
+async def fetchJSON(
+    session: aiohttp.ClientSession, method: str, url: str, headers: Dict[str, str],
 ) -> Dict[str, Any]:
-    async with session.get(url, timeout=2) as response:
+    async with session.request(method, url, timeout=2, headers=headers) as response:
         # content_type=None prevents validation of the HTTP Content-Type header
         # The EC2 metadata endpoint uses text/plain instead of application/json
         # https://github.com/aio-libs/aiohttp/blob/7f777333a4ec0043ddf2e8d67146a626089773d9/aiohttp/web_request.py#L582-L585
         return cast(Dict[str, Any], await response.json(content_type=None))
 
 
+async def fetchString(
+    session: aiohttp.ClientSession, method: str, url: str, headers: Dict[str, str]
+) -> str:
+    async with session.request(method, url, timeout=2, headers=headers) as response:
+        return await response.text()
+
+
 class EC2Environment(Environment):
     def __init__(self) -> None:
         self.sink: Optional[AgentSink] = None
@@ -48,10 +61,22 @@ async def probe(self) -> bool:
             metadata_endpoint = (
                 Config.ec2_metadata_endpoint or DEFAULT_EC2_METADATA_ENDPOINT
             )
+            token_header = {TOKEN_REQUEST_HEADER_KEY: TOKEN_REQUEST_HEADER_VALUE}
+            log.info("Fetching token for EC2 metadata request from: %s", TOKEN_ENDPOINT)
+            try:
+                token = await fetchString(session, "PUT", TOKEN_ENDPOINT, token_header)
+                log.debug("Received token for request to EC2 metadata endpoint.")
+            except Exception:
+                log.info(
+                    "Failed to fetch token for EC2 metadata request from %s", TOKEN_ENDPOINT
+                )
+                return False
+
             log.info("Fetching EC2 metadata from: %s", metadata_endpoint)
             try:
-                response_json = await fetch(session, metadata_endpoint)
-                log.debug("Received response from EC2 metdata endpoint.")
+                metadata_request_header = {METADATA_REQUEST_HEADER_KEY: token}
+                response_json = await fetchJSON(session, "GET", metadata_endpoint, metadata_request_header)
+                log.debug("Received response from EC2 metadata endpoint.")
                 self.metadata = response_json
                 return True
             except Exception:

tests/environment/test_ec2_environment.py

@@ -16,7 +16,7 @@
 @pytest.mark.asyncio
 async def test_probe_returns_true_if_fetch_succeeds(aresponses):
     # arrange
-    configure_response(aresponses, "{}")
+    configure_response(aresponses, fake.pystr(), "{}")
     env = EC2Environment()
 
     # act
@@ -55,7 +55,7 @@ def test_get_name_returns_configured_name():
 async def test_get_type_returns_ec2_instance(aresponses):
     # arrange
     expected = "AWS::EC2::Instance"
-    configure_response(aresponses, "{}")
+    configure_response(aresponses, fake.pystr(), "{}")
     env = EC2Environment()
 
     # environment MUST be detected before we can access the metadata
@@ -79,6 +79,7 @@ async def test_configure_context_adds_ec2_metadata_props(aresponses):
 
     configure_response(
         aresponses,
+        fake.pystr(),
         json.dumps(
             {
                 "imageId": image_id,
@@ -109,12 +110,20 @@ async def test_configure_context_adds_ec2_metadata_props(aresponses):
 # Test helper methods
 
 
-def configure_response(aresponses, json):
+def configure_response(aresponses, token, json):
+    aresponses.add(
+        "169.254.169.254",
+        "/latest/api/token",
+        "put",
+        aresponses.Response(text=token, headers={"X-aws-ec2-metadata-token-ttl-seconds": "21600"}),
+    )
     aresponses.add(
         "169.254.169.254",
         "/latest/dynamic/instance-identity/document",
         "get",
         # the ec2-metdata endpoint does not actually set the correct
         # content-type header, it will instead use text/plain
-        aresponses.Response(text=json, content_type="text/plain"),
+        aresponses.Response(text=json,
+                            content_type="text/plain",
+                            headers={"X-aws-ec2-metadata-token": token})
     )