From 1755ad1ce5a0165f08f1b4d778f198d652d48b9c Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Thu, 18 Jul 2019 15:05:28 +0100 Subject: [PATCH] [Project] Add possibility to modify body in milter context --- src/libserver/milter.c | 9 ++- src/libserver/milter.h | 4 +- src/libserver/protocol.c | 2 +- src/libserver/protocol_internal.h | 1 + src/rspamd_proxy.c | 91 +++++++++++++++++++++++++++---- 5 files changed, 94 insertions(+), 13 deletions(-) diff --git a/src/libserver/milter.c b/src/libserver/milter.c index 2e00e4c46..50a84a42a 100644 --- a/src/libserver/milter.c +++ b/src/libserver/milter.c @@ -1786,7 +1786,9 @@ rspamd_milter_process_milter_block (struct rspamd_milter_session *session, void rspamd_milter_send_task_results (struct rspamd_milter_session *session, - const ucl_object_t *results) + const ucl_object_t *results, + const gchar *new_body, + gsize bodylen) { const ucl_object_t *elt; struct rspamd_milter_private *priv = session->priv; @@ -1883,6 +1885,11 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session, goto cleanup; } + if (new_body) { + rspamd_milter_send_action (session, RSPAMD_MILTER_REPLBODY, + bodylen, new_body); + } + if (priv->no_action) { msg_info_milter ("do not apply action %s, no_action is set", str_action); diff --git a/src/libserver/milter.h b/src/libserver/milter.h index 10d2c3c47..df2a5efc6 100644 --- a/src/libserver/milter.h +++ b/src/libserver/milter.h @@ -163,7 +163,9 @@ struct rspamd_http_message *rspamd_milter_to_http ( * @param results */ void rspamd_milter_send_task_results (struct rspamd_milter_session *session, - const ucl_object_t *results); + const ucl_object_t *results, + const gchar *new_body, + gsize bodylen); /** * Init internal milter context diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c index 02dc73607..e8bbd979b 100644 --- a/src/libserver/protocol.c +++ b/src/libserver/protocol.c @@ -1564,7 +1564,7 @@ rspamd_protocol_http_reply (struct rspamd_http_message *msg, GString *hdr_offset = g_string_sized_new (30); rspamd_printf_gstring (hdr_offset, "%z", RSPAMD_FSTRING_LEN (reply)); - rspamd_http_message_add_header (msg, "Message-Offset", + rspamd_http_message_add_header (msg, MESSAGE_OFFSET_HEADER, hdr_offset->str); msg_debug_protocol ("write body block at position %s", hdr_offset->str); diff --git a/src/libserver/protocol_internal.h b/src/libserver/protocol_internal.h index 418af70d9..d9616e03d 100644 --- a/src/libserver/protocol_internal.h +++ b/src/libserver/protocol_internal.h @@ -88,6 +88,7 @@ extern "C" { #define CERT_ISSUER_HEADER "TLS-Cert-Issuer" #define MAILER_HEADER "Mailer" #define RAW_DATA_HEADER "Raw" +#define MESSAGE_OFFSET_HEADER "Message-Offset" #ifdef __cplusplus } diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c index 8d1f7c116..cca9f792f 100644 --- a/src/rspamd_proxy.c +++ b/src/rspamd_proxy.c @@ -854,17 +854,36 @@ proxy_backend_close_connection (struct rspamd_proxy_backend_connection *conn) static gboolean proxy_backend_parse_results (struct rspamd_proxy_session *session, - struct rspamd_proxy_backend_connection *conn, - lua_State *L, gint parser_ref, - const gchar *in, gsize inlen) + struct rspamd_proxy_backend_connection *conn, + lua_State *L, gint parser_ref, + struct rspamd_http_message *msg, + goffset *body_offset) { struct ucl_parser *parser; gint err_idx; + const gchar *in = msg->body_buf.begin; + gsize inlen = msg->body_buf.len; + const rspamd_ftok_t *offset_hdr; if (inlen == 0 || in == NULL) { return FALSE; } + offset_hdr = rspamd_http_message_find_header (msg, MESSAGE_OFFSET_HEADER); + + if (offset_hdr) { + gulong val; + + if (rspamd_strtoul (offset_hdr->begin, offset_hdr->len, &val) + && val < inlen) { + + if (body_offset) { + *body_offset = val; + } + inlen = val; + } + } + if (parser_ref != -1) { /* Call parser function */ lua_pushcfunction (L, &rspamd_lua_traceback); @@ -1300,7 +1319,7 @@ proxy_backend_mirror_finish_handler (struct rspamd_http_connection *conn, proxy_request_decompress (msg); if (!proxy_backend_parse_results (session, bk_conn, session->ctx->lua_state, - bk_conn->parser_from_ref, msg->body_buf.begin, msg->body_buf.len)) { + bk_conn->parser_from_ref, msg, NULL)) { msg_warn_session ("cannot parse results from the mirror backend %s:%s", bk_conn->name, rspamd_inet_address_to_string ( @@ -1503,11 +1522,12 @@ proxy_backend_master_error_handler (struct rspamd_http_connection *conn, GError static gint proxy_backend_master_finish_handler (struct rspamd_http_connection *conn, - struct rspamd_http_message *msg) + struct rspamd_http_message *msg) { struct rspamd_proxy_backend_connection *bk_conn = conn->ud; struct rspamd_proxy_session *session, *nsession; rspamd_fstring_t *reply; + goffset body_offset = -1; session = bk_conn->s; rspamd_http_connection_steal_msg (session->master_conn->backend_conn); @@ -1518,7 +1538,7 @@ proxy_backend_master_finish_handler (struct rspamd_http_connection *conn, rspamd_http_connection_reset (session->master_conn->backend_conn); if (!proxy_backend_parse_results (session, bk_conn, session->ctx->lua_state, - bk_conn->parser_from_ref, msg->body_buf.begin, msg->body_buf.len)) { + bk_conn->parser_from_ref, msg, &body_offset)) { msg_warn_session ("cannot parse results from the master backend"); } @@ -1549,8 +1569,17 @@ proxy_backend_master_finish_handler (struct rspamd_http_connection *conn, if (session->client_milter_conn) { nsession = proxy_session_refresh (session); - rspamd_milter_send_task_results (nsession->client_milter_conn, - session->master_conn->results); + + if (body_offset > 0) { + rspamd_milter_send_task_results (nsession->client_milter_conn, + session->master_conn->results, + msg->body_buf.begin + body_offset, + msg->body_buf.len - body_offset); + } + else { + rspamd_milter_send_task_results (nsession->client_milter_conn, + session->master_conn->results, NULL, 0); + } REF_RELEASE (session); rspamd_http_message_free (msg); } @@ -1602,8 +1631,50 @@ rspamd_proxy_scan_self_reply (struct rspamd_task *task) if (session->client_milter_conn) { nsession = proxy_session_refresh (session); - rspamd_milter_send_task_results (nsession->client_milter_conn, - session->master_conn->results); + + if (task->flags & RSPAMD_TASK_FLAG_MESSAGE_REWRITE) { + const gchar *start; + goffset len, hdr_off; + + start = task->msg.begin; + len = task->msg.len; + + hdr_off = MESSAGE_FIELD (task, raw_headers_content).len; + + if (hdr_off < len) { + start += hdr_off; + len -= hdr_off; + + /* The problem here is that we need not end of headers, we need + * start of body. + * + * Hence, we need to skip one \r\n till there is anything else in + * a line. + */ + + if (*start == '\r' && len > 0) { + start++; + len--; + } + + if (*start == '\n' && len > 0) { + start++; + len--; + } + + rspamd_milter_send_task_results (nsession->client_milter_conn, + session->master_conn->results, start, len); + } + else { + /* XXX: should never happen! */ + rspamd_milter_send_task_results (nsession->client_milter_conn, + session->master_conn->results, NULL, 0); + } + } + else { + rspamd_milter_send_task_results (nsession->client_milter_conn, + session->master_conn->results, NULL, 0); + } rspamd_http_message_free (msg); REF_RELEASE (session); } -- 2.39.5