From 6f2788b45b5982be1257832cb1c2546397117245 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Wed, 22 Dec 2021 20:47:13 +0000
Subject: [PATCH] [Fix] Try to revert back maps content on errors properly

---
 src/libserver/maps/map.c         |  10 +-
 src/libserver/maps/map.h         |   1 +
 src/libserver/maps/map_helpers.c | 164 ++++++++++++++++++++-----------
 src/lua/lua_map.c                |  96 ++++++++++--------
 4 files changed, 169 insertions(+), 102 deletions(-)

diff --git a/src/libserver/maps/map.c b/src/libserver/maps/map.c
index 531a7ce10..4d9b1584b 100644
--- a/src/libserver/maps/map.c
+++ b/src/libserver/maps/map.c
@@ -983,8 +983,8 @@ rspamd_map_periodic_dtor (struct map_periodic_cbdata *periodic)
 	map = periodic->map;
 	msg_debug_map ("periodic dtor %p", periodic);
 
-	if (periodic->need_modify) {
-		/* We are done */
+	if (periodic->need_modify || periodic->cbdata.errored) {
+		/* Need to notify the real data structure */
 		periodic->map->fin_callback (&periodic->cbdata, periodic->map->user_data);
 	}
 	else {
@@ -1138,7 +1138,6 @@ rspamd_map_schedule_periodic (struct rspamd_map *map, int how)
 	}
 
 	cbd = g_malloc0 (sizeof (*cbd));
-	cbd->cbdata.state = 0;
 	cbd->cbdata.prev_data = *map->user_data;
 	cbd->cbdata.cur_data = NULL;
 	cbd->cbdata.map = map;
@@ -2000,7 +1999,7 @@ rspamd_map_process_periodic (struct map_periodic_cbdata *cbd)
 	}
 
 	if (cbd->errored) {
-		/* We should not check other backends if some backend has failed */
+		/* We should not check other backends if some backend has failed*/
 		rspamd_map_schedule_periodic (cbd->map, RSPAMD_MAP_SCHEDULE_ERROR);
 
 		if (cbd->locked) {
@@ -2008,6 +2007,9 @@ rspamd_map_process_periodic (struct map_periodic_cbdata *cbd)
 			cbd->locked = FALSE;
 		}
 
+		/* Also set error flag for the map consumer */
+		cbd->cbdata.errored = true;
+
 		msg_debug_map ("unlocked map %s, refcount=%d", cbd->map->name,
 				cbd->ref.refcount);
 		MAP_RELEASE (cbd, "periodic");
diff --git a/src/libserver/maps/map.h b/src/libserver/maps/map.h
index 0812e1d44..6d77454fb 100644
--- a/src/libserver/maps/map.h
+++ b/src/libserver/maps/map.h
@@ -50,6 +50,7 @@ struct rspamd_map;
 struct map_cb_data {
 	struct rspamd_map *map;
 	gint state;
+	bool errored;
 	void *prev_data;
 	void *cur_data;
 };
diff --git a/src/libserver/maps/map_helpers.c b/src/libserver/maps/map_helpers.c
index a29467497..4eb6b2fee 100644
--- a/src/libserver/maps/map_helpers.c
+++ b/src/libserver/maps/map_helpers.c
@@ -940,22 +940,34 @@ rspamd_kv_list_fin (struct map_cb_data *data, void **target)
 	struct rspamd_map *map = data->map;
 	struct rspamd_hash_map_helper *htb;
 
-	if (data->cur_data) {
-		htb = (struct rspamd_hash_map_helper *)data->cur_data;
-		msg_info_map ("read hash of %d elements from %s", kh_size (htb->htb),
-				map->name);
-		data->map->traverse_function = rspamd_map_helper_traverse_hash;
-		data->map->nelts = kh_size (htb->htb);
-		data->map->digest = rspamd_cryptobox_fast_hash_final (&htb->hst);
+	if (data->errored) {
+		/* Clean up the current data and do not touch prev data */
+		if (data->cur_data) {
+			msg_info_map ("cleanup unfinished new data as error occurred for %s",
+					map->name);
+			htb = (struct rspamd_hash_map_helper *) data->cur_data;
+			rspamd_map_helper_destroy_hash(htb);
+			data->cur_data = NULL;
+		}
 	}
+	else {
+		if (data->cur_data) {
+			htb = (struct rspamd_hash_map_helper *) data->cur_data;
+			msg_info_map ("read hash of %d elements from %s", kh_size(htb->htb),
+					map->name);
+			data->map->traverse_function = rspamd_map_helper_traverse_hash;
+			data->map->nelts = kh_size (htb->htb);
+			data->map->digest = rspamd_cryptobox_fast_hash_final(&htb->hst);
+		}
 
-	if (target) {
-		*target = data->cur_data;
-	}
+		if (target) {
+			*target = data->cur_data;
+		}
 
-	if (data->prev_data) {
-		htb = (struct rspamd_hash_map_helper *)data->prev_data;
-		rspamd_map_helper_destroy_hash (htb);
+		if (data->prev_data) {
+			htb = (struct rspamd_hash_map_helper *) data->prev_data;
+			rspamd_map_helper_destroy_hash(htb);
+		}
 	}
 }
 
@@ -1000,22 +1012,34 @@ rspamd_radix_fin (struct map_cb_data *data, void **target)
 	struct rspamd_map *map = data->map;
 	struct rspamd_radix_map_helper *r;
 
-	if (data->cur_data) {
-		r = (struct rspamd_radix_map_helper *)data->cur_data;
-		msg_info_map ("read radix trie of %z elements: %s",
-				radix_get_size (r->trie), radix_get_info (r->trie));
-		data->map->traverse_function = rspamd_map_helper_traverse_radix;
-		data->map->nelts = kh_size (r->htb);
-		data->map->digest = rspamd_cryptobox_fast_hash_final (&r->hst);
+	if (data->errored) {
+		/* Clean up the current data and do not touch prev data */
+		if (data->cur_data) {
+			msg_info_map ("cleanup unfinished new data as error occurred for %s",
+					map->name);
+			r = (struct rspamd_radix_map_helper *) data->cur_data;
+			rspamd_map_helper_destroy_radix(r);
+			data->cur_data = NULL;
+		}
 	}
+	else {
+		if (data->cur_data) {
+			r = (struct rspamd_radix_map_helper *) data->cur_data;
+			msg_info_map ("read radix trie of %z elements: %s",
+					radix_get_size(r->trie), radix_get_info(r->trie));
+			data->map->traverse_function = rspamd_map_helper_traverse_radix;
+			data->map->nelts = kh_size (r->htb);
+			data->map->digest = rspamd_cryptobox_fast_hash_final(&r->hst);
+		}
 
-	if (target) {
-		*target = data->cur_data;
-	}
+		if (target) {
+			*target = data->cur_data;
+		}
 
-	if (data->prev_data) {
-		r = (struct rspamd_radix_map_helper *)data->prev_data;
-		rspamd_map_helper_destroy_radix (r);
+		if (data->prev_data) {
+			r = (struct rspamd_radix_map_helper *) data->prev_data;
+			rspamd_map_helper_destroy_radix(r);
+		}
 	}
 }
 
@@ -1494,33 +1518,45 @@ rspamd_regexp_list_fin (struct map_cb_data *data, void **target)
 	struct rspamd_regexp_map_helper *re_map = NULL, *old_re_map;
 	struct rspamd_map *map = data->map;
 
-	if (data->cur_data) {
-		re_map = data->cur_data;
-		rspamd_cryptobox_hash_final (&re_map->hst, re_map->re_digest);
-		memcpy (&data->map->digest, re_map->re_digest, sizeof (data->map->digest));
-		rspamd_re_map_finalize (re_map);
-		msg_info_map ("read regexp list of %ud elements",
-				re_map->regexps->len);
-		data->map->traverse_function = rspamd_map_helper_traverse_regexp;
-		data->map->nelts = kh_size (re_map->htb);
+	if (data->errored) {
+		/* Clean up the current data and do not touch prev data */
+		if (data->cur_data) {
+			msg_info_map ("cleanup unfinished new data as error occurred for %s",
+					map->name);
+			re_map = (struct rspamd_regexp_map_helper *)data->cur_data;
+			rspamd_map_helper_destroy_regexp (re_map);
+			data->cur_data = NULL;
+		}
 	}
+	else {
+		if (data->cur_data) {
+			re_map = data->cur_data;
+			rspamd_cryptobox_hash_final(&re_map->hst, re_map->re_digest);
+			memcpy(&data->map->digest, re_map->re_digest, sizeof(data->map->digest));
+			rspamd_re_map_finalize(re_map);
+			msg_info_map ("read regexp list of %ud elements",
+					re_map->regexps->len);
+			data->map->traverse_function = rspamd_map_helper_traverse_regexp;
+			data->map->nelts = kh_size (re_map->htb);
+		}
 
-	if (target) {
-		*target = data->cur_data;
-	}
+		if (target) {
+			*target = data->cur_data;
+		}
 
-	if (data->prev_data) {
-		old_re_map = data->prev_data;
+		if (data->prev_data) {
+			old_re_map = data->prev_data;
 
 #ifdef WITH_HYPERSCAN
-		if (re_map && memcmp (re_map->re_digest, old_re_map->re_digest,
-				sizeof (re_map->re_digest)) != 0) {
-			/* Cleanup old stuff */
-			rspamd_re_map_cache_cleanup_old (old_re_map);
-		}
+			if (re_map && memcmp(re_map->re_digest, old_re_map->re_digest,
+					sizeof(re_map->re_digest)) != 0) {
+				/* Cleanup old stuff */
+				rspamd_re_map_cache_cleanup_old(old_re_map);
+			}
 #endif
 
-		rspamd_map_helper_destroy_regexp (old_re_map);
+			rspamd_map_helper_destroy_regexp(old_re_map);
+		}
 	}
 }
 void
@@ -1889,21 +1925,33 @@ 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 (data->errored) {
+		/* Clean up the current data and do not touch prev data */
+		if (data->cur_data) {
+			msg_info_map ("cleanup unfinished new data as error occurred for %s",
+					map->name);
+			cdb_data = (struct rspamd_cdb_map_helper *) data->cur_data;
+			rspamd_map_helper_destroy_cdb(cdb_data);
+			data->cur_data = NULL;
+		}
 	}
+	else {
+		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 (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);
+		if (data->prev_data) {
+			cdb_data = (struct rspamd_cdb_map_helper *) data->prev_data;
+			rspamd_map_helper_destroy_cdb(cdb_data);
+		}
 	}
 }
 void
diff --git a/src/lua/lua_map.c b/src/lua/lua_map.c
index 72bc00fca..923b9adcc 100644
--- a/src/lua/lua_map.c
+++ b/src/lua/lua_map.c
@@ -430,60 +430,76 @@ lua_map_fin (struct map_cb_data *data, void **target)
 
 	map = data->map;
 
-	if (data->cur_data) {
-		cbdata = (struct lua_map_callback_data *)data->cur_data;
+	if (data->errored) {
+		if (data->cur_data) {
+			cbdata = (struct lua_map_callback_data *)data->cur_data;
+			if (cbdata->ref != -1) {
+				luaL_unref (cbdata->L, LUA_REGISTRYINDEX, cbdata->ref);
+			}
+
+			if (cbdata->data) {
+				rspamd_fstring_free (cbdata->data);
+			}
+
+			data->cur_data = NULL;
+		}
 	}
 	else {
-		msg_err_map ("no data read for map");
-		return;
-	}
+		if (data->cur_data) {
+			cbdata = (struct lua_map_callback_data *) data->cur_data;
+		}
+		else {
+			msg_err_map ("no data read for map");
+			return;
+		}
 
-	if (cbdata->ref == -1) {
-		msg_err_map ("map has no callback set");
-	}
-	else if (cbdata->data != NULL && cbdata->data->len != 0) {
+		if (cbdata->ref == -1) {
+			msg_err_map ("map has no callback set");
+		}
+		else if (cbdata->data != NULL && cbdata->data->len != 0) {
 
-		lua_pushcfunction (cbdata->L, &rspamd_lua_traceback);
-		int err_idx = lua_gettop (cbdata->L);
+			lua_pushcfunction (cbdata->L, &rspamd_lua_traceback);
+			int err_idx = lua_gettop(cbdata->L);
 
-		lua_rawgeti (cbdata->L, LUA_REGISTRYINDEX, cbdata->ref);
+			lua_rawgeti(cbdata->L, LUA_REGISTRYINDEX, cbdata->ref);
 
-		if (!cbdata->opaque) {
-			lua_pushlstring (cbdata->L, cbdata->data->str, cbdata->data->len);
-		}
-		else {
-			struct rspamd_lua_text *t;
+			if (!cbdata->opaque) {
+				lua_pushlstring(cbdata->L, cbdata->data->str, cbdata->data->len);
+			}
+			else {
+				struct rspamd_lua_text *t;
 
-			t = lua_newuserdata (cbdata->L, sizeof (*t));
-			rspamd_lua_setclass (cbdata->L, "rspamd{text}", -1);
-			t->flags = 0;
-			t->len = cbdata->data->len;
-			t->start = cbdata->data->str;
-		}
+				t = lua_newuserdata(cbdata->L, sizeof(*t));
+				rspamd_lua_setclass(cbdata->L, "rspamd{text}", -1);
+				t->flags = 0;
+				t->len = cbdata->data->len;
+				t->start = cbdata->data->str;
+			}
 
-		pmap = lua_newuserdata (cbdata->L, sizeof (void *));
-		*pmap = cbdata->lua_map;
-		rspamd_lua_setclass (cbdata->L, "rspamd{map}", -1);
+			pmap = lua_newuserdata(cbdata->L, sizeof(void *));
+			*pmap = cbdata->lua_map;
+			rspamd_lua_setclass(cbdata->L, "rspamd{map}", -1);
 
-		gint ret = lua_pcall (cbdata->L, 2, 0, err_idx);
+			gint ret = lua_pcall(cbdata->L, 2, 0, err_idx);
 
-		if (ret != 0) {
-			msg_info_map ("call to %s failed (%d): %s", "map fin function",
-				ret,
-				lua_tostring (cbdata->L, -1));
-		}
+			if (ret != 0) {
+				msg_info_map ("call to %s failed (%d): %s", "map fin function",
+						ret,
+						lua_tostring(cbdata->L, -1));
+			}
 
-		lua_settop (cbdata->L, err_idx - 1);
-	}
+			lua_settop(cbdata->L, err_idx - 1);
+		}
 
-	cbdata->data = rspamd_fstring_assign (cbdata->data, "", 0);
+		cbdata->data = rspamd_fstring_assign(cbdata->data, "", 0);
 
-	if (target) {
-		*target = data->cur_data;
-	}
+		if (target) {
+			*target = data->cur_data;
+		}
 
-	if (data->prev_data) {
-		data->prev_data = NULL;
+		if (data->prev_data) {
+			data->prev_data = NULL;
+		}
 	}
 }
 
-- 
2.39.5