3 files changed, 17 insertions, 21 deletions

diff --git a/src/libserver/spf.c b/src/libserver/spf.c
index 7fae50523..aa777f85c 100644
--- a/src/libserver/spf.c
+++ b/src/libserver/spf.c
@@ -60,6 +60,7 @@ struct spf_record {
 	gchar *local_part;
 	struct rspamd_task *task;
 	spf_cb_t callback;
+	gpointer cbdata;
 	gboolean done;
 };
 
@@ -356,7 +357,7 @@ rspamd_spf_maybe_return (struct spf_record *rec)
 
 	if (rec->requests_inflight == 0 && !rec->done) {
 		flat = rspamd_spf_record_flatten (rec);
-		rec->callback (flat, rec->task);
+		rec->callback (flat, rec->task, rec->cbdata);
 		REF_RELEASE (flat);
 		rec->done = TRUE;
 	}
@@ -1821,7 +1822,8 @@ rspamd_spf_get_domain (struct rspamd_task *task)
 }
 
 gboolean
-resolve_spf (struct rspamd_task *task, spf_cb_t callback)
+rspamd_spf_resolve (struct rspamd_task *task, spf_cb_t callback,
+		gpointer cbdata)
 {
 	struct spf_record *rec;
 	struct rspamd_spf_cred *cred;
@@ -1839,6 +1841,7 @@ resolve_spf (struct rspamd_task *task, spf_cb_t callback)
 	rec = rspamd_mempool_alloc0 (task->task_pool, sizeof (struct spf_record));
 	rec->task = task;
 	rec->callback = callback;
+	rec->cbdata = cbdata;
 
 	rec->resolved = g_ptr_array_sized_new (8);
 
diff --git a/src/libserver/spf.h b/src/libserver/spf.h
index 01850a392..001d3775a 100644
--- a/src/libserver/spf.h
+++ b/src/libserver/spf.h
@@ -8,7 +8,8 @@
 struct rspamd_task;
 struct spf_resolved;
 
-typedef void (*spf_cb_t)(struct spf_resolved *record, struct rspamd_task *task);
+typedef void (*spf_cb_t)(struct spf_resolved *record,
+		struct rspamd_task *task, gpointer cbdata);
 
 typedef enum spf_mech_e {
 	SPF_FAIL,
@@ -66,7 +67,8 @@ struct spf_resolved {
 /*
  * Resolve spf record for specified task and call a callback after resolution fails/succeed
  */
-gboolean resolve_spf (struct rspamd_task *task, spf_cb_t callback);
+gboolean rspamd_spf_resolve (struct rspamd_task *task, spf_cb_t callback,
+		gpointer cbdata);
 
 /*
  * Get a domain for spf for specified task
diff --git a/src/plugins/spf.c b/src/plugins/spf.c
index 5dbc535e4..5a12d3c7c 100644
--- a/src/plugins/spf.c
+++ b/src/plugins/spf.c
@@ -68,18 +68,6 @@ module_t spf_module = {
 	RSPAMD_MODULE_VER
 };
 
-static GQuark
-spf_plugin_quark (void)
-{
-	return g_quark_from_static_string ("spf-plugin");
-}
-
-static void
-spf_plugin_fin (gpointer ud)
-{
-
-}
-
 gint
 spf_module_init (struct rspamd_config *cfg, struct module_ctx **ctx)
 {
@@ -373,9 +361,11 @@ spf_check_list (struct spf_resolved *rec, struct rspamd_task *task)
 }
 
 static void
-spf_plugin_callback (struct spf_resolved *record, struct rspamd_task *task)
+spf_plugin_callback (struct spf_resolved *record, struct rspamd_task *task,
+		gpointer ud)
 {
 	struct spf_resolved *l;
+	struct rspamd_async_watcher *w = ud;
 
 	if (record && record->elts->len > 0 && record->domain) {
 
@@ -394,7 +384,7 @@ spf_plugin_callback (struct spf_resolved *record, struct rspamd_task *task)
 		spf_record_unref (l);
 	}
 
-	rspamd_session_remove_event (task->s, spf_plugin_fin, NULL);
+	rspamd_session_watcher_pop (task->s, w);
 }
 
 
@@ -403,6 +393,7 @@ spf_symbol_callback (struct rspamd_task *task, void *unused)
 {
 	const gchar *domain;
 	struct spf_resolved *l;
+	struct rspamd_async_watcher *w;
 
 	if (radix_find_compressed_addr (spf_module_ctx->whitelist_ip,
 			task->from_addr) != RADIX_NO_VALUE) {
@@ -424,13 +415,13 @@ spf_symbol_callback (struct rspamd_task *task, void *unused)
 			spf_record_unref (l);
 		}
 		else {
-			if (!resolve_spf (task, spf_plugin_callback)) {
+			w = rspamd_session_get_watcher (task->s);
+			if (!rspamd_spf_resolve (task, spf_plugin_callback, w)) {
 				msg_info_task ("cannot make spf request for [%s]",
 						task->message_id);
 			}
 			else {
-				rspamd_session_add_event (task->s, spf_plugin_fin, NULL,
-						spf_plugin_quark ());
+				rspamd_session_watcher_push (task->s);
 			}
 		}
 	}