Merge pull request #4867 from rspamd/vstakhov-reply-type

[Feature] Support reply in message pack format

7 files changed, 67 insertions, 19 deletions

diff --git a/contrib/libucl/ucl_parser.c b/contrib/libucl/ucl_parser.c
index 354bfe857..b18fd06ce 100644
--- a/contrib/libucl/ucl_parser.c
+++ b/contrib/libucl/ucl_parser.c
@@ -2984,7 +2984,7 @@ ucl_parser_add_chunk_full (struct ucl_parser *parser, const unsigned char *data,
 
 		if (parse_type == UCL_PARSE_AUTO && len > 0) {
 			/* We need to detect parse type by the first symbol */
-			if ((*data & 0x80) == 0x80 && (*data >= 0xdc && *data <= 0xdf)) {
+			if ((*data & 0x80) == 0x80) {
 				parse_type = UCL_PARSE_MSGPACK;
 			}
 			else if (*data == '(') {
diff --git a/src/client/rspamdclient.c b/src/client/rspamdclient.c
index 85f4749f1..e77f662eb 100644
--- a/src/client/rspamdclient.c
+++ b/src/client/rspamdclient.c
@@ -1,11 +1,11 @@
-/*-
- * Copyright 2016 Vsevolod Stakhov
+/*
+ * Copyright 2024 Vsevolod Stakhov
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *   http://www.apache.org/licenses/LICENSE-2.0
+ *    http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
@@ -232,7 +232,9 @@ rspamd_client_finish_handler(struct rspamd_http_connection *conn,
 		}
 
 		parser = ucl_parser_new(0);
-		if (!ucl_parser_add_chunk(parser, start, len)) {
+		if (!ucl_parser_add_chunk_full(parser, start, len,
+									   ucl_parser_get_default_priority(parser),
+									   UCL_DUPLICATE_APPEND, UCL_PARSE_AUTO)) {
 			err = g_error_new(RCLIENT_ERROR, msg->code, "Cannot parse UCL: %s",
 							  ucl_parser_get_error(parser));
 			ucl_parser_free(parser);
@@ -454,6 +456,11 @@ rspamd_client_command(struct rspamd_client_connection *conn,
 		rspamd_http_message_add_header(req->msg, "Filename", filename);
 	}
 
+	/*
+	 * Allow messagepack reply if supported
+	 */
+	rspamd_http_message_add_header(req->msg, "Accept", "application/msgpack");
+
 	req->msg->url = rspamd_fstring_append(req->msg->url, "/", 1);
 	req->msg->url = rspamd_fstring_append(req->msg->url, command, strlen(command));
 
diff --git a/src/controller.c b/src/controller.c
index c9cda1d62..e68e10ced 100644
--- a/src/controller.c
+++ b/src/controller.c
@@ -1574,7 +1574,6 @@ rspamd_controller_handle_lua_history(lua_State *L,
 				}
 
 				task->http_conn = rspamd_http_connection_ref(conn_ent->conn);
-				;
 				task->sock = -1;
 				session->task = task;
 
@@ -1904,7 +1903,6 @@ rspamd_controller_handle_lua(struct rspamd_http_connection_entry *conn_ent,
 									task);
 	task->fin_arg = conn_ent;
 	task->http_conn = rspamd_http_connection_ref(conn_ent->conn);
-	;
 	task->sock = -1;
 	session->task = task;
 
@@ -2005,16 +2003,25 @@ rspamd_controller_scan_reply(struct rspamd_task *task)
 {
 	struct rspamd_http_message *msg;
 	struct rspamd_http_connection_entry *conn_ent;
+	int out_type = UCL_EMIT_JSON_COMPACT;
+	const char *ctype = "application/json";
+	const rspamd_ftok_t *accept_hdr = rspamd_task_get_request_header(task, "Accept");
+
+	if (accept_hdr && rspamd_substring_search(accept_hdr->begin, accept_hdr->len,
+											  "application/msgpack", sizeof("application/msgpack") - 1) != -1) {
+		ctype = "application/msgpack";
+		out_type = UCL_EMIT_MSGPACK;
+	}
 
 	conn_ent = task->fin_arg;
 	msg = rspamd_http_new_message(HTTP_RESPONSE);
 	msg->date = time(NULL);
 	msg->code = 200;
-	rspamd_protocol_http_reply(msg, task, NULL);
+	rspamd_protocol_http_reply(msg, task, NULL, out_type);
 	rspamd_http_connection_reset(conn_ent->conn);
 	rspamd_http_router_insert_headers(conn_ent->rt, msg);
 	rspamd_http_connection_write_message(conn_ent->conn, msg, NULL,
-										 "application/json", conn_ent, conn_ent->rt->timeout);
+										 ctype, conn_ent, conn_ent->rt->timeout);
 	conn_ent->is_reply = TRUE;
 }
 
diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index 37a41c111..a251dcd05 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -1636,7 +1636,7 @@ rspamd_protocol_write_ucl(struct rspamd_task *task,
 }
 
 void rspamd_protocol_http_reply(struct rspamd_http_message *msg,
-								struct rspamd_task *task, ucl_object_t **pobj)
+								struct rspamd_task *task, ucl_object_t **pobj, int how)
 {
 	struct rspamd_scan_result *metric_res;
 	const struct rspamd_re_cache_stat *restat;
@@ -1695,7 +1695,7 @@ void rspamd_protocol_http_reply(struct rspamd_http_message *msg,
 
 	if (msg->method < HTTP_SYMBOLS && !RSPAMD_TASK_IS_SPAMC(task)) {
 		msg_debug_protocol("writing json reply");
-		rspamd_ucl_emit_fstring(top, UCL_EMIT_JSON_COMPACT, &reply);
+		rspamd_ucl_emit_fstring(top, how, &reply);
 	}
 	else {
 		if (RSPAMD_TASK_IS_SPAMC(task)) {
@@ -2111,6 +2111,16 @@ void rspamd_protocol_write_reply(struct rspamd_task *task, ev_tstamp timeout)
 						  MESSAGE_FIELD_CHECK(task, message_id));
 	}
 
+	const rspamd_ftok_t *accept_hdr;
+	int out_type = UCL_EMIT_JSON_COMPACT;
+	accept_hdr = rspamd_task_get_request_header(task, "Accept");
+
+	if (accept_hdr && rspamd_substring_search(accept_hdr->begin, accept_hdr->len,
+											  "application/msgpack", sizeof("application/msgpack") - 1) != -1) {
+		ctype = "application/msgpack";
+		out_type = UCL_EMIT_MSGPACK;
+	}
+
 	/* Compatibility */
 	if (task->cmd == CMD_CHECK_RSPAMC) {
 		msg->method = HTTP_SYMBOLS;
@@ -2134,15 +2144,15 @@ void rspamd_protocol_write_reply(struct rspamd_task *task, ev_tstamp timeout)
 							  ucl_object_fromstring(g_quark_to_string(task->err->domain)),
 							  "error_domain", 0, false);
 		reply = rspamd_fstring_sized_new(256);
-		rspamd_ucl_emit_fstring(top, UCL_EMIT_JSON_COMPACT, &reply);
+		rspamd_ucl_emit_fstring(top, out_type, &reply);
 		ucl_object_unref(top);
 
 		/* We also need to validate utf8 */
-		if (rspamd_fast_utf8_validate(reply->str, reply->len) != 0) {
+		if (out_type != UCL_EMIT_MSGPACK && rspamd_fast_utf8_validate(reply->str, reply->len) != 0) {
 			gsize valid_len;
 			gchar *validated;
 
-			/* We copy reply several times here but it should be a rare case */
+			/* We copy reply several times here, but it should be a rare case */
 			validated = rspamd_str_make_utf_valid(reply->str, reply->len,
 												  &valid_len, task->task_pool);
 			rspamd_http_message_set_body(msg, validated, valid_len);
@@ -2161,7 +2171,7 @@ void rspamd_protocol_write_reply(struct rspamd_task *task, ev_tstamp timeout)
 		case CMD_CHECK_SPAMC:
 		case CMD_SKIP:
 		case CMD_CHECK_V2:
-			rspamd_protocol_http_reply(msg, task, NULL);
+			rspamd_protocol_http_reply(msg, task, NULL, out_type);
 			rspamd_protocol_write_log_pipe(task);
 			break;
 		case CMD_PING:
diff --git a/src/libserver/protocol.h b/src/libserver/protocol.h
index 0e3c18744..38d9cef4b 100644
--- a/src/libserver/protocol.h
+++ b/src/libserver/protocol.h
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2024 Vsevolod Stakhov
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 /**
  * @file protocol.h
  * Rspamd protocol definition
@@ -70,7 +86,8 @@ gboolean rspamd_protocol_handle_request(struct rspamd_task *task,
  * @param task
  */
 void rspamd_protocol_http_reply(struct rspamd_http_message *msg,
-								struct rspamd_task *task, ucl_object_t **pobj);
+								struct rspamd_task *task, ucl_object_t **pobj,
+								int how);
 
 /**
  * Write data to log pipes
diff --git a/src/lua/lua_common.c b/src/lua/lua_common.c
index 92e26417a..cf4ca8023 100644
--- a/src/lua/lua_common.c
+++ b/src/lua/lua_common.c
@@ -1974,8 +1974,7 @@ rspamd_lua_check_udata_common(lua_State *L, gint pos, const gchar *classname,
 							  gboolean fatal)
 {
 	void *p = lua_touserdata(L, pos);
-	guint i, top = lua_gettop(L);
-	khiter_t k;
+	gint i, top = lua_gettop(L);
 
 	if (p == NULL) {
 		goto err;
diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c
index 9139866b5..531a27be7 100644
--- a/src/rspamd_proxy.c
+++ b/src/rspamd_proxy.c
@@ -1713,7 +1713,15 @@ rspamd_proxy_scan_self_reply(struct rspamd_task *task)
 	struct rspamd_http_message *msg;
 	struct rspamd_proxy_session *session = task->fin_arg, *nsession;
 	ucl_object_t *rep = NULL;
+	int out_type = UCL_EMIT_JSON_COMPACT;
 	const char *ctype = "application/json";
+	const rspamd_ftok_t *accept_hdr = rspamd_task_get_request_header(task, "Accept");
+
+	if (accept_hdr && rspamd_substring_search(accept_hdr->begin, accept_hdr->len,
+											  "application/msgpack", sizeof("application/msgpack") - 1) != -1) {
+		ctype = "application/msgpack";
+		out_type = UCL_EMIT_MSGPACK;
+	}
 
 	msg = rspamd_http_new_message(HTTP_RESPONSE);
 	msg->date = time(NULL);
@@ -1726,7 +1734,7 @@ rspamd_proxy_scan_self_reply(struct rspamd_task *task)
 	case CMD_CHECK_SPAMC:
 	case CMD_CHECK_V2:
 		rspamd_task_set_finish_time(task);
-		rspamd_protocol_http_reply(msg, task, &rep);
+		rspamd_protocol_http_reply(msg, task, &rep, out_type);
 		rspamd_protocol_write_log_pipe(task);
 		break;
 	case CMD_PING: