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, users are able to make
their workloads work with MMDS out of the box.

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

Author: zulinx86

CHANGELOG.md

@@ -16,6 +16,10 @@ and this project adheres to
 - [#5175](https://github.com/firecracker-microvm/firecracker/pull/5175): Allow
   including a custom cpu template directly in the json configuration file passed
   to `--config-file` under the `cpu_config` key.
+- [#XXXX](https://github.com/firecracker-microvm/firecracker/pull/XXXX):
+  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
 

src/vmm/src/mmds/token_headers.rs

@@ -16,13 +16,17 @@ pub struct XMetadataToken(pub Option<String>);
 
 // Defined in lowercase since HTTP headers are case-insensitive.
 const X_METADATA_TOKEN_HEADER: &str = "x-metadata-token";
+const X_AWS_EC2_METADATA_TOKEN_HEADER: &str = "x-aws-ec2-metadata-token";
 
 impl From<&HashMap<String, String>> for XMetadataToken {
     fn from(custom_headers: &HashMap<String, String>) -> Self {
         Self(
             custom_headers
                 .iter()
-                .find(|(k, _)| k.to_lowercase() == X_METADATA_TOKEN_HEADER)
+                .find(|(k, _)| {
+                    let k = k.to_lowercase();
+                    k == X_METADATA_TOKEN_HEADER || k == X_AWS_EC2_METADATA_TOKEN_HEADER
+                })
                 .map(|(_, v)| v.to_string()),
         )
     }
@@ -35,14 +39,19 @@ pub struct XMetadataTokenTtlSeconds(pub Option<u32>);
 
 // Defined in lowercase since HTTP headers are case-insensitive.
 const X_METADATA_TOKEN_TTL_SECONDS_HEADER: &str = "x-metadata-token-ttl-seconds";
+const X_AWS_EC2_METADATA_TOKEN_SSL_SECONDS_HEADER: &str = "x-aws-ec2-metadata-token-ttl-seconds";
 
 impl TryFrom<&HashMap<String, String>> for XMetadataTokenTtlSeconds {
     type Error = RequestError;
 
     fn try_from(custom_headers: &HashMap<String, String>) -> Result<Self, RequestError> {
         let seconds = custom_headers
             .iter()
-            .find(|(k, _)| k.to_lowercase() == X_METADATA_TOKEN_TTL_SECONDS_HEADER)
+            .find(|(k, _)| {
+                let k = k.to_lowercase();
+                k == X_METADATA_TOKEN_TTL_SECONDS_HEADER
+                    || k == X_AWS_EC2_METADATA_TOKEN_SSL_SECONDS_HEADER
+            })
             .map(|(k, v)| {
                 v.parse::<u32>().map_err(|_| {
                     RequestError::HeaderError(HttpHeaderError::InvalidValue(
@@ -89,26 +98,28 @@ mod tests {
         let x_metadata_token = XMetadataToken::from(&custom_headers);
         assert!(x_metadata_token.0.is_none());
 
-        // Valid header
-        let token = "THIS_IS_TOKEN";
-        let custom_headers = HashMap::from([(X_METADATA_TOKEN_HEADER.into(), token.into())]);
-        let x_metadata_token = XMetadataToken::from(&custom_headers);
-        assert_eq!(&x_metadata_token.0.unwrap(), token);
-
-        // 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(), token.into()),
-        ]);
-        let x_metadata_token = XMetadataToken::from(&custom_headers);
-        assert_eq!(&x_metadata_token.0.unwrap(), token);
-
-        // Test case-insensitiveness
-        let custom_headers =
-            HashMap::from([(to_mixed_case(X_METADATA_TOKEN_HEADER), token.into())]);
-        let x_metadata_token = XMetadataToken::from(&custom_headers);
-        assert_eq!(&x_metadata_token.0.unwrap(), token);
+        for header in [X_METADATA_TOKEN_HEADER, X_AWS_EC2_METADATA_TOKEN_HEADER] {
+            let token = "THIS_IS_TOKEN";
+
+            // Valid header
+            let custom_headers = HashMap::from([(header.into(), token.into())]);
+            let x_metadata_token = XMetadataToken::from(&custom_headers);
+            assert_eq!(&x_metadata_token.0.unwrap(), token);
+
+            // Valid header in unrelated custom headers
+            let custom_headers = HashMap::from([
+                ("Some-Header".into(), "10".into()),
+                ("Another-Header".into(), "value".into()),
+                (header.into(), token.into()),
+            ]);
+            let x_metadata_token = XMetadataToken::from(&custom_headers);
+            assert_eq!(&x_metadata_token.0.unwrap(), token);
+
+            // Test case-insensitiveness
+            let custom_headers = HashMap::from([(to_mixed_case(header), token.into())]);
+            let x_metadata_token = XMetadataToken::from(&custom_headers);
+            assert_eq!(&x_metadata_token.0.unwrap(), token);
+        }
     }
 
     #[test]
@@ -128,48 +139,46 @@ mod tests {
             XMetadataTokenTtlSeconds::try_from(&custom_headers).unwrap();
         assert!(x_metadata_token_ttl_seconds.0.is_none());
 
-        // Valid header
-        let seconds = 60;
-        let custom_headers = HashMap::from([(
-            X_METADATA_TOKEN_TTL_SECONDS_HEADER.into(),
-            seconds.to_string(),
-        )]);
-        let x_metadata_token_ttl_seconds =
-            XMetadataTokenTtlSeconds::try_from(&custom_headers).unwrap();
-        assert_eq!(x_metadata_token_ttl_seconds.0.unwrap(), seconds);
-
-        // Valid header in unrelated custom headers
-        let custom_headers = HashMap::from([
-            ("Some-Header".into(), "10".into()),
-            ("Another-Header".into(), "value".into()),
-            (
-                X_METADATA_TOKEN_TTL_SECONDS_HEADER.into(),
-                seconds.to_string(),
-            ),
-        ]);
-        let x_metadata_token_ttl_seconds =
-            XMetadataTokenTtlSeconds::try_from(&custom_headers).unwrap();
-        assert_eq!(x_metadata_token_ttl_seconds.0.unwrap(), seconds);
-
-        // Test case-insensitiveness
-        let custom_headers = HashMap::from([(
-            to_mixed_case(X_METADATA_TOKEN_TTL_SECONDS_HEADER),
-            seconds.to_string(),
-        )]);
-        let x_metadata_token_ttl_seconds =
-            XMetadataTokenTtlSeconds::try_from(&custom_headers).unwrap();
-        assert_eq!(x_metadata_token_ttl_seconds.0.unwrap(), seconds);
-
-        // Invalid value
-        let header_name = "X-metadata-token-ttl-seconds";
-        let invalid_seconds = "-60";
-        let custom_headers = HashMap::from([(header_name.into(), invalid_seconds.to_string())]);
-        assert_eq!(
-            XMetadataTokenTtlSeconds::try_from(&custom_headers).unwrap_err(),
-            RequestError::HeaderError(HttpHeaderError::InvalidValue(
-                header_name.into(),
-                invalid_seconds.to_string()
-            ))
-        );
+        for header in [
+            X_METADATA_TOKEN_TTL_SECONDS_HEADER,
+            X_AWS_EC2_METADATA_TOKEN_SSL_SECONDS_HEADER,
+        ] {
+            let seconds = 60;
+
+            // Valid header
+            let custom_headers = HashMap::from([(header.into(), seconds.to_string())]);
+            let x_metadata_token_ttl_seconds =
+                XMetadataTokenTtlSeconds::try_from(&custom_headers).unwrap();
+            assert_eq!(x_metadata_token_ttl_seconds.0.unwrap(), seconds);
+
+            // Valid header in unrelated custom headers
+            let custom_headers = HashMap::from([
+                ("Some-Header".into(), "10".into()),
+                ("Another-Header".into(), "value".into()),
+                (header.into(), seconds.to_string()),
+            ]);
+            let x_metadata_token_ttl_seconds =
+                XMetadataTokenTtlSeconds::try_from(&custom_headers).unwrap();
+            assert_eq!(x_metadata_token_ttl_seconds.0.unwrap(), seconds);
+
+            // Test case-insensitiveness
+            let custom_headers = HashMap::from([(to_mixed_case(header), seconds.to_string())]);
+            let x_metadata_token_ttl_seconds =
+                XMetadataTokenTtlSeconds::try_from(&custom_headers).unwrap();
+            assert_eq!(x_metadata_token_ttl_seconds.0.unwrap(), seconds);
+
+            // Invalid value
+            let mixed_case_header = to_mixed_case(header);
+            let invalid_seconds = "-60";
+            let custom_headers =
+                HashMap::from([(mixed_case_header.clone(), invalid_seconds.to_string())]);
+            assert_eq!(
+                XMetadataTokenTtlSeconds::try_from(&custom_headers).unwrap_err(),
+                RequestError::HeaderError(HttpHeaderError::InvalidValue(
+                    mixed_case_header,
+                    invalid_seconds.to_string()
+                ))
+            );
+        }
     }
 }

tests/integration_tests/functional/test_mmds.py

@@ -647,7 +647,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 +664,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)