Add optional limit on request size

Signed-off-by: AlexandruCihodaru <cihodar@amazon.com>

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<T> {
     /// The latest file that has been received and which must be associated
     /// with the pending request.
     file: Option<File>,
+    /// Optional payload max size.
+    payload_max_size: usize,
 }
 
 impl<T: Read + Write + ScmSocket> HttpConnection<T> {
@@ -71,9 +74,16 @@ impl<T: Read + Write + ScmSocket> HttpConnection<T> {
             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<T: Read + Write + ScmSocket> HttpConnection<T> {
                 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();

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<T> = std::result::Result<T, ServerError>;
 
@@ -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<RawFd, ClientConnection<UnixStream>>,
+    /// 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;