diff --git a/micro_http/src/common/headers.rs b/micro_http/src/common/headers.rs index 1bf5d53dd..1effae299 100644 --- a/micro_http/src/common/headers.rs +++ b/micro_http/src/common/headers.rs @@ -16,6 +16,8 @@ pub enum Header { Expect, /// Header `Transfer-Encoding`. TransferEncoding, + /// Header `Server`. + Server, } impl Header { @@ -25,6 +27,7 @@ impl Header { Header::ContentType => b"Content-Type", Header::Expect => b"Expect", Header::TransferEncoding => b"Transfer-Encoding", + Header::Server => b"Server", } } @@ -35,6 +38,7 @@ impl Header { "Content-Type" => Ok(Header::ContentType), "Expect" => Ok(Header::Expect), "Transfer-Encoding" => Ok(Header::TransferEncoding), + "Server" => Ok(Header::Server), _ => Err(RequestError::InvalidHeader), } } else { @@ -130,6 +134,7 @@ impl Headers { } _ => Err(RequestError::InvalidHeader), }, + Header::Server => Ok(()), } } else { Err(RequestError::UnsupportedHeader) diff --git a/micro_http/src/response.rs b/micro_http/src/response.rs index 722b44a2c..160c716c8 100644 --- a/micro_http/src/response.rs +++ b/micro_http/src/response.rs @@ -70,16 +70,29 @@ impl StatusLine { /// Wrapper over the list of headers associated with a HTTP Response. /// When creating a ResponseHeaders object, the content type is initialized to `text/plain`. /// The content type can be updated with a call to `set_content_type`. -#[derive(Default)] pub struct ResponseHeaders { content_length: i32, content_type: MediaType, + server: String, +} + +impl Default for ResponseHeaders { + fn default() -> Self { + ResponseHeaders { + content_length: Default::default(), + content_type: Default::default(), + server: "Firecracker API".to_string(), + } + } } impl ResponseHeaders { /// Writes the headers to `buf` using the HTTP specification. pub fn write_all(&self, buf: &mut T) -> Result<(), WriteError> { - buf.write_all(b"Server: Firecracker API")?; + buf.write_all(Header::Server.raw())?; + buf.write_all(&[COLON, SP])?; + buf.write_all(self.server.as_bytes())?; + buf.write_all(&[CR, LF])?; buf.write_all(b"Connection: keep-alive")?; buf.write_all(&[CR, LF])?; @@ -104,6 +117,11 @@ impl ResponseHeaders { self.content_length = content_length; } + /// Sets the HTTP response header server. + pub fn set_server(&mut self, server: &str) { + self.server = String::from(server); + } + /// Sets the content type to be written in the HTTP response. #[allow(unused)] pub fn set_content_type(&mut self, content_type: MediaType) { @@ -141,6 +159,11 @@ impl Response { self.body = Some(body); } + /// Sets the HTTP response server. + pub fn set_server(&mut self, server: &str) { + self.headers.set_server(server); + } + fn write_body(&self, mut buf: T) -> Result<(), WriteError> { if let Some(ref body) = self.body { buf.write_all(body.raw())?; @@ -219,6 +242,35 @@ mod tests { assert!(response.write_all(&mut response_buf.as_mut()).is_err()); } + #[test] + fn test_set_server() { + let mut response = Response::new(Version::Http10, StatusCode::OK); + let body = "This is a test"; + let server = "rust-vmm API"; + response.set_body(Body::new(body)); + response.set_server(server); + + assert!(response.status() == StatusCode::OK); + assert_eq!(response.body().unwrap(), Body::new(body)); + assert_eq!(response.http_version(), Version::Http10); + assert_eq!(response.content_length(), 14); + assert_eq!(response.content_type(), MediaType::PlainText); + + let expected_response = format!( + "HTTP/1.0 200 \r\n\ + Server: {}\r\n\ + Connection: keep-alive\r\n\ + Content-Type: text/plain\r\n\ + Content-Length: 14\r\n\r\n\ + This is a test", + server + ); + + let mut response_buf: [u8; 123] = [0; 123]; + assert!(response.write_all(&mut response_buf.as_mut()).is_ok()); + assert!(response_buf.as_ref() == expected_response.as_bytes()); + } + #[test] fn test_status_code() { assert_eq!(StatusCode::Continue.raw(), b"100");