[Fix] Set max size on per connection basis

4 files changed, 39 insertions, 19 deletions

diff --git a/src/libutil/http.c b/src/libutil/http.c
index 03d28db77..d7f10d20f 100644
--- a/src/libutil/http.c
+++ b/src/libutil/http.c
@@ -102,7 +102,6 @@ static const rspamd_ftok_t last_modified_header = {
 		.begin = "Last-Modified",
 		.len = 13
 };
-static gsize rspamd_http_global_max_size = 0;
 
 static void rspamd_http_message_storage_cleanup (struct rspamd_http_message *msg);
 static gboolean rspamd_http_message_grow_body (struct rspamd_http_message *msg,
@@ -617,16 +616,23 @@ rspamd_http_on_headers_complete (http_parser * parser)
 		priv->flags &= ~RSPAMD_HTTP_CONN_FLAG_NEW_HEADER;
 	}
 
-	if (conn->type == RSPAMD_HTTP_SERVER &&
-			rspamd_http_global_max_size > 0 &&
-			parser->content_length > rspamd_http_global_max_size) {
-		/* Too large message */
-		priv->flags |= RSPAMD_HTTP_CONN_FLAG_TOO_LARGE;
-		return -1;
-	}
+	/*
+	 * HTTP parser sets content length to (-1) when it doesn't know the real
+	 * length, for example, in case of chunked encoding.
+	 *
+	 * Hence, we skip body setup here
+	 */
+	if (parser->content_length != ULLONG_MAX) {
+		if (conn->max_size > 0 &&
+				parser->content_length > conn->max_size) {
+			/* Too large message */
+			priv->flags |= RSPAMD_HTTP_CONN_FLAG_TOO_LARGE;
+			return -1;
+		}
 
-	if (!rspamd_http_message_set_body (msg, NULL, parser->content_length)) {
-		return -1;
+		if (!rspamd_http_message_set_body (msg, NULL, parser->content_length)) {
+			return -1;
+		}
 	}
 
 	if (parser->flags & F_SPAMC) {
@@ -663,9 +669,14 @@ rspamd_http_on_body (http_parser * parser, const gchar *at, size_t length)
 	pbuf = priv->buf;
 	p = at;
 
-	if (conn->type == RSPAMD_HTTP_SERVER &&
-			rspamd_http_global_max_size > 0 &&
-			msg->body_buf.len + length > rspamd_http_global_max_size) {
+	if (!(msg->flags & RSPAMD_HTTP_FLAG_HAS_BODY)) {
+		if (!rspamd_http_message_set_body (msg, NULL, parser->content_length)) {
+			return -1;
+		}
+	}
+
+	if (conn->max_size > 0 &&
+			msg->body_buf.len + length > conn->max_size) {
 		/* Body length overflow */
 		priv->flags |= RSPAMD_HTTP_CONN_FLAG_TOO_LARGE;
 		return -1;
@@ -1097,7 +1108,8 @@ rspamd_http_event_handler (int fd, short what, gpointer ud)
 					d, r) != (size_t)r || priv->parser.http_errno != 0) {
 				if (priv->flags & RSPAMD_HTTP_CONN_FLAG_TOO_LARGE) {
 					err = g_error_new (HTTP_ERROR, 413,
-							"Request entity too large");
+							"Request entity too large: %zu",
+							(size_t)priv->parser.content_length);
 				}
 				else {
 					err = g_error_new (HTTP_ERROR, priv->parser.http_errno,
@@ -2331,6 +2343,8 @@ rspamd_http_message_set_body (struct rspamd_http_message *msg,
 		msg->body_buf.allocated_len = storage->normal->allocated;
 	}
 
+	msg->flags |= RSPAMD_HTTP_FLAG_HAS_BODY;
+
 	return TRUE;
 }
 
@@ -2536,9 +2550,10 @@ rspamd_http_message_storage_cleanup (struct rspamd_http_message *msg)
 }
 
 void
-rspamd_http_message_set_max_size (gsize sz)
+rspamd_http_connection_set_max_size (struct rspamd_http_connection *conn,
+		gsize sz)
 {
-	rspamd_http_global_max_size = sz;
+	conn->max_size = sz;
 }
 
 void
diff --git a/src/libutil/http.h b/src/libutil/http.h
index 8223feb17..9dc6302c9 100644
--- a/src/libutil/http.h
+++ b/src/libutil/http.h
@@ -63,6 +63,10 @@ struct rspamd_storage_shmem {
  * Use tls for this message
  */
 #define RSPAMD_HTTP_FLAG_SSL (1 << 4)
+/**
+ * Body has been set for a message
+ */
+#define RSPAMD_HTTP_FLAG_HAS_BODY (1 << 5)
 
 /**
  * Options for HTTP connection
@@ -104,6 +108,7 @@ struct rspamd_http_connection {
 	rspamd_http_finish_handler_t finish_handler;
 	struct rspamd_keypair_cache *cache;
 	gpointer ud;
+	gsize max_size;
 	unsigned opts;
 	enum rspamd_http_connection_type type;
 	gboolean finished;
@@ -389,7 +394,8 @@ void rspamd_http_message_free (struct rspamd_http_message *msg);
  * Sets global maximum size for HTTP message being processed
  * @param sz
  */
-void rspamd_http_message_set_max_size (gsize sz);
+void rspamd_http_connection_set_max_size (struct rspamd_http_connection *conn,
+		gsize sz);
 
 /**
  * Increase refcount for shared file (if any) to prevent early memory unlinking
diff --git a/src/libutil/util.c b/src/libutil/util.c
index ca18a7fcc..762c8d557 100644
--- a/src/libutil/util.c
+++ b/src/libutil/util.c
@@ -2118,8 +2118,6 @@ rspamd_config_libs (struct rspamd_external_libs_ctx *ctx,
 			magic_load (ctx->libmagic, cfg->magic_file);
 		}
 	}
-
-	rspamd_http_message_set_max_size (cfg->max_message);
 }
 
 void
diff --git a/src/worker.c b/src/worker.c
index 98473b003..60c39b2af 100644
--- a/src/worker.c
+++ b/src/worker.c
@@ -340,6 +340,7 @@ accept_socket (gint fd, short what, void *arg)
 			RSPAMD_HTTP_SERVER,
 			ctx->keys_cache,
 			NULL);
+	rspamd_http_connection_set_max_size (task->http_conn, task->cfg->max_message);
 	task->ev_base = ctx->ev_base;
 	worker->nconns++;
 	rspamd_mempool_add_destructor (task->task_pool,