[Feature] Allow multiple formats of DKIM signing key

3 files changed, 149 insertions, 91 deletions

diff --git a/src/libserver/dkim.c b/src/libserver/dkim.c
index 9254443ba..29cd21ecb 100644
--- a/src/libserver/dkim.c
+++ b/src/libserver/dkim.c
@@ -121,8 +121,9 @@ struct rspamd_dkim_sign_context_s {
 };
 
 struct rspamd_dkim_sign_key_s {
+	enum rspamd_dkim_sign_key_type type;
 	guint8 *keydata;
-	guint keylen;
+	gsize keylen;
 	RSA *key_rsa;
 	BIO *key_bio;
 	EVP_PKEY *key_evp;
@@ -1037,7 +1038,13 @@ rspamd_dkim_sign_key_free (rspamd_dkim_sign_key_t *key)
 	}
 
 	if (key->keydata && key->keylen > 0) {
-		munmap (key->keydata, key->keylen);
+
+		if (key->type == RSPAMD_DKIM_SIGN_KEY_FILE) {
+			munmap (key->keydata, key->keylen);
+		}
+		else {
+			g_free (key->keydata);
+		}
 	}
 
 	g_slice_free1 (sizeof (rspamd_dkim_sign_key_t), key);
@@ -1976,89 +1983,99 @@ rspamd_dkim_get_dns_key (rspamd_dkim_context_t *ctx)
 }
 
 rspamd_dkim_sign_key_t*
-rspamd_dkim_sign_key_load (const gchar *path, GError **err)
+rspamd_dkim_sign_key_load (const gchar *what, gsize len,
+		enum rspamd_dkim_sign_key_type type,
+		GError **err)
 {
 	gpointer map;
-	gsize len = 0;
+	gsize map_len = 0;
 	rspamd_dkim_sign_key_t *nkey;
+	guint tmp;
 
-	map = rspamd_file_xmap (path, PROT_READ, &len);
+	if (type == RSPAMD_DKIM_SIGN_KEY_FILE) {
+		gchar fpath[PATH_MAX];
 
-	if (map == NULL) {
-		g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
-				"cannot map private key %s: %s",
-				path, strerror (errno));
+		rspamd_snprintf (fpath, sizeof (fpath), "%*s", (gint)len, what);
+		map = rspamd_file_xmap (fpath, PROT_READ, &map_len);
 
-		return NULL;
-	}
+		if (map == NULL) {
+			g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
+					"cannot map private key %s: %s",
+					fpath, strerror (errno));
 
-	nkey = g_slice_alloc0 (sizeof (*nkey));
-	(void)mlock (map, len);
-	nkey->keydata = map;
-	nkey->keylen = len;
-
-	nkey->key_bio = BIO_new_mem_buf (map, len);
-
-	if (!PEM_read_bio_PrivateKey (nkey->key_bio, &nkey->key_evp, NULL, NULL)) {
-		g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
-				"cannot read private key from %s: %s",
-				path, ERR_error_string (ERR_get_error (), NULL));
-		rspamd_dkim_sign_key_free (nkey);
-
-		return NULL;
-	}
-
-	nkey->key_rsa = EVP_PKEY_get1_RSA (nkey->key_evp);
-	if (nkey->key_rsa == NULL) {
-		g_set_error (err,
-				DKIM_ERROR,
-				DKIM_SIGERROR_KEYFAIL,
-				"cannot extract rsa key from evp key");
-		rspamd_dkim_sign_key_free (nkey);
-
-		return NULL;
+			return NULL;
+		}
 	}
 
-	REF_INIT_RETAIN (nkey, rspamd_dkim_sign_key_free);
-
-	return nkey;
-}
-
-rspamd_dkim_sign_key_t*
-rspamd_dkim_sign_key_from_memory (const guchar *data,
-		gsize len, GError **err)
-{
-	rspamd_dkim_sign_key_t *nkey;
-	gpointer map;
+	nkey = g_slice_alloc0 (sizeof (*nkey));
+	nkey->type = type;
+
+	switch (type) {
+	case RSPAMD_DKIM_SIGN_KEY_FILE:
+		(void)mlock (map, len);
+		nkey->keydata = map;
+		nkey->keylen = map_len;
+		break;
+	case RSPAMD_DKIM_SIGN_KEY_BASE64:
+		nkey->keydata = g_malloc (len);
+		nkey->keylen = len;
+
+		if (!rspamd_cryptobox_base64_decode (what, len, nkey->keydata,
+				&nkey->keylen)) {
+			g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
+					"cannot decode base64 encoded private key");
+			rspamd_dkim_sign_key_free (nkey);
 
-	map = mmap (NULL, len, PROT_READ | PROT_WRITE,
-			MAP_PRIVATE | MAP_ANON, -1, 0);
+			return NULL;
+		}
+		break;
+	case RSPAMD_DKIM_SIGN_KEY_DER:
+	case RSPAMD_DKIM_SIGN_KEY_PEM:
+		nkey->keydata = g_malloc (len);
+		memcpy (nkey->keydata, what, len);
+		nkey->keylen = len;
+	}
+
+	(void)mlock (nkey->keydata, nkey->keylen);
+	nkey->key_bio = BIO_new_mem_buf (nkey->keydata, nkey->keylen);
+
+	if (type == RSPAMD_DKIM_SIGN_KEY_DER || type == RSPAMD_DKIM_SIGN_KEY_BASE64) {
+		if (d2i_PrivateKey_bio (nkey->key_bio, &nkey->key_evp) == NULL) {
+			if (type == RSPAMD_DKIM_SIGN_KEY_FILE) {
+				g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
+						"cannot read private key from %*s: %s",
+						(gint)len, what,
+						ERR_error_string (ERR_get_error (), NULL));
+			}
+			else {
+				g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
+						"cannot read private key from string: %s",
+						ERR_error_string (ERR_get_error (), NULL));
+			}
 
-	if (map == MAP_FAILED) {
-		g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
-				"cannot map in-memory private key %s",
-				strerror (errno));
+			rspamd_dkim_sign_key_free (nkey);
 
-		return NULL;
+			return NULL;
+		}
 	}
+	else {
+		if (!PEM_read_bio_PrivateKey (nkey->key_bio, &nkey->key_evp, NULL, NULL)) {
+			if (type == RSPAMD_DKIM_SIGN_KEY_FILE) {
+				g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
+						"cannot read private key from %*s: %s",
+						(gint)len, what,
+						ERR_error_string (ERR_get_error (), NULL));
+			}
+			else {
+				g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
+						"cannot read private key from string: %s",
+						ERR_error_string (ERR_get_error (), NULL));
+			}
 
-	/* Initialize memory and set the appropriate access mode */
-	memcpy (map, data, len);
-	(void)mprotect (map, len, PROT_READ);
-	(void)mlock (map, len);
-
-	nkey = g_slice_alloc0 (sizeof (*nkey));
-	nkey->keydata = map;
-	nkey->keylen = len;
-	nkey->key_bio = BIO_new_mem_buf (map, len);
-
-	if (!PEM_read_bio_PrivateKey (nkey->key_bio, &nkey->key_evp, NULL, NULL)) {
-		g_set_error (err, dkim_error_quark (), DKIM_SIGERROR_KEYFAIL,
-				"cannot read private key from memory: %s",
-				ERR_error_string (ERR_get_error (), NULL));
-		rspamd_dkim_sign_key_free (nkey);
+			rspamd_dkim_sign_key_free (nkey);
 
-		return NULL;
+			return NULL;
+		}
 	}
 
 	nkey->key_rsa = EVP_PKEY_get1_RSA (nkey->key_evp);
diff --git a/src/libserver/dkim.h b/src/libserver/dkim.h
index ac3f04233..a2bd4d0e2 100644
--- a/src/libserver/dkim.h
+++ b/src/libserver/dkim.h
@@ -102,6 +102,13 @@ typedef struct rspamd_dkim_sign_key_s rspamd_dkim_sign_key_t;
 
 struct rspamd_task;
 
+enum rspamd_dkim_sign_key_type {
+	RSPAMD_DKIM_SIGN_KEY_FILE = 0,
+	RSPAMD_DKIM_SIGN_KEY_PEM,
+	RSPAMD_DKIM_SIGN_KEY_BASE64,
+	RSPAMD_DKIM_SIGN_KEY_DER
+};
+
 /* Err MUST be freed if it is not NULL, key is allocated by slice allocator */
 typedef void (*dkim_key_handler_f)(rspamd_dkim_key_t *key, gsize keylen,
 	rspamd_dkim_context_t *ctx, gpointer ud, GError *err);
@@ -134,21 +141,14 @@ rspamd_dkim_sign_context_t * rspamd_create_dkim_sign_context (struct rspamd_task
 		GError **err);
 
 /**
- * Load dkim key from a file
- * @param path
- * @param err
- * @return
- */
-rspamd_dkim_sign_key_t* rspamd_dkim_sign_key_load (const gchar *path, GError **err);
-
-/**
- * Load dkim key from memory chunk
+ * Load dkim key
  * @param path
  * @param err
  * @return
  */
-rspamd_dkim_sign_key_t* rspamd_dkim_sign_key_from_memory (const guchar *data,
-		gsize len, GError **err);
+rspamd_dkim_sign_key_t* rspamd_dkim_sign_key_load (const gchar *what, gsize len,
+		enum rspamd_dkim_sign_key_type type,
+		GError **err);
 
 /**
  * Make DNS request for specified context and obtain and parse key
diff --git a/src/plugins/dkim_check.c b/src/plugins/dkim_check.c
index 87203b5d5..aa51e7592 100644
--- a/src/plugins/dkim_check.c
+++ b/src/plugins/dkim_check.c
@@ -587,7 +587,8 @@ lua_dkim_sign_handler (lua_State *L)
 				key, time (NULL));
 
 		if (dkim_key == NULL) {
-			dkim_key = rspamd_dkim_sign_key_load (key, &err);
+			dkim_key = rspamd_dkim_sign_key_load (key, strlen (key),
+					RSPAMD_DKIM_SIGN_KEY_FILE, &err);
 
 			if (dkim_key == NULL) {
 				msg_err_task ("cannot load dkim key %s: %e",
@@ -614,7 +615,8 @@ lua_dkim_sign_handler (lua_State *L)
 				hex_hash, time (NULL));
 
 		if (dkim_key == NULL) {
-			dkim_key = rspamd_dkim_sign_key_from_memory (rawkey, rawlen, &err);
+			dkim_key = rspamd_dkim_sign_key_load (rawkey, rawlen,
+					RSPAMD_DKIM_SIGN_KEY_PEM, &err);
 
 			if (dkim_key == NULL) {
 				msg_err_task ("cannot load dkim key %s: %e",
@@ -996,11 +998,16 @@ dkim_sign_callback (struct rspamd_task *task, void *unused)
 	struct rspamd_task **ptask;
 	gboolean sign = FALSE;
 	gint err_idx;
+	gsize len;
 	GString *tb, *hdr;
 	GError *err = NULL;
-	const gchar *selector = NULL, *domain = NULL, *key = NULL;
+	const gchar *selector = NULL, *domain = NULL, *key = NULL, *type = NULL,
+			*lru_key;
 	rspamd_dkim_sign_context_t *ctx;
 	rspamd_dkim_sign_key_t *dkim_key;
+	enum rspamd_dkim_sign_key_type sign_type = RSPAMD_DKIM_SIGN_KEY_FILE;
+	guchar h[rspamd_cryptobox_HASHBYTES],
+		hex_hash[rspamd_cryptobox_HASHBYTES * 2 + 1];
 
 	if (dkim_module_ctx->sign_condition_ref != -1) {
 		sign = FALSE;
@@ -1027,8 +1034,8 @@ dkim_sign_callback (struct rspamd_task *task, void *unused)
 				 * - key
 				 */
 				if (!rspamd_lua_parse_table_arguments (L, -1, &err,
-						"*key=S;*domain=S;*selector=S",
-						&key, &domain, &selector)) {
+						"*key=V;*domain=S;*selector=S;type=S",
+						&len, &key, &domain, &selector, &type)) {
 					msg_err_task ("invalid return value from sign condition: %e",
 							err);
 					g_error_free (err);
@@ -1036,11 +1043,44 @@ dkim_sign_callback (struct rspamd_task *task, void *unused)
 					return;
 				}
 
-				dkim_key = rspamd_lru_hash_lookup (dkim_module_ctx->dkim_sign_hash,
-						key, time (NULL));
+				if (type) {
+					if (strcmp (type, "file") == 0) {
+						sign_type = RSPAMD_DKIM_SIGN_KEY_FILE;
+					}
+					else if (strcmp (type, "base64") == 0) {
+						sign_type = RSPAMD_DKIM_SIGN_KEY_BASE64;
+					}
+					else if (strcmp (type, "pem") == 0) {
+						sign_type = RSPAMD_DKIM_SIGN_KEY_PEM;
+					}
+					else if (strcmp (type, "der") == 0 ||
+							strcmp (type, "raw") == 0) {
+						sign_type = RSPAMD_DKIM_SIGN_KEY_DER;
+					}
+				}
+
+				if (sign_type == RSPAMD_DKIM_SIGN_KEY_FILE) {
+
+					dkim_key = rspamd_lru_hash_lookup (
+							dkim_module_ctx->dkim_sign_hash,
+							key, time (NULL));
+					lru_key = key;
+				}
+				else {
+					/* Prehash */
+					memset (hex_hash, 0, sizeof (hex_hash));
+					rspamd_cryptobox_hash (h, key, len, NULL, 0);
+					rspamd_encode_hex_buf (h, sizeof (h),
+							hex_hash, sizeof (hex_hash));
+					dkim_key = rspamd_lru_hash_lookup (
+							dkim_module_ctx->dkim_sign_hash,
+							hex_hash, time (NULL));
+					lru_key = hex_hash;
+				}
 
 				if (dkim_key == NULL) {
-					dkim_key = rspamd_dkim_sign_key_load (key, &err);
+					dkim_key = rspamd_dkim_sign_key_load (key, len,
+							sign_type, &err);
 
 					if (dkim_key == NULL) {
 						msg_err_task ("cannot load dkim key %s: %e",
@@ -1051,7 +1091,7 @@ dkim_sign_callback (struct rspamd_task *task, void *unused)
 					}
 
 					rspamd_lru_hash_insert (dkim_module_ctx->dkim_sign_hash,
-							g_strdup (key), dkim_key,
+							g_strdup (lru_key), dkim_key,
 							time (NULL), 0);
 				}
 
@@ -1070,7 +1110,8 @@ dkim_sign_callback (struct rspamd_task *task, void *unused)
 				hdr = rspamd_dkim_sign (task, selector, domain, 0, 0, ctx);
 
 				if (hdr) {
-					rspamd_mempool_set_variable (task->task_pool, "dkim-signature",
+					rspamd_mempool_set_variable (task->task_pool,
+							"dkim-signature",
 							hdr, rspamd_gstring_free_hard);
 				}