5 files changed, 287 insertions, 66 deletions

diff --git a/src/libserver/maps/map.c b/src/libserver/maps/map.c
index ff3a38f90..a837b11ac 100644
--- a/src/libserver/maps/map.c
+++ b/src/libserver/maps/map.c
@@ -788,72 +788,79 @@ read_map_file (struct rspamd_map *map, struct file_map_data *data,
 	}
 
 	if (len > 0) {
-		if (bk->is_compressed) {
-			bytes = rspamd_file_xmap (data->filename, PROT_READ, &len, TRUE);
-
-			if (bytes == NULL) {
-				msg_err_map ("can't open map %s: %s", data->filename, strerror (errno));
-				return FALSE;
-			}
-
-			ZSTD_DStream *zstream;
-			ZSTD_inBuffer zin;
-			ZSTD_outBuffer zout;
-			guchar *out;
-			gsize outlen, r;
+		if (map->no_file_read) {
+			/* We just call read callback with backend name */
+			map->read_callback (data->filename, strlen (data->filename),
+					&periodic->cbdata, TRUE);
+		}
+		else {
+			if (bk->is_compressed) {
+				bytes = rspamd_file_xmap (data->filename, PROT_READ, &len, TRUE);
 
-			zstream = ZSTD_createDStream ();
-			ZSTD_initDStream (zstream);
+				if (bytes == NULL) {
+					msg_err_map ("can't open map %s: %s", data->filename, strerror (errno));
+					return FALSE;
+				}
 
-			zin.pos = 0;
-			zin.src = bytes;
-			zin.size = len;
+				ZSTD_DStream *zstream;
+				ZSTD_inBuffer zin;
+				ZSTD_outBuffer zout;
+				guchar *out;
+				gsize outlen, r;
 
-			if ((outlen = ZSTD_getDecompressedSize (zin.src, zin.size)) == 0) {
-				outlen = ZSTD_DStreamOutSize ();
-			}
+				zstream = ZSTD_createDStream ();
+				ZSTD_initDStream (zstream);
 
-			out = g_malloc (outlen);
+				zin.pos = 0;
+				zin.src = bytes;
+				zin.size = len;
 
-			zout.dst = out;
-			zout.pos = 0;
-			zout.size = outlen;
+				if ((outlen = ZSTD_getDecompressedSize (zin.src, zin.size)) == 0) {
+					outlen = ZSTD_DStreamOutSize ();
+				}
 
-			while (zin.pos < zin.size) {
-				r = ZSTD_decompressStream (zstream, &zout, &zin);
+				out = g_malloc (outlen);
 
-				if (ZSTD_isError (r)) {
-					msg_err_map ("%s: cannot decompress data: %s",
-							data->filename,
-							ZSTD_getErrorName (r));
-					ZSTD_freeDStream (zstream);
-					g_free (out);
-					munmap (bytes, len);
-					return FALSE;
-				}
+				zout.dst = out;
+				zout.pos = 0;
+				zout.size = outlen;
+
+				while (zin.pos < zin.size) {
+					r = ZSTD_decompressStream (zstream, &zout, &zin);
+
+					if (ZSTD_isError (r)) {
+						msg_err_map ("%s: cannot decompress data: %s",
+								data->filename,
+								ZSTD_getErrorName (r));
+						ZSTD_freeDStream (zstream);
+						g_free (out);
+						munmap (bytes, len);
+						return FALSE;
+					}
 
-				if (zout.pos == zout.size) {
-					/* We need to extend output buffer */
-					zout.size = zout.size * 2 + 1;
-					out = g_realloc (zout.dst, zout.size);
-					zout.dst = out;
+					if (zout.pos == zout.size) {
+						/* We need to extend output buffer */
+						zout.size = zout.size * 2 + 1;
+						out = g_realloc (zout.dst, zout.size);
+						zout.dst = out;
+					}
 				}
-			}
 
-			ZSTD_freeDStream (zstream);
-			msg_info_map ("%s: read map data, %z bytes compressed, "
-					"%z uncompressed)", data->filename,
-					len, zout.pos);
-			map->read_callback (out, zout.pos, &periodic->cbdata, TRUE);
-			g_free (out);
+				ZSTD_freeDStream (zstream);
+				msg_info_map ("%s: read map data, %z bytes compressed, "
+							  "%z uncompressed)", data->filename,
+						len, zout.pos);
+				map->read_callback (out, zout.pos, &periodic->cbdata, TRUE);
+				g_free (out);
 
-			munmap (bytes, len);
-		}
-		else {
-			/* Perform buffered read: fail-safe */
-			if (!read_map_file_chunks (map, &periodic->cbdata, data->filename,
-					len, 0)) {
-				return FALSE;
+				munmap (bytes, len);
+			}
+			else {
+				/* Perform buffered read: fail-safe */
+				if (!read_map_file_chunks (map, &periodic->cbdata, data->filename,
+						len, 0)) {
+					return FALSE;
+				}
 			}
 		}
 	}
@@ -2610,7 +2617,8 @@ rspamd_map_add (struct rspamd_config *cfg,
 				map_fin_cb_t fin_callback,
 				map_dtor_t dtor,
 				void **user_data,
-				struct rspamd_worker *worker)
+				struct rspamd_worker *worker,
+				int flags)
 {
 	struct rspamd_map *map;
 	struct rspamd_map_backend *bk;
@@ -2642,6 +2650,7 @@ rspamd_map_add (struct rspamd_config *cfg,
 			map->backends);
 	g_ptr_array_add (map->backends, bk);
 	map->name = rspamd_mempool_strdup (cfg->cfg_pool, map_line);
+	map->no_file_read = (flags & RSPAMD_MAP_FILE_NO_READ);
 
 	if (bk->protocol == MAP_PROTO_FILE) {
 		map->poll_timeout = (cfg->map_timeout * cfg->map_file_watch_multiplier);
@@ -2685,7 +2694,8 @@ rspamd_map_add_from_ucl (struct rspamd_config *cfg,
 						 map_fin_cb_t fin_callback,
 						 map_dtor_t dtor,
 						 void **user_data,
-						 struct rspamd_worker *worker)
+						 struct rspamd_worker *worker,
+						 gint flags)
 {
 	ucl_object_iter_t it = NULL;
 	const ucl_object_t *cur, *elt;
@@ -2698,7 +2708,7 @@ rspamd_map_add_from_ucl (struct rspamd_config *cfg,
 	if (ucl_object_type (obj) == UCL_STRING) {
 		/* Just a plain string */
 		return rspamd_map_add (cfg, ucl_object_tostring (obj), description,
-				read_callback, fin_callback, dtor, user_data, worker);
+				read_callback, fin_callback, dtor, user_data, worker, flags);
 	}
 
 	map = rspamd_mempool_alloc0 (cfg->cfg_pool, sizeof (struct rspamd_map));
@@ -2712,6 +2722,7 @@ rspamd_map_add_from_ucl (struct rspamd_config *cfg,
 			rspamd_mempool_alloc0_shared (cfg->cfg_pool, sizeof (gint));
 	map->backends = g_ptr_array_new ();
 	map->wrk = worker;
+	map->no_file_read = (flags & RSPAMD_MAP_FILE_NO_READ);
 	rspamd_mempool_add_destructor (cfg->cfg_pool, rspamd_ptr_array_free_hard,
 			map->backends);
 	map->poll_timeout = cfg->map_timeout;
diff --git a/src/libserver/maps/map.h b/src/libserver/maps/map.h
index ce49bacbb..2d3883e11 100644
--- a/src/libserver/maps/map.h
+++ b/src/libserver/maps/map.h
@@ -61,6 +61,12 @@ struct map_cb_data {
  */
 gboolean rspamd_map_is_map (const gchar *map_line);
 
+enum rspamd_map_flags {
+	RSPAMD_MAP_DEFAULT = 0,
+	RSPAMD_MAP_FILE_ONLY = 1u << 0u,
+	RSPAMD_MAP_FILE_NO_READ = 1u << 1u,
+};
+
 /**
  * Add map from line
  */
@@ -71,7 +77,8 @@ struct rspamd_map *rspamd_map_add (struct rspamd_config *cfg,
 								   map_fin_cb_t fin_callback,
 								   map_dtor_t dtor,
 								   void **user_data,
-								   struct rspamd_worker *worker);
+								   struct rspamd_worker *worker,
+								   int flags);
 
 /**
  * Add map from ucl
@@ -83,7 +90,8 @@ struct rspamd_map *rspamd_map_add_from_ucl (struct rspamd_config *cfg,
 											map_fin_cb_t fin_callback,
 											map_dtor_t dtor,
 											void **user_data,
-											struct rspamd_worker *worker);
+											struct rspamd_worker *worker,
+											int flags);
 
 enum rspamd_map_watch_type {
 	RSPAMD_MAP_WATCH_MIN = 9,
diff --git a/src/libserver/maps/map_helpers.c b/src/libserver/maps/map_helpers.c
index d179d44f5..02dae0bea 100644
--- a/src/libserver/maps/map_helpers.c
+++ b/src/libserver/maps/map_helpers.c
@@ -21,6 +21,7 @@
 #include "rspamd.h"
 #include "cryptobox.h"
 #include "contrib/fastutf8/fastutf8.h"
+#include "contrib/cdb/cdb.h"
 
 #ifdef WITH_HYPERSCAN
 #include "hs.h"
@@ -58,6 +59,12 @@ struct rspamd_hash_map_helper {
 	rspamd_cryptobox_fast_hash_state_t hst;
 };
 
+struct rspamd_cdb_map_helper {
+	GQueue cdbs;
+	rspamd_cryptobox_fast_hash_state_t hst;
+	gsize total_size;
+};
+
 struct rspamd_regexp_map_helper {
 	rspamd_mempool_t *pool;
 	struct rspamd_map *map;
@@ -1332,7 +1339,8 @@ rspamd_match_regexp_map_all (struct rspamd_regexp_map_helper *map,
 }
 
 gconstpointer
-rspamd_match_hash_map (struct rspamd_hash_map_helper *map, const gchar *in)
+rspamd_match_hash_map (struct rspamd_hash_map_helper *map, const gchar *in,
+		gsize len)
 {
 	khiter_t k;
 	struct rspamd_map_helper_value *val;
@@ -1394,4 +1402,174 @@ rspamd_match_radix_map_addr (struct rspamd_radix_map_helper *map,
 	}
 
 	return NULL;
+}
+
+
+/*
+ * CBD stuff
+ */
+
+struct rspamd_cdb_map_helper *
+rspamd_map_helper_new_cdb (struct rspamd_map *map)
+{
+	struct rspamd_cdb_map_helper *n;
+
+	n = g_malloc0 (sizeof (*n));
+	n->cdbs = (GQueue)G_QUEUE_INIT;
+
+	rspamd_cryptobox_fast_hash_init (&n->hst, map_hash_seed);
+
+	return n;
+}
+
+void
+rspamd_map_helper_destroy_cdb (struct rspamd_cdb_map_helper *c)
+{
+	if (c == NULL) {
+		return;
+	}
+
+	GList *cur = c->cdbs.head;
+
+	while (cur) {
+		struct cdb *cdb = (struct cdb *)cur->data;
+
+		cdb_free (cdb);
+		g_free (cdb->filename);
+		close (cdb->cdb_fd);
+		g_free (cdb);
+
+		cur = g_list_next (cur);
+	}
+
+	g_queue_clear (&c->cdbs);
+
+	g_free (c);
+}
+
+gchar *
+rspamd_cdb_list_read (gchar *chunk,
+					  gint len,
+					  struct map_cb_data *data,
+					  gboolean final)
+{
+	struct rspamd_cdb_map_helper *cdb_data;
+	struct cdb *found = NULL;
+	struct rspamd_map *map = data->map;
+
+	g_assert (map->no_file_read);
+
+	if (data->cur_data == NULL) {
+		cdb_data = rspamd_map_helper_new_cdb (data->map);
+		data->cur_data = cdb_data;
+	}
+	else {
+		cdb_data = (struct rspamd_cdb_map_helper *)data->cur_data;
+	}
+
+	GList *cur = cdb_data->cdbs.head;
+
+	while (cur) {
+		struct cdb *elt = (struct cdb *)cur->data;
+
+		if (strcmp (elt->filename, chunk) == 0) {
+			found = elt;
+			break;
+		}
+
+		cur = g_list_next (cur);
+	}
+
+	if (found == NULL) {
+		/* New cdb */
+		gint fd;
+		struct cdb *cdb;
+
+		fd = rspamd_file_xopen (chunk, O_RDONLY, 0, TRUE);
+
+		if (fd == -1) {
+			msg_err_map ("cannot open cdb map from %s: %s", chunk, strerror (errno));
+
+			return NULL;
+		}
+
+		cdb = g_malloc0 (sizeof (struct cdb));
+
+		if (cdb_init (cdb, fd) == -1) {
+			msg_err_map ("cannot init cdb map from %s: %s", chunk, strerror (errno));
+
+			return NULL;
+		}
+
+		cdb->filename = g_strdup (chunk);
+		g_queue_push_tail (&cdb_data->cdbs, cdb);
+		cdb_data->total_size += cdb->cdb_fsize;
+		rspamd_cryptobox_fast_hash_update (&cdb_data->hst, chunk, len);
+	}
+
+	return chunk + len;
+}
+
+void
+rspamd_cdb_list_fin (struct map_cb_data *data, void **target)
+{
+	struct rspamd_map *map = data->map;
+	struct rspamd_cdb_map_helper *cdb_data;
+
+	if (data->cur_data) {
+		cdb_data = (struct rspamd_cdb_map_helper *)data->cur_data;
+		msg_info_map ("read cdb of %Hz size", cdb_data->total_size);
+		data->map->traverse_function = NULL;
+		data->map->nelts = 0;
+		data->map->digest = rspamd_cryptobox_fast_hash_final (&cdb_data->hst);
+	}
+
+	if (target) {
+		*target = data->cur_data;
+	}
+
+	if (data->prev_data) {
+		cdb_data = (struct rspamd_cdb_map_helper *)data->prev_data;
+		rspamd_map_helper_destroy_cdb (cdb_data);
+	}
+}
+void
+rspamd_cdb_list_dtor (struct map_cb_data *data)
+{
+	if (data->cur_data) {
+		rspamd_map_helper_destroy_cdb (data->cur_data);
+	}
+}
+
+gconstpointer
+rspamd_match_cdb_map (struct rspamd_cdb_map_helper *map,
+					  const gchar *in, gsize inlen)
+{
+	if (map == NULL || map->cdbs.head == NULL) {
+		return NULL;
+	}
+
+	GList *cur = map->cdbs.head;
+	static rspamd_ftok_t found;
+
+	while (cur) {
+		struct cdb *cdb = (struct cdb *)cur->data;
+
+		if (cdb_find (cdb, in, inlen) > 0) {
+			/* Extract and push value to lua as string */
+			unsigned vlen;
+			gconstpointer vpos;
+
+			vpos = cdb->cdb_mem + cdb_datapos (cdb);
+			vlen = cdb_datalen (cdb);
+			found.len = vlen;
+			found.begin = vpos;
+
+			return &found; /* Do not reuse! */
+		}
+
+		cur = g_list_next (cur);
+	}
+
+	return NULL;
 }
\ No newline at end of file
diff --git a/src/libserver/maps/map_helpers.h b/src/libserver/maps/map_helpers.h
index 4f7b5b804..b5d8cf774 100644
--- a/src/libserver/maps/map_helpers.h
+++ b/src/libserver/maps/map_helpers.h
@@ -38,6 +38,7 @@ extern "C" {
 struct rspamd_radix_map_helper;
 struct rspamd_hash_map_helper;
 struct rspamd_regexp_map_helper;
+struct rspamd_cdb_map_helper;
 struct rspamd_map_helper_value;
 
 enum rspamd_regexp_map_flags {
@@ -76,6 +77,18 @@ void rspamd_kv_list_fin (struct map_cb_data *data, void **target);
 void rspamd_kv_list_dtor (struct map_cb_data *data);
 
 /**
+ * Cdb is a cdb mapped file with shared data
+ * chunk must be filename!
+ */
+gchar *rspamd_cdb_list_read (
+		gchar *chunk,
+		gint len,
+		struct map_cb_data *data,
+		gboolean final);
+void rspamd_cdb_list_fin (struct map_cb_data *data, void **target);
+void rspamd_cdb_list_dtor (struct map_cb_data *data);
+
+/**
  * Regexp list is a list of regular expressions
  */
 
@@ -149,7 +162,17 @@ GPtrArray *rspamd_match_regexp_map_all (struct rspamd_regexp_map_helper *map,
  * @return
  */
 gconstpointer rspamd_match_hash_map (struct rspamd_hash_map_helper *map,
-									 const gchar *in);
+									 const gchar *in, gsize len);
+
+/**
+ * Find value matching specific key in a cdb map
+ * @param map
+ * @param in
+ * @param len
+ * @return rspamd_ftok_t pointer (allocated in a static buffer!)
+ */
+gconstpointer rspamd_match_cdb_map (struct rspamd_cdb_map_helper *map,
+									 const gchar *in, gsize len);
 
 /**
  * Find value matching specific key in a hash map
diff --git a/src/libserver/maps/map_private.h b/src/libserver/maps/map_private.h
index 347f63538..37def8c18 100644
--- a/src/libserver/maps/map_private.h
+++ b/src/libserver/maps/map_private.h
@@ -155,10 +155,11 @@ struct rspamd_map {
 	ev_tstamp timeout;
 	gdouble poll_timeout;
 	time_t next_check;
-	gboolean active_http;
-	gboolean non_trivial; /* E.g. has http backends in active mode */
-	gboolean file_only; /* No HTTP backends found */
-	gboolean static_only; /* No need to check */
+	bool active_http;
+	bool non_trivial; /* E.g. has http backends in active mode */
+	bool file_only; /* No HTTP backends found */
+	bool static_only; /* No need to check */
+	bool no_file_read; /* Do not read files */
 	/* Shared lock for temporary disabling of map reading (e.g. when this map is written by UI) */
 	gint *locked;
 	gchar tag[MEMPOOL_UID_LEN];