diff --git a/CHANGELOG.md b/CHANGELOG.md index 87a1086..7558f5a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,10 +2,11 @@ ## Added -- Implemented `Eq` for `common::headers::Encoding`, `common::headers::MediaType`, +- Implemented `Eq` for `common::headers::Encoding`, `common::headers::MediaType`, `common::headers::Headers`, `common::HttpHeaderError`, `common::Body`, `common::Version`, `common::RequestError`, `request::Uri`, `request::RequestLine`, `response::StatusCode`, `response::ResponseHeaders` +- Allowed to set custom headers in HTTP responses. ## Changed diff --git a/src/common/mod.rs b/src/common/mod.rs index b211b8c..56267b6 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -23,6 +23,8 @@ pub enum HttpHeaderError { InvalidUtf8String(Utf8Error), ///The value specified is not valid. InvalidValue(String, String), + /// Non-ASCII character found. + NonAsciiCharacter(String, String), /// The content length specified is longer than the limit imposed by Micro Http. SizeLimitExceeded(String), /// The requested feature is not currently supported. @@ -45,6 +47,13 @@ impl Display for HttpHeaderError { Self::InvalidValue(header_name, value) => { write!(f, "Invalid value. Key:{}; Value:{}", header_name, value) } + Self::NonAsciiCharacter(header_name, value) => { + write!( + f, + "Non-ASCII character found. Key: {}; Value: {}", + header_name, value + ) + } Self::SizeLimitExceeded(inner) => { write!(f, "Invalid content length. Header: {}", inner) } diff --git a/src/response.rs b/src/response.rs index d60a0e5..f082415 100644 --- a/src/response.rs +++ b/src/response.rs @@ -1,11 +1,13 @@ // Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashMap; use std::io::{Error as WriteError, Write}; use crate::ascii::{COLON, CR, LF, SP}; use crate::common::{Body, Version}; use crate::headers::{Header, MediaType}; +use crate::HttpHeaderError; use crate::Method; /// Wrapper over a response status code. @@ -95,6 +97,7 @@ pub struct ResponseHeaders { server: String, allow: Vec, accept_encoding: bool, + custom_headers: HashMap, } impl Default for ResponseHeaders { @@ -106,6 +109,7 @@ impl Default for ResponseHeaders { server: String::from("Firecracker API"), allow: Vec::new(), accept_encoding: false, + custom_headers: HashMap::default(), } } } @@ -141,6 +145,18 @@ impl ResponseHeaders { buf.write_all(&[CR, LF]) } + // The logic pertaining to custom headers writing. + fn write_custom_headers(&self, buf: &mut T) -> Result<(), WriteError> { + // Note that all the custom headers have already been validated as US-ASCII. + for (header, value) in &self.custom_headers { + buf.write_all(header.as_bytes())?; + buf.write_all(&[COLON, SP])?; + buf.write_all(value.as_bytes())?; + buf.write_all(&[CR, LF])?; + } + Ok(()) + } + /// Writes the headers to `buf` using the HTTP specification. pub fn write_all(&self, buf: &mut T) -> Result<(), WriteError> { buf.write_all(Header::Server.raw())?; @@ -153,6 +169,7 @@ impl ResponseHeaders { self.write_allow_header(buf)?; self.write_deprecation_header(buf)?; + self.write_custom_headers(buf)?; if let Some(content_length) = self.content_length { buf.write_all(Header::ContentType.raw())?; @@ -203,6 +220,26 @@ impl ResponseHeaders { pub fn set_encoding(&mut self) { self.accept_encoding = true; } + + /// Sets custom headers to be written in the HTTP response. + pub fn set_custom_headers( + &mut self, + custom_headers: &HashMap, + ) -> Result<(), HttpHeaderError> { + // https://datatracker.ietf.org/doc/html/rfc7230 + // HTTP headers MUST be US-ASCII. + if let Some((k, v)) = custom_headers + .iter() + .find(|(k, v)| !k.is_ascii() || !v.is_ascii()) + { + return Err(HttpHeaderError::NonAsciiCharacter( + k.to_owned(), + v.to_owned(), + )); + } + self.custom_headers = custom_headers.to_owned(); + Ok(()) + } } /// Wrapper over an HTTP Response. @@ -294,6 +331,19 @@ impl Response { self.headers.set_server(server); } + /// Sets the custom headers. + pub fn set_custom_headers( + &mut self, + custom_headers: &HashMap, + ) -> Result<(), HttpHeaderError> { + self.headers.set_custom_headers(custom_headers) + } + + /// Gets a reference to the custom headers. + pub fn custom_headers(&self) -> &HashMap { + &self.headers.custom_headers + } + /// Sets the HTTP allowed methods. pub fn set_allow(&mut self, methods: Vec) { self.headers.allow = methods; @@ -554,4 +604,29 @@ mod tests { assert!(response.write_all(&mut response_buf.as_mut()).is_ok()); assert_eq!(response_buf.as_ref(), expected_response); } + + #[test] + fn test_custom_headers() { + // Valid custom headers. + let mut response = Response::new(Version::Http10, StatusCode::OK); + let custom_headers = [("Foo".into(), "Bar".into())].into(); + response.set_custom_headers(&custom_headers).unwrap(); + let expected_response = b"HTTP/1.0 200 \r\n\ + Server: Firecracker API\r\n\ + Connection: keep-alive\r\n\ + Foo: Bar\r\n\ + Content-Type: application/json\r\n\ + Content-Length: 0\r\n\r\n"; + let mut response_buf: [u8; 127] = [0; 127]; + response.write_all(&mut response_buf.as_mut()).unwrap(); + assert_eq!(response_buf.as_ref(), expected_response); + + // Should fail to set custom headers including non-ASCII character. + let mut response = Response::new(Version::Http10, StatusCode::OK); + let custom_headers = [("Greek capital delta".into(), "Δ".into())].into(); + assert_eq!( + response.set_custom_headers(&custom_headers).unwrap_err(), + HttpHeaderError::NonAsciiCharacter("Greek capital delta".into(), "Δ".into()) + ); + } }