diff --git a/src/connection.rs b/src/connection.rs index 2bf53b4..99b3b24 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -11,6 +11,7 @@ pub use crate::common::{ConnectionError, HttpHeaderError, RequestError}; use crate::headers::Headers; use crate::request::{find, Request, RequestLine}; use crate::response::{Response, StatusCode}; +use crate::server::MAX_PAYLOAD_SIZE; use vmm_sys_util::sock_ctrl_msg::ScmSocket; const BUFFER_SIZE: usize = 1024; @@ -54,6 +55,8 @@ pub struct HttpConnection { /// The latest file that has been received and which must be associated /// with the pending request. file: Option, + /// Optional payload max size. + payload_max_size: usize, } impl HttpConnection { @@ -71,9 +74,16 @@ impl HttpConnection { response_queue: VecDeque::new(), response_buffer: None, file: None, + payload_max_size: MAX_PAYLOAD_SIZE, } } + /// This function sets the limit for PUT/PATCH requests. It overwrites the + /// default limit of 0.05MiB with the one allowed by server. + pub fn set_payload_max_size(&mut self, request_payload_max_size: usize) { + self.payload_max_size = request_payload_max_size; + } + /// Tries to read new bytes from the stream and automatically update the request. /// Meant to be used only with non-blocking streams and an `EPOLL` structure. /// Should be called whenever an `EPOLLIN` event is signaled. @@ -245,6 +255,14 @@ impl HttpConnection { if request.headers.content_length() == 0 { self.state = ConnectionState::RequestReady; } else { + if request.headers.content_length() as usize > self.payload_max_size { + return Err(ConnectionError::ParseError( + RequestError::SizeLimitExceeded( + self.payload_max_size, + request.headers.content_length() as usize, + ), + )); + } if request.headers.expect() { // Send expect. let expect_response = @@ -504,6 +522,7 @@ mod tests { use super::*; use crate::common::{Method, Version}; + use crate::server::MAX_PAYLOAD_SIZE; #[test] fn test_try_read_expect() { @@ -933,6 +952,24 @@ mod tests { ); } + #[test] + fn test_payload_size_limit() { + let (mut sender, receiver) = UnixStream::pair().unwrap(); + receiver.set_nonblocking(true).expect("Can't modify socket"); + let mut conn = HttpConnection::new(receiver); + conn.set_payload_max_size(5); + sender + .write_all( + b"PUT http://localhost/home HTTP/1.1\r\n\ + Content-Length: 51200\r\n\r\naaaaaa", + ) + .unwrap(); + assert_eq!( + conn.try_read().unwrap_err(), + ConnectionError::ParseError(RequestError::SizeLimitExceeded(5, MAX_PAYLOAD_SIZE)) + ); + } + #[test] fn test_read_bytes() { let (mut sender, receiver) = UnixStream::pair().unwrap(); diff --git a/src/server.rs b/src/server.rs index c534ca8..eeff881 100644 --- a/src/server.rs +++ b/src/server.rs @@ -21,6 +21,8 @@ static SERVER_FULL_ERROR_MESSAGE: &[u8] = b"HTTP/1.1 503\r\n\ Connection: close\r\n\ Content-Length: 40\r\n\r\n{ \"error\": \"Too many open connections\" }"; const MAX_CONNECTIONS: usize = 10; +/// Payload max size +pub(crate) const MAX_PAYLOAD_SIZE: usize = 51200; type Result = std::result::Result; @@ -259,6 +261,8 @@ pub struct HttpServer { /// We use the file descriptor of the stream as the key for mapping /// connections because the 1-to-1 relation is guaranteed by the OS. connections: HashMap>, + /// Payload max size + payload_max_size: usize, } impl HttpServer { @@ -275,6 +279,7 @@ impl HttpServer { socket, epoll, connections: HashMap::new(), + payload_max_size: MAX_PAYLOAD_SIZE, }) } @@ -295,9 +300,16 @@ impl HttpServer { socket, epoll, connections: HashMap::new(), + payload_max_size: MAX_PAYLOAD_SIZE, }) } + /// This function sets the limit for PUT/PATCH requests. It overwrites the + /// default limit of 0.05MiB with the one allowed by server. + pub fn set_payload_max_size(&mut self, request_payload_max_size: usize) { + self.payload_max_size = request_payload_max_size; + } + /// Starts the HTTP Server. pub fn start_server(&mut self) -> Result<()> { // Add the socket on which we listen for new connections to the @@ -573,12 +585,12 @@ impl HttpServer { }) .and_then(|stream| { // Add the stream to the `epoll` structure and listen for bytes to be read. - Self::epoll_add(&self.epoll, stream.as_raw_fd())?; + let raw_fd = stream.as_raw_fd(); + Self::epoll_add(&self.epoll, raw_fd)?; + let mut conn = HttpConnection::new(stream); + conn.set_payload_max_size(self.payload_max_size); // Then add it to our open connections. - self.connections.insert( - stream.as_raw_fd(), - ClientConnection::new(HttpConnection::new(stream)), - ); + self.connections.insert(raw_fd, ClientConnection::new(conn)); Ok(()) }) } @@ -676,6 +688,73 @@ mod tests { assert!(socket.read(&mut buf[..]).unwrap() > 0); } + #[test] + fn test_connection_size_limit_exceeded() { + let path_to_socket = get_temp_socket_file(); + + let mut server = HttpServer::new(path_to_socket.as_path()).unwrap(); + server.start_server().unwrap(); + + // Test one incoming connection. + let mut socket = UnixStream::connect(path_to_socket.as_path()).unwrap(); + assert!(server.requests().unwrap().is_empty()); + + socket + .write_all( + b"PATCH /machine-config HTTP/1.1\r\n\ + Content-Length: 51201\r\n\ + Content-Type: application/json\r\n\r\naaaaa", + ) + .unwrap(); + assert!(server.requests().unwrap().is_empty()); + assert!(server.requests().unwrap().is_empty()); + let mut buf: [u8; 265] = [0; 265]; + assert!(socket.read(&mut buf[..]).unwrap() > 0); + let error_message = b"HTTP/1.1 400 \r\n\ + Server: Firecracker API\r\n\ + Connection: keep-alive\r\n\ + Content-Type: application/json\r\n\ + Content-Length: 149\r\n\r\n{ \"error\": \"\ + Request payload with size 51201 is larger than \ + the limit of 51200 allowed by server.\nAll \ + previous unanswered requests will be dropped."; + assert_eq!(&buf[..], &error_message[..]); + } + + #[test] + fn test_set_payload_size() { + let path_to_socket = get_temp_socket_file(); + + let mut server = HttpServer::new(path_to_socket.as_path()).unwrap(); + server.start_server().unwrap(); + server.set_payload_max_size(4); + + // Test one incoming connection. + let mut socket = UnixStream::connect(path_to_socket.as_path()).unwrap(); + assert!(server.requests().unwrap().is_empty()); + + socket + .write_all( + b"PATCH /machine-config HTTP/1.1\r\n\ + Content-Length: 5\r\n\ + Content-Type: application/json\r\n\r\naaaaa", + ) + .unwrap(); + assert!(server.requests().unwrap().is_empty()); + assert!(server.requests().unwrap().is_empty()); + let mut buf: [u8; 260] = [0; 260]; + assert!(socket.read(&mut buf[..]).unwrap() > 0); + let error_message = b"HTTP/1.1 400 \r\n\ + Server: Firecracker API\r\n\ + Connection: keep-alive\r\n\ + Content-Type: application/json\r\n\ + Content-Length: 141\r\n\r\n{ \"error\": \"\ + Request payload with size 5 is larger than the \ + limit of 4 allowed by server.\nAll previous \ + unanswered requests will be dropped.\" }"; + assert_eq!(&buf[..], &error_message[..]); + } + #[test] fn test_wait_one_fd_connection() { use std::os::unix::io::IntoRawFd;