feature: max_body_size check and fixing is_done()

Author: StokesAsselborn

src/handler/upload_handler.cpp

@@ -18,7 +18,8 @@ UploadHandler::UploadHandler(const std::string& path, const RouteConfig& rc, con
       req_(req),
       out_off_(0),
       bytes_written_(0),
-      fd_(-1)
+      fd_(-1),
+      eof_(false)
 
 {
     if (req.content_length == 0) {
@@ -60,20 +61,28 @@ size_t UploadHandler::read_output(char* buf, size_t n)
 size_t UploadHandler::write_input(const char* buf, size_t n)
 {
     ssize_t bytes = write(fd_, buf, n);
+    if (bytes == 0) {
+        // should never happen ?
+    }
     if (bytes < 0) {
         if (out_buf_.empty()) {
             set_error(HttpResponse::kStatusDiskFull);
         }
-        bytes_written_ = req_.content_length; // to ensure needs_input returns false
+        eof_ = true; // to ensure needs_input returns false
         return 0;
     }
     bytes_written_ += bytes;
-    if (bytes_written_ < req_.content_length) {
+    if (bytes_written_ < rc_.shared.max_body_size && bytes_written_ < req_.content_length) {
         return (bytes);
     }
+    if (bytes_written_ >= rc_.shared.max_body_size) {
+        if (out_buf_.empty())
+            set_error(HttpResponse::kStatusBadRequest);
+    }
     if (out_buf_.empty()) {
         res_ = HttpResponse::make_response_headers_only(HttpResponse::kStatusCreated, "", 0, req_);
         out_buf_ = res_.to_string();
+        eof_ = true;
     }
     return (0);
 }
@@ -82,4 +91,5 @@ void UploadHandler::set_error(const HttpResponse::Status code)
 {
     res_ = HttpResponse::make_error(code, rc_.shared.error_pages, req_);
     out_buf_ = res_.to_string();
+    eof_ = true;
 }

src/handler/upload_handler.hpp

@@ -20,7 +20,7 @@ class UploadHandler : public Handler {
     virtual bool is_regular_file() const { return true; }
     virtual bool has_output() const { return out_off_ < out_buf_.size(); }
     virtual bool needs_input() const { return bytes_written_ < req_.content_length; }
-    virtual bool is_done() const { return out_off_ >= out_buf_.size(); }
+    virtual bool is_done() const { return eof_ && out_off_ >= out_buf_.size(); }
 
     virtual int cgi_read_fd() const { return -1; }
     virtual int cgi_write_fd() const { return -1; }
@@ -44,6 +44,7 @@ class UploadHandler : public Handler {
     // other handler specifc variables
     size_t bytes_written_;
     int fd_;
+    bool eof_;
 };
 
 #endif

tests/upload_handler_unittest.cpp

@@ -3,17 +3,30 @@
 
 #include <cstdio>
 
-static void write_all(Handler& handler, const std::string& content)
+static size_t write_all(Handler& handler, const std::string& content)
 {
     size_t written = 0;
 
-    while (handler.needs_input()) {
+    while (!handler.is_done() && handler.needs_input()) {
         size_t n = handler.write_input(content.c_str() + written, content.size() - written);
         written += n;
 
         if ((n == 0 && handler.needs_input()))
             break;
     }
+    return written;
+}
+
+UTEST(UploadHandlerTest, is_done_upload_ongoing)
+{
+    std::string upload_path = "www/example/needs_input_upload_ongoing.txt";
+    std::string content = "some content to be written to the file";
+    RouteConfig rc;
+    HttpRequest req;
+    req.content_length = content.size();
+    UploadHandler handler(upload_path, rc, req);
+    ASSERT_FALSE(handler.is_done());
+    std::remove(upload_path.c_str());
 }
 
 UTEST(UploadHandlerTest, needs_input_upload_ongoing)
@@ -135,3 +148,18 @@ UTEST(UploadHandlerTest, return_507)
 {
     UTEST_SKIP("TODO: Test Disk full without making your machine explode");
 }
+
+UTEST(UploadHandlerTest, max_body_size)
+{
+    std::string upload_path = "www/example/needs_input_upload_done.txt";
+    std::string content = "some content to be written to the file";
+    RouteConfig rc;
+    rc.shared.max_body_size = 5;
+    HttpRequest req;
+    req.content_length = content.size();
+    ASSERT_TRUE(req.content_length > rc.shared.max_body_size);
+    UploadHandler handler(upload_path, rc, req);
+    size_t n = write_all(handler, content);
+    ASSERT_TRUE(n < req.content_length);
+    std::remove(upload_path.c_str());
+}