[Rework] Move LRU SPF cache from spf plugin

3 files changed, 100 insertions, 97 deletions

diff --git a/src/libserver/spf.c b/src/libserver/spf.c
index 3106a4ffd..0d508ec15 100644
--- a/src/libserver/spf.c
+++ b/src/libserver/spf.c
@@ -67,6 +67,7 @@ struct rspamd_spf_library_ctx {
 	guint max_dns_requests;
 	guint min_cache_ttl;
 	gboolean disable_ipv6;
+	rspamd_lru_hash_t *spf_hash;
 };
 
 struct rspamd_spf_library_ctx *spf_lib_ctx = NULL;
@@ -149,6 +150,9 @@ RSPAMD_CONSTRUCTOR(rspamd_spf_lib_ctx_ctor) {
 }
 
 RSPAMD_DESTRUCTOR(rspamd_spf_lib_ctx_dtor) {
+	if (spf_lib_ctx->spf_hash) {
+		rspamd_lru_hash_destroy (spf_lib_ctx->spf_hash);
+	}
 	g_free (spf_lib_ctx);
 	spf_lib_ctx = NULL;
 }
@@ -188,6 +192,27 @@ spf_library_config (const ucl_object_t *obj)
 		}
 	}
 
+	if ((value = ucl_object_find_key (obj, "disable_ipv6")) != NULL) {
+		if (ucl_object_toboolean_safe (value, &bval)) {
+			spf_lib_ctx->disable_ipv6 = bval;
+		}
+	}
+
+	if ((value = ucl_object_find_key (obj, "spf_cache_size")) != NULL) {
+		if (ucl_object_toint_safe (value, &ival) && ival > 0) {
+			spf_lib_ctx->spf_hash = rspamd_lru_hash_new (
+					ival,
+					NULL,
+					(GDestroyNotify) spf_record_unref);
+		}
+	}
+	else {
+		/* Preserve compatibility */
+		spf_lib_ctx->spf_hash = rspamd_lru_hash_new (
+				2048,
+				NULL,
+				(GDestroyNotify) spf_record_unref);
+	}
 }
 
 static gboolean start_spf_parse (struct spf_record *rec,
@@ -554,10 +579,29 @@ static void
 rspamd_spf_maybe_return (struct spf_record *rec)
 {
 	struct spf_resolved *flat;
+	struct rspamd_task *task = rec->task;
 
 	if (rec->requests_inflight == 0 && !rec->done) {
 		flat = rspamd_spf_record_flatten (rec);
 		rspamd_spf_record_postprocess (flat, rec->task);
+
+		if (flat->ttl > 0 && flat->flags == 0) {
+
+			if (spf_lib_ctx->spf_hash) {
+				rspamd_lru_hash_insert (spf_lib_ctx->spf_hash,
+						flat->domain, spf_record_ref (flat),
+						flat->timestamp, flat->ttl);
+
+				msg_info_task ("stored record for %s (0x%xuL) in LRU cache for %d seconds, "
+							   "%d/%d elements in the cache",
+						flat->domain,
+						flat->digest,
+						flat->ttl,
+						rspamd_lru_hash_size (spf_lib_ctx->spf_hash),
+						rspamd_lru_hash_capacity (spf_lib_ctx->spf_hash));
+			}
+		}
+
 		rec->callback (flat, rec->task, rec->cbdata);
 		REF_RELEASE (flat);
 		rec->done = TRUE;
@@ -2293,13 +2337,7 @@ spf_dns_callback (struct rdns_reply *reply, gpointer arg)
 	rspamd_spf_maybe_return (rec);
 }
 
-struct rspamd_spf_cred {
-	gchar *local_part;
-	gchar *domain;
-	gchar *sender;
-};
-
-struct rspamd_spf_cred *
+static struct rspamd_spf_cred *
 rspamd_spf_cache_domain (struct rspamd_task *task)
 {
 	struct rspamd_email_address *addr;
@@ -2344,10 +2382,9 @@ rspamd_spf_cache_domain (struct rspamd_task *task)
 	return cred;
 }
 
-const gchar *
-rspamd_spf_get_domain (struct rspamd_task *task)
+struct rspamd_spf_cred *
+rspamd_spf_get_cred (struct rspamd_task *task)
 {
-	gchar *domain = NULL;
 	struct rspamd_spf_cred *cred;
 
 	cred = rspamd_mempool_get_variable (task->task_pool,
@@ -2357,6 +2394,17 @@ rspamd_spf_get_domain (struct rspamd_task *task)
 		cred = rspamd_spf_cache_domain (task);
 	}
 
+	return cred;
+}
+
+const gchar *
+rspamd_spf_get_domain (struct rspamd_task *task)
+{
+	gchar *domain = NULL;
+	struct rspamd_spf_cred *cred;
+
+	cred = rspamd_spf_get_cred (task);
+
 	if (cred) {
 		domain = cred->domain;
 	}
@@ -2366,22 +2414,31 @@ rspamd_spf_get_domain (struct rspamd_task *task)
 
 gboolean
 rspamd_spf_resolve (struct rspamd_task *task, spf_cb_t callback,
-		gpointer cbdata)
+		gpointer cbdata, struct rspamd_spf_cred *cred)
 {
 	struct spf_record *rec;
-	struct rspamd_spf_cred *cred;
-
-	cred = rspamd_mempool_get_variable (task->task_pool,
-			RSPAMD_MEMPOOL_SPF_DOMAIN);
-
-	if (!cred) {
-		cred = rspamd_spf_cache_domain (task);
-	}
 
 	if (!cred || !cred->domain) {
 		return FALSE;
 	}
 
+	/* First lookup in the hash */
+	if (spf_lib_ctx->spf_hash) {
+		struct spf_resolved *cached;
+
+		cached = rspamd_lru_hash_lookup (spf_lib_ctx->spf_hash, cred->domain,
+				task->task_timestamp);
+
+		if (cached) {
+			cached->flags |= RSPAMD_SPF_FLAG_CACHED;
+		}
+
+		callback (cached, task, cbdata);
+
+		return TRUE;
+	}
+
+
 	rec = rspamd_mempool_alloc0 (task->task_pool, sizeof (struct spf_record));
 	rec->task = task;
 	rec->callback = callback;
diff --git a/src/libserver/spf.h b/src/libserver/spf.h
index 9fcb01d1e..1a0d8bfcf 100644
--- a/src/libserver/spf.h
+++ b/src/libserver/spf.h
@@ -45,6 +45,7 @@ typedef enum spf_action_e {
 #define RSPAMD_SPF_FLAG_NA (1u << 9u)
 #define RSPAMD_SPF_FLAG_PERMFAIL (1u << 10u)
 #define RSPAMD_SPF_FLAG_RESOLVED (1u << 11u)
+#define RSPAMD_SPF_FLAG_CACHED (1u << 12u)
 
 /** Default SPF limits for avoiding abuse **/
 #define SPF_MAX_NESTING 10
@@ -84,19 +85,26 @@ struct spf_resolved {
 	ref_entry_t ref; /* Refcounting */
 };
 
+struct rspamd_spf_cred {
+	gchar *local_part;
+	gchar *domain;
+	gchar *sender;
+};
 
 /*
  * Resolve spf record for specified task and call a callback after resolution fails/succeed
  */
-gboolean rspamd_spf_resolve (struct rspamd_task *task, spf_cb_t callback,
-							 gpointer cbdata);
+gboolean rspamd_spf_resolve (struct rspamd_task *task,
+							 spf_cb_t callback,
+							 gpointer cbdata,
+							 struct rspamd_spf_cred *cred);
 
 /*
  * Get a domain for spf for specified task
  */
 const gchar *rspamd_spf_get_domain (struct rspamd_task *task);
 
-
+struct rspamd_spf_cred *rspamd_spf_get_cred (struct rspamd_task *task);
 /*
  * Increase refcount
  */
diff --git a/src/plugins/spf.c b/src/plugins/spf.c
index cc52fbd0b..521fa3480 100644
--- a/src/plugins/spf.c
+++ b/src/plugins/spf.c
@@ -58,7 +58,6 @@ struct spf_ctx {
 	const gchar *symbol_permfail;
 
 	struct rspamd_radix_map_helper *whitelist_ip;
-	rspamd_lru_hash_t *spf_hash;
 
 	gboolean check_local;
 	gboolean check_authed;
@@ -321,13 +320,6 @@ spf_module_config (struct rspamd_config *cfg)
 	else {
 		spf_module_ctx->symbol_permfail = DEFAULT_SYMBOL_PERMFAIL;
 	}
-	if ((value =
-		rspamd_config_get_module_opt (cfg, "spf", "spf_cache_size")) != NULL) {
-		cache_size = ucl_obj_toint (value);
-	}
-	else {
-		cache_size = DEFAULT_CACHE_SIZE;
-	}
 
 	spf_library_config (ucl_obj_get_key (cfg->rcl_obj, "spf"));
 
@@ -390,15 +382,6 @@ spf_module_config (struct rspamd_config *cfg)
 			SYMBOL_TYPE_VIRTUAL,
 			cb_id);
 
-	if (cache_size > 0) {
-		spf_module_ctx->spf_hash = rspamd_lru_hash_new (
-				cache_size,
-				NULL,
-				(GDestroyNotify) spf_record_unref);
-		rspamd_mempool_add_destructor (cfg->cfg_pool,
-				(rspamd_mempool_destruct_t)rspamd_lru_hash_destroy,
-				spf_module_ctx->spf_hash);
-	}
 
 	rspamd_mempool_add_destructor (cfg->cfg_pool,
 			(rspamd_mempool_destruct_t)rspamd_map_helper_destroy_radix,
@@ -560,16 +543,12 @@ spf_check_list (struct spf_resolved *rec, struct rspamd_task *task, gboolean cac
 {
 	guint i;
 	struct spf_addr *addr;
-	struct spf_ctx *spf_module_ctx = spf_get_context (task->cfg);
 
 	if (cached) {
-		msg_info_task ("use cached record for %s (0x%xuL) in LRU cache for %d seconds, "
-					   "%d/%d elements in the cache",
+		msg_info_task ("use cached record for %s (0x%xuL) in LRU cache for %d seconds",
 				rec->domain,
 				rec->digest,
-				rec->ttl - (guint)(task->task_timestamp - rec->timestamp),
-				rspamd_lru_hash_size (spf_module_ctx->spf_hash),
-				rspamd_lru_hash_capacity (spf_module_ctx->spf_hash));
+				rec->ttl - (guint)(task->task_timestamp - rec->timestamp));
 	}
 
 	for (i = 0; i < rec->elts->len; i ++) {
@@ -584,7 +563,6 @@ static void
 spf_plugin_callback (struct spf_resolved *record, struct rspamd_task *task,
 		gpointer ud)
 {
-	struct spf_resolved *l = NULL;
 	struct rspamd_symcache_item *item = (struct rspamd_symcache_item *)ud;
 	struct spf_ctx *spf_module_ctx = spf_get_context (task->cfg);
 
@@ -613,37 +591,8 @@ spf_plugin_callback (struct spf_resolved *record, struct rspamd_task *task,
 				NULL);
 	}
 	else if (record && record->domain) {
-
 		spf_record_ref (record);
-
-		if (!spf_module_ctx->spf_hash ||
-			(l = rspamd_lru_hash_lookup (spf_module_ctx->spf_hash,
-					record->domain, task->task_timestamp)) == NULL) {
-			l = record;
-
-			if (record->ttl > 0 && record->flags == 0) {
-
-				if (spf_module_ctx->spf_hash) {
-					rspamd_lru_hash_insert (spf_module_ctx->spf_hash,
-							record->domain, spf_record_ref (l),
-							record->timestamp, record->ttl);
-
-					msg_info_task ("stored record for %s (0x%xuL) in LRU cache for %d seconds, "
-								   "%d/%d elements in the cache",
-							record->domain,
-							record->digest,
-							record->ttl,
-							rspamd_lru_hash_size (spf_module_ctx->spf_hash),
-							rspamd_lru_hash_capacity (spf_module_ctx->spf_hash));
-				}
-			}
-
-		}
-
-		spf_record_ref (l);
-		spf_check_list (l, task, FALSE);
-		spf_record_unref (l);
-
+		spf_check_list (record, task, record->flags & RSPAMD_SPF_FLAG_CACHED);
 		spf_record_unref (record);
 	}
 
@@ -656,8 +605,7 @@ spf_symbol_callback (struct rspamd_task *task,
 					 struct rspamd_symcache_item *item,
 					 void *unused)
 {
-	const gchar *domain;
-	struct spf_resolved *l;
+	struct rspamd_spf_cred *spf_cred;
 	gint *dmarc_checks;
 	struct spf_ctx *spf_module_ctx = spf_get_context (task->cfg);
 
@@ -692,29 +640,19 @@ spf_symbol_callback (struct rspamd_task *task,
 		return;
 	}
 
-	domain = rspamd_spf_get_domain (task);
+	spf_cred = rspamd_spf_get_cred (task);
 	rspamd_symcache_item_async_inc (task, item, M);
 
-	if (domain) {
-		if (spf_module_ctx->spf_hash &&
-				(l = rspamd_lru_hash_lookup (spf_module_ctx->spf_hash, domain,
-					task->task_timestamp)) != NULL) {
-			spf_record_ref (l);
-			spf_check_list (l, task, TRUE);
-			spf_record_unref (l);
+	if (spf_cred && spf_cred->domain) {
+		if (!rspamd_spf_resolve (task, spf_plugin_callback, item, spf_cred)) {
+			msg_info_task ("cannot make spf request for %s", spf_cred->domain);
+			rspamd_task_insert_result (task,
+					spf_module_ctx->symbol_dnsfail,
+					1,
+					"(SPF): spf DNS fail");
 		}
 		else {
-
-			if (!rspamd_spf_resolve (task, spf_plugin_callback, item)) {
-				msg_info_task ("cannot make spf request for %s", domain);
-				rspamd_task_insert_result (task,
-						spf_module_ctx->symbol_dnsfail,
-						1,
-						"(SPF): spf DNS fail");
-			}
-			else {
-				rspamd_symcache_item_async_inc (task, item, M);
-			}
+			rspamd_symcache_item_async_inc (task, item, M);
 		}
 	}