Use refcounting logic for dkim keys to avoid some races

3 files changed, 45 insertions, 14 deletions

diff --git a/src/libserver/dkim.c b/src/libserver/dkim.c
index 3b8fc63b5..d1542c1a0 100644
--- a/src/libserver/dkim.c
+++ b/src/libserver/dkim.c
@@ -840,6 +840,7 @@ rspamd_dkim_make_key (rspamd_dkim_context_t *ctx, const gchar *keydata,
 		msg_err_dkim ("DKIM key is too short to be valid");
 		return NULL;
 	}
+
 	key = g_slice_alloc0 (sizeof (rspamd_dkim_key_t));
 	key->keydata = g_slice_alloc (keylen + 1);
 	rspamd_strlcpy (key->keydata, keydata, keylen + 1);
@@ -855,6 +856,8 @@ rspamd_dkim_make_key (rspamd_dkim_context_t *ctx, const gchar *keydata,
 #else
 	g_base64_decode_inplace (key->keydata, &key->decoded_len);
 #endif
+	REF_INIT_RETAIN (key, rspamd_dkim_key_free);
+
 #ifdef HAVE_OPENSSL
 	key->key_bio = BIO_new_mem_buf (key->keydata, key->decoded_len);
 	if (key->key_bio == NULL) {
@@ -862,7 +865,8 @@ rspamd_dkim_make_key (rspamd_dkim_context_t *ctx, const gchar *keydata,
 			DKIM_ERROR,
 			DKIM_SIGERROR_KEYFAIL,
 			"cannot make ssl bio from key");
-		rspamd_dkim_key_free (key);
+		REF_RELEASE (key);
+
 		return NULL;
 	}
 
@@ -872,7 +876,8 @@ rspamd_dkim_make_key (rspamd_dkim_context_t *ctx, const gchar *keydata,
 			DKIM_ERROR,
 			DKIM_SIGERROR_KEYFAIL,
 			"cannot extract pubkey from bio");
-		rspamd_dkim_key_free (key);
+		REF_RELEASE (key);
+
 		return NULL;
 	}
 
@@ -882,10 +887,10 @@ rspamd_dkim_make_key (rspamd_dkim_context_t *ctx, const gchar *keydata,
 			DKIM_ERROR,
 			DKIM_SIGERROR_KEYFAIL,
 			"cannot extract rsa key from evp key");
-		rspamd_dkim_key_free (key);
+		REF_RELEASE (key);
+
 		return NULL;
 	}
-
 #endif
 
 	return key;
@@ -899,12 +904,12 @@ void
 rspamd_dkim_key_free (rspamd_dkim_key_t *key)
 {
 #ifdef HAVE_OPENSSL
-	if (key->key_rsa) {
-		RSA_free (key->key_rsa);
-	}
 	if (key->key_evp) {
 		EVP_PKEY_free (key->key_evp);
 	}
+	if (key->key_rsa) {
+		RSA_free (key->key_rsa);
+	}
 	if (key->key_bio) {
 		BIO_free (key->key_bio);
 	}
@@ -914,8 +919,8 @@ rspamd_dkim_key_free (rspamd_dkim_key_t *key)
 }
 
 static rspamd_dkim_key_t *
-rspamd_dkim_parse_key (rspamd_dkim_context_t *ctx, const gchar *txt, gsize
-*keylen, GError **err)
+rspamd_dkim_parse_key (rspamd_dkim_context_t *ctx, const gchar *txt,
+		gsize *keylen, GError **err)
 {
 	const gchar *c, *p, *end;
 	gint state = 0;
@@ -943,6 +948,11 @@ rspamd_dkim_parse_key (rspamd_dkim_context_t *ctx, const gchar *txt, gsize
 			/* State when we got p= and looking for some public key */
 			if ((*p == ';' || p == end) && p > c) {
 				len = p - c;
+
+				if (keylen) {
+					*keylen = len;
+				}
+
 				return rspamd_dkim_make_key (ctx, c, len, err);
 			}
 			else {
diff --git a/src/libserver/dkim.h b/src/libserver/dkim.h
index 97f54b912..59735aefe 100644
--- a/src/libserver/dkim.h
+++ b/src/libserver/dkim.h
@@ -28,6 +28,7 @@
 #include "config.h"
 #include "event.h"
 #include "dns.h"
+#include "ref.h"
 #ifdef HAVE_OPENSSL
 #include <openssl/rsa.h>
 #include <openssl/engine.h>
@@ -162,8 +163,8 @@ typedef struct rspamd_dkim_key_s {
 	BIO *key_bio;
 	EVP_PKEY *key_evp;
 #endif
-}
-rspamd_dkim_key_t;
+	ref_entry_t ref;
+} rspamd_dkim_key_t;
 
 struct rspamd_task;
 
diff --git a/src/plugins/dkim_check.c b/src/plugins/dkim_check.c
index 8aabe666e..b6b42b579 100644
--- a/src/plugins/dkim_check.c
+++ b/src/plugins/dkim_check.c
@@ -95,6 +95,14 @@ module_t dkim_module = {
 	NULL
 };
 
+static void
+dkim_module_key_dtor (gpointer k)
+{
+	rspamd_dkim_key_t *key = k;
+
+	REF_RELEASE (key);
+}
+
 gint
 dkim_module_init (struct rspamd_config *cfg, struct module_ctx **ctx)
 {
@@ -250,8 +258,8 @@ dkim_module_config (struct rspamd_config *cfg)
 		dkim_module_ctx->dkim_hash = rspamd_lru_hash_new (
 				cache_size,
 				cache_expire,
-				g_free,
-				(GDestroyNotify)rspamd_dkim_key_free);
+				g_free, /* Keys are just C-strings */
+				dkim_module_key_dtor);
 
 		msg_info_config ("init internal dkim module");
 #ifndef HAVE_OPENSSL
@@ -407,11 +415,19 @@ dkim_module_key_handler (rspamd_dkim_key_t *key,
 	struct dkim_check_result *res = ud;
 
 	if (key != NULL) {
-		/* Add new key to the lru cache */
+		/*
+		 * We actually receive key with refcount = 1, so we just assume that
+		 * lru hash owns this object now
+		 */
 		rspamd_lru_hash_insert (dkim_module_ctx->dkim_hash,
 			g_strdup (ctx->dns_key),
 			key, res->task->tv.tv_sec, key->ttl);
+		/* Another ref belongs to the check context */
 		res->key = key;
+		REF_RETAIN (res->key);
+		/* Release key when task is processed */
+		rspamd_mempool_add_destructor (res->task->task_pool,
+				dkim_module_key_dtor, res->key);
 	}
 	else {
 		/* Insert tempfail symbol */
@@ -511,6 +527,10 @@ dkim_symbol_callback (struct rspamd_task *task, void *unused)
 					if (key != NULL) {
 						debug_task ("found key for %s in cache", ctx->dns_key);
 						cur->key = key;
+						REF_RETAIN (cur->key);
+						/* Release key when task is processed */
+						rspamd_mempool_add_destructor (task->task_pool,
+								dkim_module_key_dtor, cur->key);
 					}
 					else {
 						debug_task ("request key for %s from DNS", ctx->dns_key);