Add a way to limit size of uploads

wlandry · wlandry · commit 369707bd195b · 2016-04-17T22:07:00.000-07:00
diff --git a/src/httpserver/create_webserver.hpp b/src/httpserver/create_webserver.hpp
@@ -52,6 +52,7 @@ class create_webserver
             _max_threads(0),
             _max_connections(0),
             _memory_limit(0),
+            _content_size_limit(0),
             _connection_timeout(DEFAULT_WS_TIMEOUT),
             _per_IP_connection_limit(0),
             _log_access(0x0),
@@ -93,6 +94,7 @@ class create_webserver
             _max_threads(0),
             _max_connections(0),
             _memory_limit(0),
+            _content_size_limit(0),
             _connection_timeout(DEFAULT_WS_TIMEOUT),
             _per_IP_connection_limit(0),
             _log_access(0x0),
@@ -147,6 +149,10 @@ class create_webserver
         {
             _memory_limit = memory_limit; return *this;
         }
+        create_webserver& content_size_limit(size_t content_size_limit)
+        {
+            _content_size_limit = content_size_limit; return *this;
+        }
         create_webserver& connection_timeout(int connection_timeout)
         {
             _connection_timeout = connection_timeout; return *this;
@@ -333,6 +339,7 @@ class create_webserver
         int _max_threads;
         int _max_connections;
         int _memory_limit;
+        size_t _content_size_limit;
         int _connection_timeout;
         int _per_IP_connection_limit;
         log_access_ptr _log_access;
diff --git a/src/httpserver/http_request.hpp b/src/httpserver/http_request.hpp
@@ -467,9 +467,14 @@ class http_request
          * @param content The content to append.
          * @param size The size of the data to append.
         **/
-        void grow_content(const char* content, size_t size)
+        void grow_content(const char* content, size_t size,
+                          size_t content_size_limit)
         {
             this->content.append(content, size);
+            if (this->content.size() > content_size_limit)
+              {
+                this->content.resize (content_size_limit);
+              }
         }
         /**
          * Method used to set the path requested.
diff --git a/src/httpserver/webserver.hpp b/src/httpserver/webserver.hpp
@@ -179,6 +179,7 @@ class webserver
         const int max_threads;
         const int max_connections;
         const int memory_limit;
+        const size_t content_size_limit;
         const int connection_timeout;
         const int per_IP_connection_limit;
         log_access_ptr log_access;
diff --git a/src/webserver.cpp b/src/webserver.cpp
@@ -141,6 +141,7 @@ webserver::webserver(const create_webserver& params):
     max_threads(params._max_threads),
     max_connections(params._max_connections),
     memory_limit(params._memory_limit),
+    content_size_limit(params._content_size_limit),
     connection_timeout(params._connection_timeout),
     per_IP_connection_limit(params._per_IP_connection_limit),
     log_access(params._log_access),
@@ -698,7 +699,7 @@ int webserver::bodyfull_requests_answer_second_step(
 #ifdef DEBUG
     cout << "Writing content: " << upload_data << endl;
 #endif //DEBUG
-    mr->dhr->grow_content(upload_data, *upload_data_size);
+    mr->dhr->grow_content(upload_data, *upload_data_size, content_size_limit);
 
     if (mr->pp != NULL) MHD_post_process(mr->pp, upload_data, *upload_data_size);
     *upload_data_size = 0;
