[Project] Add possibility to modify body in milter context

diff --git a/src/libserver/milter.c b/src/libserver/milter.c
index 2e00e4c46f9165500b8aa95af771f68f9a6f797a..50a84a42ad104cd486ea785317f2007b10415b5f 100644 (file)
--- 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 10d2c3c47477ca977dff370a856babf0c07125b8..df2a5efc6555aae24a78bfceec2cb7d847de1d18 100644 (file)
--- 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 02dc7360790bfcfb26d93c36ce1a874905dda2b3..e8bbd979ba7f4dda53ea5aa7a922d36e8d58b215 100644 (file)
--- 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 418af70d965d6edbf0ed249b8ecf5b5723dcdb14..d9616e03dcced7622ab3862747272526ace9ba6c 100644 (file)
--- 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 8d1f7c116bba444dfd9adf170a339dd22f05088b..cca9f792fbd7d56dc2998be12f2d47fc34e2c140 100644 (file)
--- 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);
        }