[Feature] Enable sessions cache tracking for milter connections

diff --git a/src/libserver/milter.c b/src/libserver/milter.c
index b543a9b63275f7bf4f88564643308cbfe5008d01..c547f7ea96a81745f2cbcaaacffad1c9c30015c3 100644 (file)
--- a/src/libserver/milter.c
+++ b/src/libserver/milter.c
@@ -26,6 +26,7 @@
 #include "libutil/http_private.h"
 #include "libserver/protocol_internal.h"
 #include "libmime/filter.h"
+#include "libserver/worker_util.h"
 #include "utlist.h"
 
 #define msg_err_milter(...) rspamd_default_log_function(G_LOG_LEVEL_CRITICAL, \
@@ -47,6 +48,7 @@
 
 struct rspamd_milter_context {
        gchar *spam_header;
+       void *sessions_cache;
 };
 
 static struct rspamd_milter_context *milter_ctx = NULL;
@@ -188,6 +190,11 @@ rspamd_milter_session_dtor (struct rspamd_milter_session *session)
                        g_hash_table_destroy (priv->headers);
                }
 
+               if (milter_ctx->sessions_cache) {
+                       rspamd_worker_session_cache_remove (milter_ctx->sessions_cache,
+                                       session);
+               }
+
                rspamd_mempool_delete (priv->pool);
                g_free (priv);
                g_free (session);
@@ -996,11 +1003,7 @@ rspamd_milter_handle_socket (gint fd, const struct timeval *tv,
 
        g_assert (finish_cb != NULL);
        g_assert (error_cb != NULL);
-
-       if (G_UNLIKELY (milter_ctx == NULL)) {
-               milter_ctx = g_malloc (sizeof (*milter_ctx));
-               milter_ctx->spam_header = g_strdup (RSPAMD_MILTER_SPAM_HEADER);
-       }
+       g_assert (milter_ctx != NULL);
 
        session = g_malloc0 (sizeof (*session));
        priv = g_malloc0 (sizeof (*priv));
@@ -1033,6 +1036,11 @@ rspamd_milter_handle_socket (gint fd, const struct timeval *tv,
        session->priv = priv;
        REF_INIT_RETAIN (session, rspamd_milter_session_dtor);
 
+       if (milter_ctx->sessions_cache) {
+               rspamd_worker_session_cache_add (milter_ctx->sessions_cache,
+                               priv->pool->tag.uid, &session->ref.refcount, session);
+       }
+
        return rspamd_milter_handle_session (session, priv);
 }
 
@@ -1629,7 +1637,7 @@ cleanup:
 }
 
 void
-rspamd_milter_init_library (const gchar *spam_header)
+rspamd_milter_init_library (const gchar *spam_header, void *sessions_cache)
 {
        if (milter_ctx) {
                g_free (milter_ctx->spam_header);
@@ -1637,7 +1645,15 @@ rspamd_milter_init_library (const gchar *spam_header)
        }
 
        milter_ctx = g_malloc (sizeof (*milter_ctx));
-       milter_ctx->spam_header = g_strdup (spam_header);
+
+       if (spam_header) {
+               milter_ctx->spam_header = g_strdup (spam_header);
+       }
+       else {
+               milter_ctx->spam_header = g_strdup (RSPAMD_MILTER_SPAM_HEADER);
+       }
+
+       milter_ctx->sessions_cache = sessions_cache;
 }
 
 rspamd_mempool_t *

diff --git a/src/libserver/milter.h b/src/libserver/milter.h
index 3e1c7fff75f51807837071edaadf3342e2a5f6e5..8550ecc9778f9a8046542398d846a9a2b25fec9a 100644 (file)
--- a/src/libserver/milter.h
+++ b/src/libserver/milter.h
@@ -148,7 +148,8 @@ void rspamd_milter_send_task_results (struct rspamd_milter_session *session,
  * Init internal milter context
  * @param spam_header spam header name (must NOT be NULL)
  */
-void rspamd_milter_init_library (const gchar *spam_header);
+void rspamd_milter_init_library (const gchar *spam_header,
+               void *sessions_cache);
 
 /**
  * Returns pool for a session

diff --git a/src/libserver/worker_util.c b/src/libserver/worker_util.c
index 30d1a09936f8efaa43aef9cab428581a0cbf86c7..bb6f9221d54ecf4ab2000bc2f0dc2942764f4614 100644 (file)
--- a/src/libserver/worker_util.c
+++ b/src/libserver/worker_util.c
@@ -697,7 +697,7 @@ rspamd_worker_is_normal (struct rspamd_worker *w)
 
 struct rspamd_worker_session_elt {
        void *ptr;
-       gint *pref;
+       guint *pref;
        const gchar *tag;
        time_t when;
 };
@@ -781,7 +781,7 @@ rspamd_worker_session_cache_new (struct rspamd_worker *w,
 
 void
 rspamd_worker_session_cache_add (void *cache, const gchar *tag,
-               gint *pref, void *ptr)
+               guint *pref, void *ptr)
 {
        struct rspamd_worker_session_cache *c = cache;
        struct rspamd_worker_session_elt *elt;

diff --git a/src/libserver/worker_util.h b/src/libserver/worker_util.h
index 0659b34ff3028341ab9a6d901672e66d7f6371cd..b21a92bfcccda37da5fafced5073563f076c2caf 100644 (file)
--- a/src/libserver/worker_util.h
+++ b/src/libserver/worker_util.h
@@ -152,7 +152,7 @@ void * rspamd_worker_session_cache_new (struct rspamd_worker *w,
  * @param ptr
  */
 void rspamd_worker_session_cache_add (void *cache, const gchar *tag,
-               gint *pref, void *ptr);
+               guint *pref, void *ptr);
 
 /**
  * Removes session from cache

diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c
index f96e008b710d9f3c4adff2f8fc8aeb235f54e174..81ef2a8aa2d1979db87f20adb6c6c6fe362ab17a 100644 (file)
--- a/src/rspamd_proxy.c
+++ b/src/rspamd_proxy.c
@@ -133,6 +133,8 @@ struct rspamd_proxy_ctx {
        gboolean milter;
        /* Milter spam header */
        gchar *spam_header;
+       /* Sessions cache */
+       void *sessions_cache;
 };
 
 enum rspamd_backend_flags {
@@ -986,6 +988,11 @@ proxy_session_dtor (struct rspamd_proxy_session *session)
                close (session->client_sock);
        }
 
+       if (session->ctx->sessions_cache) {
+               rspamd_worker_session_cache_remove (session->ctx->sessions_cache,
+                               session);
+       }
+
        if (session->pool) {
                rspamd_mempool_delete (session->pool);
        }
@@ -1119,6 +1126,11 @@ proxy_session_refresh (struct rspamd_proxy_session *session)
 
        REF_INIT_RETAIN (nsession, proxy_session_dtor);
 
+       if (nsession->ctx->sessions_cache) {
+               rspamd_worker_session_cache_add (nsession->ctx->sessions_cache,
+                               nsession->pool->tag.uid, &nsession->ref.refcount, nsession);
+       }
+
        return nsession;
 }
 
@@ -1951,6 +1963,11 @@ proxy_accept_socket (gint fd, short what, void *arg)
        session->ctx = ctx;
        session->worker = worker;
 
+       if (ctx->sessions_cache) {
+               rspamd_worker_session_cache_add (ctx->sessions_cache,
+                               session->pool->tag.uid, &session->ref.refcount, session);
+       }
+
        if (!ctx->milter) {
                session->client_conn = rspamd_http_connection_new (NULL,
                                proxy_client_error_handler,
@@ -2041,10 +2058,13 @@ start_rspamd_proxy (struct rspamd_worker *worker) {
                rspamd_worker_init_scanner (worker, ctx->ev_base, ctx->resolver);
        }
 
-       if (ctx->spam_header) {
-               rspamd_milter_init_library (ctx->spam_header);
+       if (worker->srv->cfg->enable_sessions_cache) {
+               ctx->sessions_cache = rspamd_worker_session_cache_new (worker,
+                               ctx->ev_base);
        }
 
+       rspamd_milter_init_library (ctx->spam_header, ctx->sessions_cache);
+
        event_base_loop (ctx->ev_base, 0);
        rspamd_worker_block_signals ();