From b14cc05b8ca305bbfacb017ae7c2d005b8354a89 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Tue, 4 Jul 2017 08:44:54 +0100 Subject: [PATCH] [Feature] Enable sessions cache tracking for milter connections --- src/libserver/milter.c | 30 +++++++++++++++++++++++------- src/libserver/milter.h | 3 ++- src/libserver/worker_util.c | 4 ++-- src/libserver/worker_util.h | 2 +- src/rspamd_proxy.c | 24 ++++++++++++++++++++++-- 5 files changed, 50 insertions(+), 13 deletions(-) diff --git a/src/libserver/milter.c b/src/libserver/milter.c index b543a9b63..c547f7ea9 100644 --- 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 3e1c7fff7..8550ecc97 100644 --- 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 30d1a0993..bb6f9221d 100644 --- 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 0659b34ff..b21a92bfc 100644 --- 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 f96e008b7..81ef2a8aa 100644 --- 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 (); -- 2.39.5