feat(mmds): Accept EC2 IMDS-compatible custom headers

The custom headers supported by MMDS (X-metadata-token and
X-metadata-token-ttl-seconds) are different from what EC2 IMDS supports
(X-aws-ec2-metadata-token and X-aws-ec2-metadata-token-ttl-seconds).
Supporting EC2 IMDS-compatible custom headers in MMDS, users become able
to make their workloads work with MMDS out of the box.

Signed-off-by: Takahiro Itazuri <[email protected]>

Author: zulinx86

CHANGELOG.md

@@ -20,6 +20,10 @@ and this project adheres to
   taking diff snapshots even if dirty page tracking is disabled, by using
   `mincore(2)` to overapproximate the set of dirty pages. Only works if swap is
   disabled.
+- [#5290](https://github.com/firecracker-microvm/firecracker/pull/5290):
+  Extended MMDS to support the EC2 IMDS-compatible session token headers (i.e.
+  "X-aws-ec2-metadata-token" and "X-aws-ec2-metadata-token-ttl-seconds")
+  alongside the MMDS-specific ones.
 
 ### Changed
 

docs/mmds/mmds-user-guide.md

@@ -252,9 +252,9 @@ token. In order to be successful, the request must respect the following
 constraints:
 
 - must be directed towards `/latest/api/token` path
-- must contain a `X-metadata-token-ttl-seconds` header specifying the token
-  lifetime in seconds. The value cannot be lower than 1 or greater than 21600 (6
-  hours).
+- must contain a `X-metadata-token-ttl-seconds` or
+  `X-aws-ec2-metadata-token-ttl-seconds` header specifying the token lifetime in
+  seconds. The value cannot be lower than 1 or greater than 21600 (6 hours).
 - must not contain a `X-Forwarded-For` header.
 
 ```bash
@@ -266,8 +266,8 @@ TOKEN=`curl -X PUT "http://${MMDS_IPV4_ADDR}/latest/api/token" \
 The HTTP response from MMDS is a plaintext containing the session token.
 
 During the duration specified by the token's time to live value, all subsequent
-`GET` requests must specify the session token through the `X-metadata-token`
-header in order to fetch data from MMDS.
+`GET` requests must specify the session token through the `X-metadata-token` or
+`X-aws-ec2-metadata-token` header in order to fetch data from MMDS.
 
 ```bash
 MMDS_IPV4_ADDR=169.254.170.2

src/vmm/src/mmds/mod.rs

@@ -21,8 +21,8 @@ use serde_json::{Map, Value};
 use crate::mmds::data_store::{Mmds, MmdsDatastoreError as MmdsError, MmdsVersion, OutputFormat};
 use crate::mmds::token::PATH_TO_TOKEN;
 use crate::mmds::token_headers::{
-    REJECTED_HEADER, X_METADATA_TOKEN_HEADER, X_METADATA_TOKEN_TTL_SECONDS_HEADER,
-    get_header_value_pair,
+    REJECTED_HEADER, X_AWS_EC2_METADATA_TOKEN_HEADER, X_AWS_EC2_METADATA_TOKEN_SSL_SECONDS_HEADER,
+    X_METADATA_TOKEN_HEADER, X_METADATA_TOKEN_TTL_SECONDS_HEADER, get_header_value_pair,
 };
 
 #[rustfmt::skip]
@@ -35,9 +35,9 @@ pub enum VmmMmdsError {
     InvalidURI,
     /// Not allowed HTTP method.
     MethodNotAllowed,
-    /// No MMDS token provided. Use `X-metadata-token` header to specify the session token.
+    /// No MMDS token provided. Use `X-metadata-token` or `X-aws-ec2-metadata-token` header to specify the session token.
     NoTokenProvided,
-    /// Token time to live value not found. Use `X-metadata-token-ttl-seconds` header to specify the token's lifetime.
+    /// Token time to live value not found. Use `X-metadata-token-ttl-seconds` or `X-aws-ec2-metadata-token-ttl-seconds` header to specify the token's lifetime.
     NoTtlProvided,
     /// Resource not found: {0}.
     ResourceNotFound(String),
@@ -164,19 +164,21 @@ fn respond_to_request_mmdsv2(mmds: &mut Mmds, request: Request) -> Response {
 
 fn respond_to_get_request_checked(mmds: &Mmds, request: Request) -> Response {
     // Check whether a token exists.
-    let token =
-        match get_header_value_pair(request.headers.custom_entries(), X_METADATA_TOKEN_HEADER) {
-            Some((_, token)) => token,
-            None => {
-                let error_msg = VmmMmdsError::NoTokenProvided.to_string();
-                return build_response(
-                    request.http_version(),
-                    StatusCode::Unauthorized,
-                    MediaType::PlainText,
-                    Body::new(error_msg),
-                );
-            }
-        };
+    let token = match get_header_value_pair(
+        request.headers.custom_entries(),
+        &[X_METADATA_TOKEN_HEADER, X_AWS_EC2_METADATA_TOKEN_HEADER],
+    ) {
+        Some((_, token)) => token,
+        None => {
+            let error_msg = VmmMmdsError::NoTokenProvided.to_string();
+            return build_response(
+                request.http_version(),
+                StatusCode::Unauthorized,
+                MediaType::PlainText,
+                Body::new(error_msg),
+            );
+        }
+    };
 
     // Validate the token.
     match mmds.is_valid_token(token) {
@@ -267,36 +269,41 @@ fn respond_to_put_request(mmds: &mut Mmds, request: Request) -> Response {
     }
 
     // Get token lifetime value.
-    let ttl_seconds =
-        match get_header_value_pair(custom_headers, X_METADATA_TOKEN_TTL_SECONDS_HEADER) {
-            // Header found
-            Some((k, v)) => match v.parse::<u32>() {
-                Ok(ttl_seconds) => ttl_seconds,
-                Err(_) => {
-                    return build_response(
-                        request.http_version(),
-                        StatusCode::BadRequest,
-                        MediaType::PlainText,
-                        Body::new(
-                            RequestError::HeaderError(HttpHeaderError::InvalidValue(
-                                k.into(),
-                                v.into(),
-                            ))
-                            .to_string(),
-                        ),
-                    );
-                }
-            },
-            // Header not found
-            None => {
+    let ttl_seconds = match get_header_value_pair(
+        custom_headers,
+        &[
+            X_METADATA_TOKEN_TTL_SECONDS_HEADER,
+            X_AWS_EC2_METADATA_TOKEN_SSL_SECONDS_HEADER,
+        ],
+    ) {
+        // Header found
+        Some((k, v)) => match v.parse::<u32>() {
+            Ok(ttl_seconds) => ttl_seconds,
+            Err(_) => {
                 return build_response(
                     request.http_version(),
                     StatusCode::BadRequest,
                     MediaType::PlainText,
-                    Body::new(VmmMmdsError::NoTtlProvided.to_string()),
+                    Body::new(
+                        RequestError::HeaderError(HttpHeaderError::InvalidValue(
+                            k.into(),
+                            v.into(),
+                        ))
+                        .to_string(),
+                    ),
                 );
             }
-        };
+        },
+        // Header not found
+        None => {
+            return build_response(
+                request.http_version(),
+                StatusCode::BadRequest,
+                MediaType::PlainText,
+                Body::new(VmmMmdsError::NoTtlProvided.to_string()),
+            );
+        }
+    };
 
     // Generate token.
     let result = mmds.generate_token(ttl_seconds);
@@ -854,13 +861,14 @@ mod tests {
 
         assert_eq!(
             VmmMmdsError::NoTokenProvided.to_string(),
-            "No MMDS token provided. Use `X-metadata-token` header to specify the session token."
+            "No MMDS token provided. Use `X-metadata-token` or `X-aws-ec2-metadata-token` header \
+             to specify the session token."
         );
 
         assert_eq!(
             VmmMmdsError::NoTtlProvided.to_string(),
-            "Token time to live value not found. Use `X-metadata-token-ttl-seconds` header to \
-             specify the token's lifetime."
+            "Token time to live value not found. Use `X-metadata-token-ttl-seconds` or \
+             `X-aws-ec2-metadata-token-ttl-seconds` header to specify the token's lifetime."
         );
 
         assert_eq!(

src/vmm/src/mmds/token_headers.rs

@@ -8,16 +8,21 @@ pub const REJECTED_HEADER: &str = "X-Forwarded-For";
 
 // `X-metadata-token`
 pub(crate) const X_METADATA_TOKEN_HEADER: &str = "x-metadata-token";
+// `X-aws-ec2-metadata-token`
+pub(crate) const X_AWS_EC2_METADATA_TOKEN_HEADER: &str = "x-aws-ec2-metadata-token";
 // `X-metadata-token-ttl-seconds`
 pub(crate) const X_METADATA_TOKEN_TTL_SECONDS_HEADER: &str = "x-metadata-token-ttl-seconds";
+// `X-aws-ec2-metadata-token-ttl-seconds`
+pub(crate) const X_AWS_EC2_METADATA_TOKEN_SSL_SECONDS_HEADER: &str =
+    "x-aws-ec2-metadata-token-ttl-seconds";
 
 pub(crate) fn get_header_value_pair<'a>(
     custom_headers: &'a HashMap<String, String>,
-    header: &'static str,
+    headers: &'a [&'static str],
 ) -> Option<(&'a String, &'a String)> {
     custom_headers
         .iter()
-        .find(|(k, _)| k.eq_ignore_ascii_case(header))
+        .find(|(k, _)| headers.iter().any(|header| k.eq_ignore_ascii_case(header)))
 }
 
 #[cfg(test)]
@@ -39,38 +44,42 @@ mod tests {
 
     #[test]
     fn test_get_header_value_pair() {
+        let headers = [X_METADATA_TOKEN_HEADER, X_AWS_EC2_METADATA_TOKEN_HEADER];
+
         // No custom headers
         let custom_headers = HashMap::default();
-        let token = get_header_value_pair(&custom_headers, X_METADATA_TOKEN_HEADER);
+        let token = get_header_value_pair(&custom_headers, &headers);
         assert!(token.is_none());
 
         // Unrelated custom headers
         let custom_headers = HashMap::from([
             ("Some-Header".into(), "10".into()),
             ("Another-Header".into(), "value".into()),
         ]);
-        let token = get_header_value_pair(&custom_headers, X_METADATA_TOKEN_HEADER);
+        let token = get_header_value_pair(&custom_headers, &headers);
         assert!(token.is_none());
 
-        // Valid header
-        let expected = "THIS_IS_TOKEN";
-        let custom_headers = HashMap::from([(X_METADATA_TOKEN_HEADER.into(), expected.into())]);
-        let token = get_header_value_pair(&custom_headers, X_METADATA_TOKEN_HEADER).unwrap();
-        assert_eq!(token, (&X_METADATA_TOKEN_HEADER.into(), &expected.into()));
+        for header in headers {
+            // Valid header
+            let expected = "THIS_IS_TOKEN";
+            let custom_headers = HashMap::from([(header.into(), expected.into())]);
+            let token = get_header_value_pair(&custom_headers, &headers).unwrap();
+            assert_eq!(token, (&header.into(), &expected.into()));
 
-        // Valid header in unrelated custom headers
-        let custom_headers = HashMap::from([
-            ("Some-Header".into(), "10".into()),
-            ("Another-Header".into(), "value".into()),
-            (X_METADATA_TOKEN_HEADER.into(), expected.into()),
-        ]);
-        let token = get_header_value_pair(&custom_headers, X_METADATA_TOKEN_HEADER).unwrap();
-        assert_eq!(token, (&X_METADATA_TOKEN_HEADER.into(), &expected.into()));
+            // Valid header in unrelated custom headers
+            let custom_headers = HashMap::from([
+                ("Some-Header".into(), "10".into()),
+                ("Another-Header".into(), "value".into()),
+                (header.into(), expected.into()),
+            ]);
+            let token = get_header_value_pair(&custom_headers, &headers).unwrap();
+            assert_eq!(token, (&header.into(), &expected.into()));
 
-        // Test case-insensitiveness
-        let header = to_mixed_case(X_METADATA_TOKEN_HEADER);
-        let custom_headers = HashMap::from([(header.clone(), expected.into())]);
-        let token = get_header_value_pair(&custom_headers, X_METADATA_TOKEN_HEADER).unwrap();
-        assert_eq!(token, (&header, &expected.into()));
+            // Test case-insensitiveness
+            let header = to_mixed_case(header);
+            let custom_headers = HashMap::from([(header.clone(), expected.into())]);
+            let token = get_header_value_pair(&custom_headers, &headers).unwrap();
+            assert_eq!(token, (&header, &expected.into()));
+        }
     }
 }

tests/framework/utils.py

@@ -397,25 +397,35 @@ def get_kernel_version(level=2):
     return linux_version
 
 
-def generate_mmds_session_token(ssh_connection, ipv4_address, token_ttl):
+def generate_mmds_session_token(
+    ssh_connection, ipv4_address, token_ttl, imds_compat=False
+):
     """Generate session token used for MMDS V2 requests."""
     cmd = "curl -m 2 -s"
     cmd += " -X PUT"
-    cmd += ' -H  "X-metadata-token-ttl-seconds: {}"'.format(token_ttl)
+    if imds_compat:
+        cmd += ' -H "X-aws-ec2-metadata-token-ttl-seconds: {}"'.format(token_ttl)
+    else:
+        cmd += ' -H "X-metadata-token-ttl-seconds: {}"'.format(token_ttl)
     cmd += " http://{}/latest/api/token".format(ipv4_address)
     _, stdout, _ = ssh_connection.run(cmd)
     token = stdout
 
     return token
 
 
-def generate_mmds_get_request(ipv4_address, token=None, app_json=True):
+def generate_mmds_get_request(
+    ipv4_address, token=None, app_json=True, imds_compat=False
+):
     """Build `GET` request to fetch metadata from MMDS."""
     cmd = "curl -m 2 -s"
 
     if token is not None:
         cmd += " -X GET"
-        cmd += ' -H  "X-metadata-token: {}"'.format(token)
+        if imds_compat:
+            cmd += ' -H "X-aws-ec2-metadata-token: {}"'.format(token)
+        else:
+            cmd += ' -H "X-metadata-token: {}"'.format(token)
 
     if app_json:
         cmd += ' -H "Accept: application/json"'

tests/integration_tests/functional/test_mmds.py

@@ -128,6 +128,35 @@ def _validate_mmds_snapshot(
     run_guest_cmd(ssh_connection, cmd, data_store, use_json=True)
 
 
+@pytest.mark.parametrize("version", MMDS_VERSIONS)
+@pytest.mark.parametrize("imds_compat", [True, False])
+def test_token_generation(uvm_plain, version, imds_compat):
+    """
+    Test MMDS token generation.
+    """
+    test_microvm = uvm_plain
+    test_microvm.spawn()
+
+    test_microvm.add_net_iface()
+    configure_mmds(test_microvm, iface_ids=["eth0"], version=version)
+    populate_data_store(test_microvm, {"foo": "bar"})
+
+    test_microvm.basic_config(vcpu_count=1)
+    test_microvm.start()
+    ssh_connection = test_microvm.ssh
+
+    cmd = "ip route add {} dev eth0".format(DEFAULT_IPV4)
+    run_guest_cmd(ssh_connection, cmd, "")
+
+    token = generate_mmds_session_token(ssh_connection, DEFAULT_IPV4, 60, imds_compat)
+    if version == "V1":
+        assert token == "Not allowed HTTP method."
+        # V1 accepts GET request even with an invalid token. So keep going.
+
+    cmd = generate_mmds_get_request(DEFAULT_IPV4, token, False, imds_compat) + "foo"
+    run_guest_cmd(ssh_connection, cmd, "bar")
+
+
 @pytest.mark.parametrize("version", MMDS_VERSIONS)
 def test_custom_ipv4(uvm_plain, version):
     """
@@ -647,7 +676,7 @@ def test_mmds_v2_negative(uvm_plain):
     # Check `GET` request fails when token is not provided.
     cmd = generate_mmds_get_request(DEFAULT_IPV4)
     expected = (
-        "No MMDS token provided. Use `X-metadata-token` header "
+        "No MMDS token provided. Use `X-metadata-token` or `X-aws-ec2-metadata-token` header "
         "to specify the session token."
     )
     run_guest_cmd(ssh_connection, cmd, expected)
@@ -664,9 +693,8 @@ def test_mmds_v2_negative(uvm_plain):
     # Check `PUT` request fails when token TTL is not provided.
     cmd = f"curl -m 2 -s -X PUT http://{DEFAULT_IPV4}/latest/api/token"
     expected = (
-        "Token time to live value not found. Use "
-        "`X-metadata-token-ttl-seconds` header to specify "
-        "the token's lifetime."
+        "Token time to live value not found. Use `X-metadata-token-ttl-seconds` or "
+        "`X-aws-ec2-metadata-token-ttl-seconds` header to specify the token's lifetime."
     )
     run_guest_cmd(ssh_connection, cmd, expected)