[Feature] Allow forced and non-forced upstreams selection

2 files changed, 35 insertions, 4 deletions

diff --git a/src/libutil/upstream.c b/src/libutil/upstream.c
index 9187ba9b9..f22092289 100644
--- a/src/libutil/upstream.c
+++ b/src/libutil/upstream.c
@@ -786,10 +786,10 @@ rspamd_upstream_get_hashed (struct upstream_list *ups, const guint8 *key, guint
 	return g_ptr_array_index (ups->alive, idx);
 }
 
-struct upstream*
-rspamd_upstream_get (struct upstream_list *ups,
+static struct upstream*
+rspamd_upstream_get_common (struct upstream_list *ups,
 		enum rspamd_upstream_rotation default_type,
-		const guchar *key, gsize keylen)
+		const guchar *key, gsize keylen, gboolean forced)
 {
 	enum rspamd_upstream_rotation type;
 
@@ -800,7 +800,12 @@ rspamd_upstream_get (struct upstream_list *ups,
 	}
 	rspamd_mutex_unlock (ups->lock);
 
-	type = ups->rot_alg != RSPAMD_UPSTREAM_UNDEF ? ups->rot_alg : default_type;
+	if (!forced) {
+		type = ups->rot_alg != RSPAMD_UPSTREAM_UNDEF ? ups->rot_alg : default_type;
+	}
+	else {
+		type = default_type != RSPAMD_UPSTREAM_UNDEF ? default_type : ups->rot_alg;
+	}
 
 	if (type == RSPAMD_UPSTREAM_HASHED && (keylen == 0 || key == NULL)) {
 		/* Cannot use hashed rotation when no key is specified, switch to random */
@@ -830,6 +835,22 @@ rspamd_upstream_get (struct upstream_list *ups,
 	return NULL;
 }
 
+struct upstream*
+rspamd_upstream_get (struct upstream_list *ups,
+		enum rspamd_upstream_rotation default_type,
+		const guchar *key, gsize keylen)
+{
+	return rspamd_upstream_get_common (ups, default_type, key, keylen, FALSE);
+}
+
+struct upstream*
+rspamd_upstream_get_forced (struct upstream_list *ups,
+		enum rspamd_upstream_rotation forced_type,
+		const guchar *key, gsize keylen)
+{
+	return rspamd_upstream_get_common (ups, forced_type, key, keylen, TRUE);
+}
+
 void
 rspamd_upstream_reresolve (struct upstream_ctx *ctx)
 {
diff --git a/src/libutil/upstream.h b/src/libutil/upstream.h
index c3cbf685a..07b9302e8 100644
--- a/src/libutil/upstream.h
+++ b/src/libutil/upstream.h
@@ -191,6 +191,16 @@ struct upstream* rspamd_upstream_get (struct upstream_list *ups,
 		const guchar *key, gsize keylen);
 
 /**
+ * Get new upstream from the list
+ * @param ups upstream list
+ * @param type type of rotation algorithm, for `RSPAMD_UPSTREAM_HASHED` it is required to specify `key` and `keylen` as arguments
+ * @return
+ */
+struct upstream* rspamd_upstream_get_forced (struct upstream_list *ups,
+		enum rspamd_upstream_rotation forced_type,
+		const guchar *key, gsize keylen);
+
+/**
  * Re-resolve addresses for all upstreams registered
  */
 void rspamd_upstream_reresolve (struct upstream_ctx *ctx);