Add auth support and db selection for redis stats

2 files changed, 81 insertions, 0 deletions

diff --git a/src/libstat/backends/redis_backend.c b/src/libstat/backends/redis_backend.c
index a862ace6d..6e1767e8d 100644
--- a/src/libstat/backends/redis_backend.c
+++ b/src/libstat/backends/redis_backend.c
@@ -39,6 +39,8 @@ struct redis_stat_ctx {
 	struct upstream_list *write_servers;
 	struct rspamd_stat_async_elt *stat_elt;
 	const gchar *redis_object;
+	const gchar *password;
+	const gchar *dbname;
 	gdouble timeout;
 	gboolean enable_users;
 	gint cbref_user;
@@ -335,6 +337,17 @@ rspamd_redis_expand_object (const gchar *pattern,
 	return tlen;
 }
 
+static void
+rspamd_redis_maybe_auth (struct redis_stat_ctx *ctx, redisAsyncContext *redis)
+{
+	if (ctx->password) {
+		redisAsyncCommand (redis, NULL, NULL, "AUTH %s", ctx->password);
+	}
+	if (ctx->dbname) {
+		redisAsyncCommand (redis, NULL, NULL, "SELECT %s", ctx->dbname);
+	}
+}
+
 static rspamd_fstring_t *
 rspamd_redis_tokens_to_query (struct rspamd_task *task, GPtrArray *tokens,
 		const gchar *arg0, const gchar *arg1, gboolean learn, gint idx,
@@ -643,6 +656,7 @@ rspamd_redis_async_stat_cb (struct rspamd_stat_async_elt *elt, gpointer d)
 
 	/* XXX: deal with timeouts maybe */
 	/* Get keys in redis that match our symbol */
+	rspamd_redis_maybe_auth (ctx, cbdata->redis);
 	redisAsyncCommand (cbdata->redis, rspamd_redis_stat_keys, cbdata,
 			"KEYS %s*",
 			ctx->stcf->symbol);
@@ -962,6 +976,22 @@ rspamd_redis_init (struct rspamd_stat_ctx *ctx,
 		backend->timeout = REDIS_DEFAULT_TIMEOUT;
 	}
 
+	elt = ucl_object_lookup (stf->opts, "password");
+	if (elt) {
+		backend->password = ucl_object_tostring (elt);
+	}
+	else {
+		backend->password = NULL;
+	}
+
+	elt = ucl_object_lookup_any (stf->opts, "db", "database", NULL);
+	if (elt) {
+		backend->dbname = ucl_object_tostring (elt);
+	}
+	else {
+		backend->dbname = NULL;
+	}
+
 	stf->clcf->flags |= RSPAMD_FLAG_CLASSIFIER_INCREMENTING_BACKEND;
 	backend->stcf = stf;
 
@@ -1040,6 +1070,7 @@ rspamd_redis_runtime (struct rspamd_task *task,
 	double_to_tv (ctx->timeout, &tv);
 	event_add (&rt->timeout_event, &tv);
 
+	rspamd_redis_maybe_auth (ctx, rt->redis);
 	redisAsyncCommand (rt->redis, rspamd_redis_connected, rt, "HGET %s %s",
 			rt->redis_object_expanded, "learns");
 
diff --git a/src/libstat/learn_cache/redis_cache.c b/src/libstat/learn_cache/redis_cache.c
index 6dc9d5336..e76e6b2b6 100644
--- a/src/libstat/learn_cache/redis_cache.c
+++ b/src/libstat/learn_cache/redis_cache.c
@@ -32,6 +32,8 @@ struct rspamd_redis_cache_ctx {
 	struct rspamd_statfile_config *stcf;
 	struct upstream_list *read_servers;
 	struct upstream_list *write_servers;
+	const gchar *password;
+	const gchar *dbname;
 	const gchar *redis_object;
 	gdouble timeout;
 };
@@ -50,6 +52,18 @@ rspamd_stat_cache_redis_quark (void)
 	return g_quark_from_static_string ("redis-statistics");
 }
 
+static void
+rspamd_redis_cache_maybe_auth (struct rspamd_redis_cache_ctx *ctx,
+		redisAsyncContext *redis)
+{
+	if (ctx->password) {
+		redisAsyncCommand (redis, NULL, NULL, "AUTH %s", ctx->password);
+	}
+	if (ctx->dbname) {
+		redisAsyncCommand (redis, NULL, NULL, "SELECT %s", ctx->dbname);
+	}
+}
+
 /* Called on connection termination */
 static void
 rspamd_redis_cache_fin (gpointer data)
@@ -198,6 +212,41 @@ rspamd_stat_cache_redis_init (struct rspamd_stat_ctx *ctx,
 
 			return NULL;
 		}
+
+		elt = ucl_object_lookup (st->classifier->cfg->opts, "password");
+		if (elt) {
+			cache_ctx->password = ucl_object_tostring (elt);
+		}
+		else {
+			cache_ctx->password = NULL;
+		}
+
+		elt = ucl_object_lookup_any (st->classifier->cfg->opts,
+				"db", "database", NULL);
+		if (elt) {
+			cache_ctx->dbname = ucl_object_tostring (elt);
+		}
+		else {
+			cache_ctx->dbname = NULL;
+		}
+	}
+	else {
+		elt = ucl_object_lookup (stf->opts, "password");
+		if (elt) {
+			cache_ctx->password = ucl_object_tostring (elt);
+		}
+		else {
+			cache_ctx->password = NULL;
+		}
+
+		elt = ucl_object_lookup_any (stf->opts,
+				"db", "database", NULL);
+		if (elt) {
+			cache_ctx->dbname = ucl_object_tostring (elt);
+		}
+		else {
+			cache_ctx->dbname = NULL;
+		}
 	}
 
 	relt = elt;
@@ -308,6 +357,7 @@ rspamd_stat_cache_redis_runtime (struct rspamd_task *task,
 	/* Now check stats */
 	event_set (&rt->timeout_event, -1, EV_TIMEOUT, rspamd_redis_cache_timeout, rt);
 	event_base_set (task->ev_base, &rt->timeout_event);
+	rspamd_redis_cache_maybe_auth (ctx, rt->redis);
 
 	if (!learn) {
 		rspamd_stat_cache_redis_generate_id (task);