[Feature] Allow to open message from a shared memory segment

1 files changed, 97 insertions, 35 deletions

diff --git a/src/libserver/task.c b/src/libserver/task.c
index ce95b927d..16e8cc681 100644
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -289,6 +289,64 @@ rspamd_task_load_message (struct rspamd_task *task,
 		rspamd_protocol_handle_headers (task, msg);
 	}
 
+	srch.begin = "shm";
+	srch.len = 3;
+	tok = g_hash_table_lookup (task->request_headers, &srch);
+
+	if (tok) {
+		/* Shared memory part */
+		r = rspamd_strlcpy (filepath, tok->begin,
+				MIN (sizeof (filepath), tok->len + 1));
+
+		rspamd_decode_url (filepath, filepath, r + 1);
+		flen = strlen (filepath);
+
+		if (filepath[0] == '"' && flen > 2) {
+			/* We need to unquote filepath */
+			fp = &filepath[1];
+			fp[flen - 2] = '\0';
+		}
+		else {
+			fp = &filepath[0];
+		}
+
+		fd = shm_open (fp, O_RDONLY);
+
+		if (fd == -1) {
+			g_set_error (&task->err, rspamd_task_quark(), RSPAMD_PROTOCOL_ERROR,
+					"Cannot open shm segment (%s): %s", fp, strerror (errno));
+			return FALSE;
+		}
+
+		if (fstat (fd, &st) == -1) {
+			g_set_error (&task->err, rspamd_task_quark(), RSPAMD_PROTOCOL_ERROR,
+					"Cannot stat shm segment (%s): %s", fp, strerror (errno));
+			close (fd);
+
+			return FALSE;
+		}
+
+		map = mmap (NULL, st.st_size, PROT_READ, MAP_SHARED, fd, 0);
+
+		if (map == MAP_FAILED) {
+			close (fd);
+			g_set_error (&task->err, rspamd_task_quark(), RSPAMD_PROTOCOL_ERROR,
+					"Cannot mmap file (%s): %s", fp, strerror (errno));
+			return FALSE;
+		}
+
+		close (fd);
+		task->msg.begin = map;
+		task->msg.len = st.st_size;
+		task->flags |= RSPAMD_TASK_FLAG_FILE;
+
+		msg_info_task ("loaded message from shared memory %s", fp);
+
+		rspamd_mempool_add_destructor (task->task_pool, rspamd_task_unmapper, task);
+
+		return TRUE;
+	}
+
 	srch.begin = "file";
 	srch.len = 4;
 	tok = g_hash_table_lookup (task->request_headers, &srch);
@@ -317,7 +375,7 @@ rspamd_task_load_message (struct rspamd_task *task,
 			fp = &filepath[0];
 		}
 
-		if (access (fp, R_OK) == -1 || stat (fp, &st) == -1) {
+		if (stat (fp, &st) == -1) {
 			g_set_error (&task->err, rspamd_task_quark(), RSPAMD_PROTOCOL_ERROR,
 					"Invalid file (%s): %s", fp, strerror (errno));
 			return FALSE;
@@ -346,46 +404,50 @@ rspamd_task_load_message (struct rspamd_task *task,
 		task->msg.len = st.st_size;
 		task->flags |= RSPAMD_TASK_FLAG_FILE;
 
+		msg_info_task ("loaded message from file %s", fp);
+
 		rspamd_mempool_add_destructor (task->task_pool, rspamd_task_unmapper, task);
+
+		return TRUE;
 	}
-	else {
-		debug_task ("got input of length %z", task->msg.len);
-		task->msg.begin = start;
-		task->msg.len = len;
-
-		if (task->msg.len == 0) {
-			task->flags |= RSPAMD_TASK_FLAG_EMPTY;
-		}
-
-		if (task->flags & RSPAMD_TASK_FLAG_HAS_CONTROL) {
-			/* We have control chunk, so we need to process it separately */
-			if (task->msg.len < task->message_len) {
-				msg_warn_task ("message has invalid message length: %ul and total len: %ul",
-						task->message_len, task->msg.len);
-				g_set_error (&task->err, rspamd_task_quark(), RSPAMD_PROTOCOL_ERROR,
-						"Invalid length");
-				return FALSE;
-			}
-			control_len = task->msg.len - task->message_len;
 
-			if (control_len > 0) {
-				parser = ucl_parser_new (UCL_PARSER_KEY_LOWERCASE);
+	/* Plain data */
+	debug_task ("got input of length %z", task->msg.len);
+	task->msg.begin = start;
+	task->msg.len = len;
 
-				if (!ucl_parser_add_chunk (parser, task->msg.begin, control_len)) {
-					msg_warn_task ("processing of control chunk failed: %s",
-							ucl_parser_get_error (parser));
-					ucl_parser_free (parser);
-				}
-				else {
-					control_obj = ucl_parser_get_object (parser);
-					ucl_parser_free (parser);
-					rspamd_protocol_handle_control (task, control_obj);
-					ucl_object_unref (control_obj);
-				}
+	if (task->msg.len == 0) {
+		task->flags |= RSPAMD_TASK_FLAG_EMPTY;
+	}
 
-				task->msg.begin += control_len;
-				task->msg.len -= control_len;
+	if (task->flags & RSPAMD_TASK_FLAG_HAS_CONTROL) {
+		/* We have control chunk, so we need to process it separately */
+		if (task->msg.len < task->message_len) {
+			msg_warn_task ("message has invalid message length: %ul and total len: %ul",
+					task->message_len, task->msg.len);
+			g_set_error (&task->err, rspamd_task_quark(), RSPAMD_PROTOCOL_ERROR,
+					"Invalid length");
+			return FALSE;
+		}
+		control_len = task->msg.len - task->message_len;
+
+		if (control_len > 0) {
+			parser = ucl_parser_new (UCL_PARSER_KEY_LOWERCASE);
+
+			if (!ucl_parser_add_chunk (parser, task->msg.begin, control_len)) {
+				msg_warn_task ("processing of control chunk failed: %s",
+						ucl_parser_get_error (parser));
+				ucl_parser_free (parser);
+			}
+			else {
+				control_obj = ucl_parser_get_object (parser);
+				ucl_parser_free (parser);
+				rspamd_protocol_handle_control (task, control_obj);
+				ucl_object_unref (control_obj);
 			}
+
+			task->msg.begin += control_len;
+			task->msg.len -= control_len;
 		}
 	}