@@ -51,8 +51,14 @@ impl From<MediaType> for OutputFormat {
5151} 
5252
5353// Builds the `micro_http::Response` with a given HTTP version, status code, and body. 
54- fn  build_response ( http_version :  Version ,  status_code :  StatusCode ,  body :  Body )  -> Response  { 
54+ fn  build_response ( 
55+     http_version :  Version , 
56+     status_code :  StatusCode , 
57+     content_type :  MediaType , 
58+     body :  Body , 
59+ )  -> Response  { 
5560    let  mut  response = Response :: new ( http_version,  status_code) ; 
61+     response. set_content_type ( content_type) ; 
5662    response. set_body ( body) ; 
5763    response
5864} 
@@ -105,6 +111,7 @@ pub fn convert_to_response(mmds: Arc<Mutex<Mmds>>, request: Request) -> Response
105111        return  build_response ( 
106112            request. http_version ( ) , 
107113            StatusCode :: BadRequest , 
114+             MediaType :: PlainText , 
108115            Body :: new ( VmmMmdsError :: InvalidURI . to_string ( ) ) , 
109116        ) ; 
110117    } 
@@ -125,6 +132,7 @@ fn respond_to_request_mmdsv1(mmds: &Mmds, request: Request) -> Response {
125132            let  mut  response = build_response ( 
126133                request. http_version ( ) , 
127134                StatusCode :: MethodNotAllowed , 
135+                 MediaType :: PlainText , 
128136                Body :: new ( VmmMmdsError :: MethodNotAllowed . to_string ( ) ) , 
129137            ) ; 
130138            response. allow_method ( Method :: Get ) ; 
@@ -141,6 +149,7 @@ fn respond_to_request_mmdsv2(mmds: &mut Mmds, request: Request) -> Response {
141149            return  build_response ( 
142150                request. http_version ( ) , 
143151                StatusCode :: BadRequest , 
152+                 MediaType :: PlainText , 
144153                Body :: new ( err. to_string ( ) ) , 
145154            ) ; 
146155        } 
@@ -154,6 +163,7 @@ fn respond_to_request_mmdsv2(mmds: &mut Mmds, request: Request) -> Response {
154163            let  mut  response = build_response ( 
155164                request. http_version ( ) , 
156165                StatusCode :: MethodNotAllowed , 
166+                 MediaType :: PlainText , 
157167                Body :: new ( VmmMmdsError :: MethodNotAllowed . to_string ( ) ) , 
158168            ) ; 
159169            response. allow_method ( Method :: Get ) ; 
@@ -176,6 +186,7 @@ fn respond_to_get_request_checked(
176186            return  build_response ( 
177187                request. http_version ( ) , 
178188                StatusCode :: Unauthorized , 
189+                 MediaType :: PlainText , 
179190                Body :: new ( error_msg) , 
180191            ) ; 
181192        } 
@@ -187,6 +198,7 @@ fn respond_to_get_request_checked(
187198        Ok ( false )  => build_response ( 
188199            request. http_version ( ) , 
189200            StatusCode :: Unauthorized , 
201+             MediaType :: PlainText , 
190202            Body :: new ( VmmMmdsError :: InvalidToken . to_string ( ) ) , 
191203        ) , 
192204        Err ( _)  => unreachable ! ( ) , 
@@ -200,10 +212,13 @@ fn respond_to_get_request_unchecked(mmds: &Mmds, request: Request) -> Response {
200212    // sanitize the URI. 
201213    let  json_path = sanitize_uri ( uri. to_string ( ) ) ; 
202214
203-     match  mmds. get_value ( json_path,  request. headers . accept ( ) . into ( ) )  { 
215+     let  content_type = request. headers . accept ( ) ; 
216+ 
217+     match  mmds. get_value ( json_path,  content_type. into ( ) )  { 
204218        Ok ( response_body)  => build_response ( 
205219            request. http_version ( ) , 
206220            StatusCode :: OK , 
221+             content_type, 
207222            Body :: new ( response_body) , 
208223        ) , 
209224        Err ( err)  => match  err { 
@@ -212,17 +227,20 @@ fn respond_to_get_request_unchecked(mmds: &Mmds, request: Request) -> Response {
212227                build_response ( 
213228                    request. http_version ( ) , 
214229                    StatusCode :: NotFound , 
230+                     MediaType :: PlainText , 
215231                    Body :: new ( error_msg) , 
216232                ) 
217233            } 
218234            MmdsError :: UnsupportedValueType  => build_response ( 
219235                request. http_version ( ) , 
220236                StatusCode :: NotImplemented , 
237+                 MediaType :: PlainText , 
221238                Body :: new ( err. to_string ( ) ) , 
222239            ) , 
223240            MmdsError :: DataStoreLimitExceeded  => build_response ( 
224241                request. http_version ( ) , 
225242                StatusCode :: PayloadTooLarge , 
243+                 MediaType :: PlainText , 
226244                Body :: new ( err. to_string ( ) ) , 
227245            ) , 
228246            _ => unreachable ! ( ) , 
@@ -248,6 +266,7 @@ fn respond_to_put_request(
248266        return  build_response ( 
249267            request. http_version ( ) , 
250268            StatusCode :: BadRequest , 
269+             MediaType :: PlainText , 
251270            Body :: new ( error_msg) , 
252271        ) ; 
253272    } 
@@ -262,6 +281,7 @@ fn respond_to_put_request(
262281        return  build_response ( 
263282            request. http_version ( ) , 
264283            StatusCode :: NotFound , 
284+             MediaType :: PlainText , 
265285            Body :: new ( error_msg) , 
266286        ) ; 
267287    } 
@@ -273,6 +293,7 @@ fn respond_to_put_request(
273293            return  build_response ( 
274294                request. http_version ( ) , 
275295                StatusCode :: BadRequest , 
296+                 MediaType :: PlainText , 
276297                Body :: new ( VmmMmdsError :: NoTtlProvided . to_string ( ) ) , 
277298            ) ; 
278299        } 
@@ -281,15 +302,16 @@ fn respond_to_put_request(
281302    // Generate token. 
282303    let  result = mmds. generate_token ( ttl_seconds) ; 
283304    match  result { 
284-         Ok ( token)  => { 
285-             let   mut  response = 
286-                  build_response ( request . http_version ( ) ,   StatusCode :: OK ,   Body :: new ( token ) ) ; 
287-             response . set_content_type ( MediaType :: PlainText ) ; 
288-             response 
289-         } 
305+         Ok ( token)  => build_response ( 
306+             request . http_version ( ) , 
307+             StatusCode :: OK , 
308+             MediaType :: PlainText , 
309+             Body :: new ( token ) , 
310+         ) , 
290311        Err ( err)  => build_response ( 
291312            request. http_version ( ) , 
292313            StatusCode :: BadRequest , 
314+             MediaType :: PlainText , 
293315            Body :: new ( err. to_string ( ) ) , 
294316        ) , 
295317    } 
@@ -343,6 +365,31 @@ mod tests {
343365        }"# 
344366    } 
345367
368+     fn  get_plain_text_data ( )  -> & ' static  str  { 
369+         "age\n name/\n phones/" 
370+     } 
371+ 
372+     fn  generate_request_and_expected_response ( 
373+         request_bytes :  & [ u8 ] , 
374+         media_type :  MediaType , 
375+     )  -> ( Request ,  Response )  { 
376+         let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
377+ 
378+         let  mut  response = Response :: new ( Version :: Http10 ,  StatusCode :: OK ) ; 
379+         response. set_content_type ( media_type) ; 
380+         let  body = match  media_type { 
381+             MediaType :: ApplicationJson  => { 
382+                 let  mut  body = get_json_data ( ) . to_string ( ) ; 
383+                 body. retain ( |c| !c. is_whitespace ( ) ) ; 
384+                 body
385+             } 
386+             MediaType :: PlainText  => get_plain_text_data ( ) . to_string ( ) , 
387+         } ; 
388+         response. set_body ( Body :: new ( body) ) ; 
389+ 
390+         ( request,  response) 
391+     } 
392+ 
346393    #[ test]  
347394    fn  test_sanitize_uri ( )  { 
348395        let  sanitized = "/a/b/c/d" ; 
@@ -362,6 +409,66 @@ mod tests {
362409        assert_eq ! ( sanitize_uri( "//aa//bb///cc//d" . to_owned( ) ) ,  "/aa/bb/cc/d" ) ; 
363410    } 
364411
412+     #[ test]  
413+     fn  test_request_accept_header ( )  { 
414+         // This test validates the response `Content-Type` header and the response content for 
415+         // various request `Accept` headers. 
416+ 
417+         // Populate MMDS with data. 
418+         let  mmds = populate_mmds ( ) ; 
419+ 
420+         // Test without `Accept` header. micro-http defaults to `Accept: text/plain`. 
421+         let  ( request,  expected_response)  = generate_request_and_expected_response ( 
422+             b"GET http://169.254.169.254/ HTTP/1.0\r \n \r \n " , 
423+             MediaType :: PlainText , 
424+         ) ; 
425+         assert_eq ! ( 
426+             convert_to_response( mmds. clone( ) ,  request) , 
427+             expected_response
428+         ) ; 
429+ 
430+         // Test with empty `Accept` header. micro-http defaults to `Accept: text/plain`. 
431+         let  ( request,  expected_response)  = generate_request_and_expected_response ( 
432+             b"GET http://169.254.169.254/ HTTP/1.0\r \n \"  
433+             Accept:\r \n \r \n " , 
434+             MediaType :: PlainText , 
435+         ) ; 
436+         assert_eq ! ( 
437+             convert_to_response( mmds. clone( ) ,  request) , 
438+             expected_response
439+         ) ; 
440+ 
441+         // Test with `Accept: */*` header. 
442+         let  ( request,  expected_response)  = generate_request_and_expected_response ( 
443+             b"GET http://169.254.169.254/ HTTP/1.0\r \n \"  
444+             Accept: */*\r \n \r \n " , 
445+             MediaType :: PlainText , 
446+         ) ; 
447+         assert_eq ! ( 
448+             convert_to_response( mmds. clone( ) ,  request) , 
449+             expected_response
450+         ) ; 
451+ 
452+         // Test with `Accept: text/plain`. 
453+         let  ( request,  expected_response)  = generate_request_and_expected_response ( 
454+             b"GET http://169.254.169.254/ HTTP/1.0\r \n \  
455+ \r \n \r \n ", 
456+             MediaType :: PlainText , 
457+         ) ; 
458+         assert_eq ! ( 
459+             convert_to_response( mmds. clone( ) ,  request) , 
460+             expected_response
461+         ) ; 
462+ 
463+         // Test with `Accept: application/json`. 
464+         let  ( request,  expected_response)  = generate_request_and_expected_response ( 
465+             b"GET http://169.254.169.254/ HTTP/1.0\r \n \  
466+ \r \n \r \n ", 
467+             MediaType :: ApplicationJson , 
468+         ) ; 
469+         assert_eq ! ( convert_to_response( mmds,  request) ,  expected_response) ; 
470+     } 
471+ 
365472    #[ test]  
366473    fn  test_respond_to_request_mmdsv1 ( )  { 
367474        // Populate MMDS with data. 
@@ -381,6 +488,7 @@ mod tests {
381488        let  request_bytes = b"GET http://169.254.169.254/invalid HTTP/1.0\r \n \r \n " ; 
382489        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
383490        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: NotFound ) ; 
491+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
384492        expected_response. set_body ( Body :: new ( 
385493            VmmMmdsError :: ResourceNotFound ( String :: from ( "/invalid" ) ) . to_string ( ) , 
386494        ) ) ; 
@@ -391,6 +499,7 @@ mod tests {
391499        let  request_bytes = b"GET /age HTTP/1.1\r \n \r \n " ; 
392500        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
393501        let  mut  expected_response = Response :: new ( Version :: Http11 ,  StatusCode :: NotImplemented ) ; 
502+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
394503        let  body = "Cannot retrieve value. The value has an unsupported type." . to_string ( ) ; 
395504        expected_response. set_body ( Body :: new ( body) ) ; 
396505        let  actual_response = convert_to_response ( mmds. clone ( ) ,  request) ; 
@@ -403,6 +512,7 @@ mod tests {
403512            let  request = Request :: try_from ( request_bytes. as_bytes ( ) ,  None ) . unwrap ( ) ; 
404513            let  mut  expected_response =
405514                Response :: new ( Version :: Http10 ,  StatusCode :: MethodNotAllowed ) ; 
515+             expected_response. set_content_type ( MediaType :: PlainText ) ; 
406516            expected_response. set_body ( Body :: new ( VmmMmdsError :: MethodNotAllowed . to_string ( ) ) ) ; 
407517            expected_response. allow_method ( Method :: Get ) ; 
408518            let  actual_response = convert_to_response ( mmds. clone ( ) ,  request) ; 
@@ -413,6 +523,7 @@ mod tests {
413523        let  request_bytes = b"GET http:// HTTP/1.0\r \n \r \n " ; 
414524        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
415525        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: BadRequest ) ; 
526+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
416527        expected_response. set_body ( Body :: new ( VmmMmdsError :: InvalidURI . to_string ( ) ) ) ; 
417528        let  actual_response = convert_to_response ( mmds. clone ( ) ,  request) ; 
418529        assert_eq ! ( actual_response,  expected_response) ; 
@@ -458,6 +569,7 @@ mod tests {
458569        let  request_bytes = b"PATCH http://169.254.169.255/ HTTP/1.0\r \n \r \n " ; 
459570        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
460571        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: MethodNotAllowed ) ; 
572+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
461573        expected_response. set_body ( Body :: new ( VmmMmdsError :: MethodNotAllowed . to_string ( ) ) ) ; 
462574        expected_response. allow_method ( Method :: Get ) ; 
463575        expected_response. allow_method ( Method :: Put ) ; 
@@ -470,6 +582,7 @@ mod tests {
470582                                    X-metadata-token-ttl-seconds: application/json\r \n \r \n " ; 
471583        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
472584        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: BadRequest ) ; 
585+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
473586        expected_response. set_body ( Body :: new ( 
474587            "Invalid header. Reason: Invalid value. Key:X-metadata-token-ttl-seconds; \  
475588
@@ -484,6 +597,7 @@ mod tests {
484597\r \n \r \n "; 
485598        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
486599        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: BadRequest ) ; 
600+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
487601        expected_response. set_body ( Body :: new ( 
488602            "Invalid header. Reason: Unsupported header name. Key: X-Forwarded-For" . to_string ( ) , 
489603        ) ) ; 
@@ -495,6 +609,7 @@ mod tests {
495609\r \n \r \n "; 
496610        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
497611        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: NotFound ) ; 
612+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
498613        expected_response. set_body ( Body :: new ( 
499614            VmmMmdsError :: ResourceNotFound ( String :: from ( "/token" ) ) . to_string ( ) , 
500615        ) ) ; 
@@ -511,6 +626,7 @@ mod tests {
511626            ) ; 
512627            let  request = Request :: try_from ( request_bytes. as_bytes ( ) ,  None ) . unwrap ( ) ; 
513628            let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: BadRequest ) ; 
629+             expected_response. set_content_type ( MediaType :: PlainText ) ; 
514630            let  error_msg = format ! ( 
515631                "Invalid time to live value provided for token: {}. Please provide a value \  
516632, 
@@ -525,6 +641,7 @@ mod tests {
525641        let  request_bytes = b"PUT http://169.254.169.254/latest/api/token HTTP/1.0\r \n \r \n " ; 
526642        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
527643        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: BadRequest ) ; 
644+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
528645        expected_response. set_body ( Body :: new ( VmmMmdsError :: NoTtlProvided . to_string ( ) ) ) ; 
529646        let  actual_response = convert_to_response ( mmds. clone ( ) ,  request) ; 
530647        assert_eq ! ( actual_response,  expected_response) ; 
@@ -559,6 +676,7 @@ mod tests {
559676        ) ; 
560677        let  request = Request :: try_from ( request_bytes. as_bytes ( ) ,  None ) . unwrap ( ) ; 
561678        let  mut  expected_response = Response :: new ( Version :: Http11 ,  StatusCode :: NotImplemented ) ; 
679+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
562680        let  body = "Cannot retrieve value. The value has an unsupported type." . to_string ( ) ; 
563681        expected_response. set_body ( Body :: new ( body) ) ; 
564682        let  actual_response = convert_to_response ( mmds. clone ( ) ,  request) ; 
@@ -571,6 +689,7 @@ mod tests {
571689        ) ; 
572690        let  request = Request :: try_from ( request_bytes. as_bytes ( ) ,  None ) . unwrap ( ) ; 
573691        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: NotFound ) ; 
692+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
574693        expected_response. set_body ( Body :: new ( 
575694            VmmMmdsError :: ResourceNotFound ( String :: from ( "/invalid" ) ) . to_string ( ) , 
576695        ) ) ; 
@@ -581,6 +700,7 @@ mod tests {
581700        let  request_bytes = b"GET http://169.254.169.254/ HTTP/1.0\r \n \r \n " ; 
582701        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
583702        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: Unauthorized ) ; 
703+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
584704        expected_response. set_body ( Body :: new ( VmmMmdsError :: NoTokenProvided . to_string ( ) ) ) ; 
585705        let  actual_response = convert_to_response ( mmds. clone ( ) ,  request) ; 
586706        assert_eq ! ( actual_response,  expected_response) ; 
@@ -590,6 +710,7 @@ mod tests {
590710\r \n \r \n "; 
591711        let  request = Request :: try_from ( request_bytes,  None ) . unwrap ( ) ; 
592712        let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: Unauthorized ) ; 
713+         expected_response. set_content_type ( MediaType :: PlainText ) ; 
593714        expected_response. set_body ( Body :: new ( VmmMmdsError :: InvalidToken . to_string ( ) ) ) ; 
594715        let  actual_response = convert_to_response ( mmds. clone ( ) ,  request) ; 
595716        assert_eq ! ( actual_response,  expected_response) ; 
@@ -614,6 +735,7 @@ mod tests {
614735            ) ; 
615736            let  request = Request :: try_from ( request_bytes. as_bytes ( ) ,  None ) . unwrap ( ) ; 
616737            let  mut  expected_response = Response :: new ( Version :: Http10 ,  StatusCode :: Unauthorized ) ; 
738+             expected_response. set_content_type ( MediaType :: PlainText ) ; 
617739            expected_response. set_body ( Body :: new ( VmmMmdsError :: InvalidToken . to_string ( ) ) ) ; 
618740            let  actual_response = convert_to_response ( mmds. clone ( ) ,  request) ; 
619741            assert_eq ! ( actual_response,  expected_response) ; 
0 commit comments