[Feature] Enable sessions cache tracking for milter connections

4 files changed, 28 insertions, 11 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