fix(mmds): Set Content-Type header correctly

The MMDS response "Content-Type" header was always "application/json"
regardless of the request "Accept" header, although MMDS responds
different contents depending on the "Accept" header. Without "Accept"
header or with empty "Accept:" or "Accept: */*", micro-http defaults
to "Accept: text/plain". In such a case, MMDS behaves like IMDS and
the response content is plain text, so "Content-Type" should be
"text/plain".

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

Author: zulinx86

CHANGELOG.md

@@ -31,6 +31,9 @@ and this project adheres to
 
 - [#5222](https://github.com/firecracker-microvm/firecracker/pull/5222): Fixed
   network and rng devices locking up on hosts with non 4K pages.
+- [#5226](https://github.com/firecracker-microvm/firecracker/pull/5226): Fixed
+  MMDS to set `Content-Type` header correctly (i.e. `Content-Type: text/plain`
+  for IMDS format and `Content-Type: application/json` for JSON format).
 
 ## [1.12.0]
 

src/vmm/src/mmds/mod.rs

@@ -51,8 +51,16 @@ impl From<MediaType> for OutputFormat {
 }
 
 // Builds the `micro_http::Response` with a given HTTP version, status code, and body.
-fn build_response(http_version: Version, status_code: StatusCode, body: Body) -> Response {
+fn build_response(
+    http_version: Version,
+    status_code: StatusCode,
+    content_type: Option<MediaType>,
+    body: Body,
+) -> Response {
     let mut response = Response::new(http_version, status_code);
+    if let Some(content_type) = content_type {
+        response.set_content_type(content_type);
+    }
     response.set_body(body);
     response
 }
@@ -105,6 +113,7 @@ pub fn convert_to_response(mmds: Arc<Mutex<Mmds>>, request: Request) -> Response
         return build_response(
             request.http_version(),
             StatusCode::BadRequest,
+            None,
             Body::new(VmmMmdsError::InvalidURI.to_string()),
         );
     }
@@ -125,6 +134,7 @@ fn respond_to_request_mmdsv1(mmds: &Mmds, request: Request) -> Response {
             let mut response = build_response(
                 request.http_version(),
                 StatusCode::MethodNotAllowed,
+                None,
                 Body::new(VmmMmdsError::MethodNotAllowed.to_string()),
             );
             response.allow_method(Method::Get);
@@ -141,6 +151,7 @@ fn respond_to_request_mmdsv2(mmds: &mut Mmds, request: Request) -> Response {
             return build_response(
                 request.http_version(),
                 StatusCode::BadRequest,
+                None,
                 Body::new(err.to_string()),
             );
         }
@@ -154,6 +165,7 @@ fn respond_to_request_mmdsv2(mmds: &mut Mmds, request: Request) -> Response {
             let mut response = build_response(
                 request.http_version(),
                 StatusCode::MethodNotAllowed,
+                None,
                 Body::new(VmmMmdsError::MethodNotAllowed.to_string()),
             );
             response.allow_method(Method::Get);
@@ -176,6 +188,7 @@ fn respond_to_get_request_checked(
             return build_response(
                 request.http_version(),
                 StatusCode::Unauthorized,
+                None,
                 Body::new(error_msg),
             );
         }
@@ -187,6 +200,7 @@ fn respond_to_get_request_checked(
         Ok(false) => build_response(
             request.http_version(),
             StatusCode::Unauthorized,
+            None,
             Body::new(VmmMmdsError::InvalidToken.to_string()),
         ),
         Err(_) => unreachable!(),
@@ -200,10 +214,13 @@ fn respond_to_get_request_unchecked(mmds: &Mmds, request: Request) -> Response {
     // sanitize the URI.
     let json_path = sanitize_uri(uri.to_string());
 
-    match mmds.get_value(json_path, request.headers.accept().into()) {
+    let content_type = request.headers.accept();
+
+    match mmds.get_value(json_path, content_type.into()) {
         Ok(response_body) => build_response(
             request.http_version(),
             StatusCode::OK,
+            Some(content_type),
             Body::new(response_body),
         ),
         Err(err) => match err {
@@ -212,17 +229,20 @@ fn respond_to_get_request_unchecked(mmds: &Mmds, request: Request) -> Response {
                 build_response(
                     request.http_version(),
                     StatusCode::NotFound,
+                    None,
                     Body::new(error_msg),
                 )
             }
             MmdsError::UnsupportedValueType => build_response(
                 request.http_version(),
                 StatusCode::NotImplemented,
+                None,
                 Body::new(err.to_string()),
             ),
             MmdsError::DataStoreLimitExceeded => build_response(
                 request.http_version(),
                 StatusCode::PayloadTooLarge,
+                None,
                 Body::new(err.to_string()),
             ),
             _ => unreachable!(),
@@ -248,6 +268,7 @@ fn respond_to_put_request(
         return build_response(
             request.http_version(),
             StatusCode::BadRequest,
+            None,
             Body::new(error_msg),
         );
     }
@@ -262,6 +283,7 @@ fn respond_to_put_request(
         return build_response(
             request.http_version(),
             StatusCode::NotFound,
+            None,
             Body::new(error_msg),
         );
     }
@@ -273,6 +295,7 @@ fn respond_to_put_request(
             return build_response(
                 request.http_version(),
                 StatusCode::BadRequest,
+                None,
                 Body::new(VmmMmdsError::NoTtlProvided.to_string()),
             );
         }
@@ -281,15 +304,16 @@ fn respond_to_put_request(
     // Generate token.
     let result = mmds.generate_token(ttl_seconds);
     match result {
-        Ok(token) => {
-            let mut response =
-                build_response(request.http_version(), StatusCode::OK, Body::new(token));
-            response.set_content_type(MediaType::PlainText);
-            response
-        }
+        Ok(token) => build_response(
+            request.http_version(),
+            StatusCode::OK,
+            Some(MediaType::PlainText),
+            Body::new(token),
+        ),
         Err(err) => build_response(
             request.http_version(),
             StatusCode::BadRequest,
+            None,
             Body::new(err.to_string()),
         ),
     }
@@ -343,6 +367,31 @@ mod tests {
         }"#
     }
 
+    fn get_plain_text_data() -> &'static str {
+        "age\nname/\nphones/"
+    }
+
+    fn generate_request_and_expected_response(
+        request_bytes: &[u8],
+        media_type: MediaType,
+    ) -> (Request, Response) {
+        let request = Request::try_from(request_bytes, None).unwrap();
+
+        let mut response = Response::new(Version::Http10, StatusCode::OK);
+        response.set_content_type(media_type);
+        let body = match media_type {
+            MediaType::ApplicationJson => {
+                let mut body = get_json_data().to_string();
+                body.retain(|c| !c.is_whitespace());
+                body
+            }
+            MediaType::PlainText => get_plain_text_data().to_string(),
+        };
+        response.set_body(Body::new(body));
+
+        (request, response)
+    }
+
     #[test]
     fn test_sanitize_uri() {
         let sanitized = "/a/b/c/d";
@@ -362,6 +411,66 @@ mod tests {
         assert_eq!(sanitize_uri("//aa//bb///cc//d".to_owned()), "/aa/bb/cc/d");
     }
 
+    #[test]
+    fn test_request_accept_header() {
+        // This test validates the response `Content-Type` header and the response content for
+        // various request `Accept` headers.
+
+        // Populate MMDS with data.
+        let mmds = populate_mmds();
+
+        // Test without `Accept` header. micro-http defaults to `Accept: text/plain`.
+        let (request, expected_response) = generate_request_and_expected_response(
+            b"GET http://169.254.169.254/ HTTP/1.0\r\n\r\n",
+            MediaType::PlainText,
+        );
+        assert_eq!(
+            convert_to_response(mmds.clone(), request),
+            expected_response
+        );
+
+        // Test with empty `Accept` header. micro-http defaults to `Accept: text/plain`.
+        let (request, expected_response) = generate_request_and_expected_response(
+            b"GET http://169.254.169.254/ HTTP/1.0\r\n\"
+            Accept:\r\n\r\n",
+            MediaType::PlainText,
+        );
+        assert_eq!(
+            convert_to_response(mmds.clone(), request),
+            expected_response
+        );
+
+        // Test with `Accept: */*` header.
+        let (request, expected_response) = generate_request_and_expected_response(
+            b"GET http://169.254.169.254/ HTTP/1.0\r\n\"
+            Accept: */*\r\n\r\n",
+            MediaType::PlainText,
+        );
+        assert_eq!(
+            convert_to_response(mmds.clone(), request),
+            expected_response
+        );
+
+        // Test with `Accept: text/plain`.
+        let (request, expected_response) = generate_request_and_expected_response(
+            b"GET http://169.254.169.254/ HTTP/1.0\r\n\
+            Accept: plain/text\r\n\r\n",
+            MediaType::PlainText,
+        );
+        assert_eq!(
+            convert_to_response(mmds.clone(), request),
+            expected_response
+        );
+
+        // Test with `Accept: application/json`.
+        let (request, expected_response) = generate_request_and_expected_response(
+            b"GET http://169.254.169.254/ HTTP/1.0\r\n\
+            Accept: application/json\r\n\r\n",
+            MediaType::ApplicationJson,
+        );
+        assert_eq!(convert_to_response(mmds, request), expected_response);
+    }
+
     #[test]
     fn test_respond_to_request_mmdsv1() {
         // Populate MMDS with data.