[Fix] Fix length calculations for shared memory segments

2 files changed, 26 insertions, 4 deletions

diff --git a/src/libserver/task.c b/src/libserver/task.c
index 82e4d6fa9..f8ebe81f1 100644
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -286,7 +286,7 @@ rspamd_task_load_message (struct rspamd_task *task,
 	ucl_object_t *control_obj;
 	gchar filepath[PATH_MAX], *fp;
 	gint fd, flen;
-	gulong offset = 0;
+	gulong offset = 0, shmem_size = 0;
 	rspamd_ftok_t srch, *tok;
 	gpointer map;
 	struct stat st;
@@ -360,14 +360,32 @@ rspamd_task_load_message (struct rspamd_task *task,
 			}
 		}
 
+		srch.begin = "shm-length";
+		srch.len = 10;
+		tok = g_hash_table_lookup (task->request_headers, &srch);
+		shmem_size = st.st_size;
+
+		if (tok) {
+			rspamd_strtoul (tok->begin, tok->len, &shmem_size);
+
+			if (shmem_size > (gulong)st.st_size) {
+				msg_err_task ("invalid length %ul (%ul available) for shm "
+						"segment %s", shmem_size, st.st_size, fp);
+				munmap (map, st.st_size);
+
+				return FALSE;
+			}
+		}
+
 		task->msg.begin = ((guchar *)map) + offset;
-		task->msg.len = st.st_size - offset;
+		task->msg.len = shmem_size;
 		task->flags |= RSPAMD_TASK_FLAG_FILE;
 		m = rspamd_mempool_alloc (task->task_pool, sizeof (*m));
 		m->begin = map;
 		m->len = st.st_size;
 
-		msg_info_task ("loaded message from shared memory %s", fp);
+		msg_info_task ("loaded message from shared memory %s (%ul size, %ul offset)",
+				fp, shmem_size, offset);
 
 		rspamd_mempool_add_destructor (task->task_pool, rspamd_task_unmapper, m);
 
diff --git a/src/libutil/http.c b/src/libutil/http.c
index 26fd3410d..57618653a 100644
--- a/src/libutil/http.c
+++ b/src/libutil/http.c
@@ -1233,7 +1233,7 @@ rspamd_http_connection_copy_msg (struct rspamd_http_connection *conn)
 			}
 
 			new_msg->body_buf.begin = new_msg->body_buf.str;
-			new_msg->body_buf.len = st.st_size;
+			new_msg->body_buf.len = msg->body_buf.len;
 			new_msg->body_buf.begin = new_msg->body_buf.str +
 					(msg->body_buf.begin - msg->body_buf.str);
 		}
@@ -1557,6 +1557,10 @@ rspamd_http_connection_write_message_common (struct rspamd_http_connection *conn
 					(int)(msg->body_buf.begin - msg->body_buf.str));
 			rspamd_http_message_add_header (msg, "Shm-Offset",
 					datebuf);
+			rspamd_snprintf (datebuf, sizeof (datebuf), "%z",
+					msg->body_buf.len);
+			rspamd_http_message_add_header (msg, "Shm-Length",
+					datebuf);
 		}
 	}