path: root/src/libstat/learn_cache/redis_cache.c

/*-
 * Copyright 2016 Vsevolod Stakhov
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "config.h"
#include "learn_cache.h"
#include "rspamd.h"
#include "stat_api.h"
#include "stat_internal.h"
#include "cryptobox.h"
#include "ucl.h"
#include "hiredis.h"
#include "adapters/libevent.h"

#define REDIS_DEFAULT_TIMEOUT 0.5
#define REDIS_STAT_TIMEOUT 30
#define REDIS_DEFAULT_PORT 6379
#define DEFAULT_REDIS_KEY "learned_ids"

static const gchar *M = "redis learn cache";

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;
};

struct rspamd_redis_cache_runtime {
	struct rspamd_redis_cache_ctx *ctx;
	struct rspamd_task *task;
	struct rspamd_symcache_item *item;
	struct upstream *selected;
	struct event timeout_event;
	redisAsyncContext *redis;
	gboolean has_event;
};

static GQuark
rspamd_stat_cache_redis_quark (void)
{
	return g_quark_from_static_string (M);
}

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)
{
	struct rspamd_redis_cache_runtime *rt = data;
	redisAsyncContext *redis;

	rt->has_event = FALSE;
	if (rspamd_event_pending (&rt->timeout_event, EV_TIMEOUT)) {
		event_del (&rt->timeout_event);
	}

	if (rt->redis) {
		redis = rt->redis;
		rt->redis = NULL;
		/* This calls for all callbacks pending */
		redisAsyncFree (redis);
	}
}

static void
rspamd_redis_cache_timeout (gint fd, short what, gpointer d)
{
	struct rspamd_redis_cache_runtime *rt = d;
	struct rspamd_task *task;

	task = rt->task;

	msg_err_task ("connection to redis server %s timed out",
			rspamd_upstream_name (rt->selected));
	rspamd_upstream_fail (rt->selected, FALSE);

	if (rt->has_event) {
		rspamd_session_remove_event (task->s, rspamd_redis_cache_fin, d);
	}
}

/* Called when we have checked the specified message id */
static void
rspamd_stat_cache_redis_get (redisAsyncContext *c, gpointer r, gpointer priv)
{
	struct rspamd_redis_cache_runtime *rt = priv;
	redisReply *reply = r;
	struct rspamd_task *task;
	glong val = 0;

	task = rt->task;

	if (c->err == 0) {
		if (reply) {
			if (G_LIKELY (reply->type == REDIS_REPLY_INTEGER)) {
				val = reply->integer;
			}
			else if (reply->type == REDIS_REPLY_STRING) {
				rspamd_strtol (reply->str, reply->len, &val);
			}
			else {
				if (reply->type != REDIS_REPLY_NIL) {
					msg_err_task ("bad learned type for %s: %d",
							rt->ctx->stcf->symbol, reply->type);
				}

				val = 0;
			}
		}

		if ((val > 0 && (task->flags & RSPAMD_TASK_FLAG_LEARN_SPAM)) ||
				(val < 0 && (task->flags & RSPAMD_TASK_FLAG_LEARN_HAM))) {
			/* Already learned */
			msg_info_task ("<%s> has been already "
					"learned as %s, ignore it", task->message_id,
					(task->flags & RSPAMD_TASK_FLAG_LEARN_SPAM) ? "spam" : "ham");
			task->flags |= RSPAMD_TASK_FLAG_ALREADY_LEARNED;
		}
		else if (val != 0) {
			/* Unlearn flag */
			task->flags |= RSPAMD_TASK_FLAG_UNLEARN;
		}

		rspamd_upstream_ok (rt->selected);
	}
	else {
		rspamd_upstream_fail (rt->selected, FALSE);
	}

	if (rt->has_event) {
		if (rt->item) {
			rspamd_symcache_item_async_dec_check (task, rt->item, M);
		}
		rspamd_session_remove_event (task->s, rspamd_redis_cache_fin, rt);
	}
}

/* Called when we have learned the specified message id */
static void
rspamd_stat_cache_redis_set (redisAsyncContext *c, gpointer r, gpointer priv)
{
	struct rspamd_redis_cache_runtime *rt = priv;
	struct rspamd_task *task;

	task = rt->task;

	if (c->err == 0) {
		/* XXX: we ignore results here */
		rspamd_upstream_ok (rt->selected);
	}
	else {
		rspamd_upstream_fail (rt->selected, FALSE);
	}

	if (rt->has_event) {
		if (rt->item) {
			rspamd_symcache_item_async_dec_check (task, rt->item, M);
		}
		rspamd_session_remove_event (task->s, rspamd_redis_cache_fin, rt);
	}
}

static void
rspamd_stat_cache_redis_generate_id (struct rspamd_task *task)
{
	rspamd_cryptobox_hash_state_t st;
	rspamd_token_t *tok;
	guint i;
	guchar out[rspamd_cryptobox_HASHBYTES];
	gchar *b32out;
	gchar *user = NULL;

	rspamd_cryptobox_hash_init (&st, NULL, 0);

	user = rspamd_mempool_get_variable (task->task_pool, "stat_user");
	/* Use dedicated hash space for per users cache */
	if (user != NULL) {
		rspamd_cryptobox_hash_update (&st, user, strlen (user));
	}

	for (i = 0; i < task->tokens->len; i ++) {
		tok = g_ptr_array_index (task->tokens, i);
		rspamd_cryptobox_hash_update (&st, (guchar *)&tok->data,
				sizeof (tok->data));
	}

	rspamd_cryptobox_hash_final (&st, out);

	b32out = rspamd_encode_base32 (out, sizeof (out));
	g_assert (b32out != NULL);
	rspamd_mempool_set_variable (task->task_pool, "words_hash", b32out, g_free);
}

static gboolean
rspamd_redis_cache_try_ucl (struct rspamd_redis_cache_ctx *cache_ctx,
		const ucl_object_t *obj,
		struct rspamd_config *cfg,
		const gchar *symbol)
{
	const ucl_object_t *elt, *relt;

	elt = ucl_object_lookup_any (obj, "read_servers", "servers", NULL);

	if (elt == NULL) {
		return FALSE;
	}

	cache_ctx->read_servers = rspamd_upstreams_create (cfg->ups_ctx);
	if (!rspamd_upstreams_from_ucl (cache_ctx->read_servers, elt,
			REDIS_DEFAULT_PORT, NULL)) {
		msg_err ("statfile %s cannot get read servers configuration",
				symbol);
		return FALSE;
	}

	relt = elt;

	elt = ucl_object_lookup (obj, "write_servers");
	if (elt == NULL) {
		/* Use read servers as write ones */
		g_assert (relt != NULL);
		cache_ctx->write_servers = rspamd_upstreams_create (cfg->ups_ctx);
		if (!rspamd_upstreams_from_ucl (cache_ctx->write_servers, relt,
				REDIS_DEFAULT_PORT, NULL)) {
			msg_err ("statfile %s cannot get write servers configuration",
					symbol);
			return FALSE;
		}
	}
	else {
		cache_ctx->write_servers = rspamd_upstreams_create (cfg->ups_ctx);
		if (!rspamd_upstreams_from_ucl (cache_ctx->write_servers, elt,
				REDIS_DEFAULT_PORT, NULL)) {
			msg_err ("statfile %s cannot get write servers configuration",
					symbol);
			rspamd_upstreams_destroy (cache_ctx->write_servers);
			cache_ctx->write_servers = NULL;
		}
	}


	elt = ucl_object_lookup (obj, "timeout");
	if (elt) {
		cache_ctx->timeout = ucl_object_todouble (elt);
	}
	else {
		cache_ctx->timeout = REDIS_DEFAULT_TIMEOUT;
	}

	elt = ucl_object_lookup (obj, "password");
	if (elt) {
		cache_ctx->password = ucl_object_tostring (elt);
	}
	else {
		cache_ctx->password = NULL;
	}

	elt = ucl_object_lookup_any (obj, "db", "database", "dbname", NULL);
	if (elt) {
		if (ucl_object_type (elt) == UCL_STRING) {
			cache_ctx->dbname = ucl_object_tostring (elt);
		}
		else if (ucl_object_type (elt) == UCL_INT) {
			cache_ctx->dbname = ucl_object_tostring_forced (elt);
		}
	}
	else {
		cache_ctx->dbname = NULL;
	}

	elt = ucl_object_lookup_any (obj, "cache_key", "key", NULL);
	if (elt == NULL || ucl_object_type (elt) != UCL_STRING) {
		cache_ctx->redis_object = DEFAULT_REDIS_KEY;
	}
	else {
		cache_ctx->redis_object = ucl_object_tostring (elt);
	}

	return TRUE;
}

gpointer
rspamd_stat_cache_redis_init (struct rspamd_stat_ctx *ctx,
		struct rspamd_config *cfg,
		struct rspamd_statfile *st,
		const ucl_object_t *cf)
{
	struct rspamd_redis_cache_ctx *cache_ctx;
	struct rspamd_statfile_config *stf = st->stcf;
	const ucl_object_t *obj;
	gboolean ret = FALSE;

	cache_ctx = g_malloc0 (sizeof (*cache_ctx));

	/* First search in backend configuration */
	obj = ucl_object_lookup (st->classifier->cfg->opts, "backend");
	if (obj != NULL && ucl_object_type (obj) == UCL_OBJECT) {
		ret = rspamd_redis_cache_try_ucl (cache_ctx, obj, cfg, stf->symbol);
	}

	/* Now try statfiles config */
	if (!ret) {
		ret = rspamd_redis_cache_try_ucl (cache_ctx, stf->opts, cfg, stf->symbol);
	}

	/* Now try classifier config */
	if (!ret) {
		ret = rspamd_redis_cache_try_ucl (cache_ctx, st->classifier->cfg->opts, cfg,
				stf->symbol);
	}

	/* Now try global redis settings */
	if (!ret) {
		obj = ucl_object_lookup (cfg->rcl_obj, "redis");

		if (obj) {
			const ucl_object_t *specific_obj;

			specific_obj = ucl_object_lookup (obj, "statistics");

			if (specific_obj) {
				ret = rspamd_redis_cache_try_ucl (cache_ctx, specific_obj, cfg,
						stf->symbol);
			}
			else {
				ret = rspamd_redis_cache_try_ucl (cache_ctx, obj, cfg,
						stf->symbol);
			}
		}
	}


	if (!ret) {
		msg_err_config ("cannot init redis cache for %s", stf->symbol);
		g_free (cache_ctx);
		return NULL;
	}

	cache_ctx->stcf = stf;

	return (gpointer)cache_ctx;
}

gpointer
rspamd_stat_cache_redis_runtime (struct rspamd_task *task,
		gpointer c, gboolean learn)
{
	struct rspamd_redis_cache_ctx *ctx = c;
	struct rspamd_redis_cache_runtime *rt;
	struct upstream *up;
	rspamd_inet_addr_t *addr;

	g_assert (ctx != NULL);

	if (learn && ctx->write_servers == NULL) {
		msg_err_task ("no write servers defined for %s, cannot learn",
				ctx->stcf->symbol);
		return NULL;
	}

	if (task->tokens == NULL || task->tokens->len == 0) {
		return NULL;
	}

	if (learn) {
		up = rspamd_upstream_get (ctx->write_servers,
				RSPAMD_UPSTREAM_MASTER_SLAVE,
				NULL,
				0);
	}
	else {
		up = rspamd_upstream_get (ctx->read_servers,
				RSPAMD_UPSTREAM_ROUND_ROBIN,
				NULL,
				0);
	}

	if (up == NULL) {
		msg_err_task ("no upstreams reachable");
		return NULL;
	}

	rt = rspamd_mempool_alloc0 (task->task_pool, sizeof (*rt));
	rt->selected = up;
	rt->task = task;
	rt->ctx = ctx;

	addr = rspamd_upstream_addr (up);
	g_assert (addr != NULL);

	if (rspamd_inet_address_get_af (addr) == AF_UNIX) {
		rt->redis = redisAsyncConnectUnix (rspamd_inet_address_to_string (addr));
	}
	else {
		rt->redis = redisAsyncConnect (rspamd_inet_address_to_string (addr),
				rspamd_inet_address_get_port (addr));
	}

	g_assert (rt->redis != NULL);

	redisLibeventAttach (rt->redis, task->ev_base);

	/* 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);
	}

	return rt;
}

gint
rspamd_stat_cache_redis_check (struct rspamd_task *task,
		gboolean is_spam,
		gpointer runtime)
{
	struct rspamd_redis_cache_runtime *rt = runtime;
	struct timeval tv;
	gchar *h;

	if (rspamd_session_blocked (task->s)) {
		return RSPAMD_LEARN_INGORE;
	}

	h = rspamd_mempool_get_variable (task->task_pool, "words_hash");

	if (h == NULL) {
		return RSPAMD_LEARN_INGORE;
	}

	double_to_tv (rt->ctx->timeout, &tv);

	if (redisAsyncCommand (rt->redis, rspamd_stat_cache_redis_get, rt,
			"HGET %s %s",
			rt->ctx->redis_object, h) == REDIS_OK) {
		rspamd_session_add_event (task->s,
				rspamd_redis_cache_fin,
				rt,
				M);
		rt->item = rspamd_symbols_cache_get_cur_item (task);
		event_add (&rt->timeout_event, &tv);
		rt->has_event = TRUE;
	}

	/* We need to return OK every time */
	return RSPAMD_LEARN_OK;
}

gint
rspamd_stat_cache_redis_learn (struct rspamd_task *task,
		gboolean is_spam,
		gpointer runtime)
{
	struct rspamd_redis_cache_runtime *rt = runtime;
	struct timeval tv;
	gchar *h;
	gint flag;

	if (rspamd_session_blocked (task->s)) {
		return RSPAMD_LEARN_INGORE;
	}

	h = rspamd_mempool_get_variable (task->task_pool, "words_hash");
	g_assert (h != NULL);

	double_to_tv (rt->ctx->timeout, &tv);
	flag = (task->flags & RSPAMD_TASK_FLAG_LEARN_SPAM) ? 1 : -1;

	if (redisAsyncCommand (rt->redis, rspamd_stat_cache_redis_set, rt,
			"HSET %s %s %d",
			rt->ctx->redis_object, h, flag) == REDIS_OK) {
		rspamd_session_add_event (task->s,
				rspamd_redis_cache_fin, rt, M);
		rt->item = rspamd_symbols_cache_get_cur_item (task);
		event_add (&rt->timeout_event, &tv);
		rt->has_event = TRUE;
	}

	/* We need to return OK every time */
	return RSPAMD_LEARN_OK;
}

void
rspamd_stat_cache_redis_close (gpointer c)
{

}