Exemplo n.º 1
0
static void
dkim_sign_callback (struct rspamd_task *task, void *unused)
{
	lua_State *L;
	struct rspamd_task **ptask;
	gboolean sign = FALSE;
	gint err_idx;
	GString *tb, *hdr;
	GError *err = NULL;
	const gchar *selector = NULL, *domain = NULL, *key = NULL;
	rspamd_dkim_sign_context_t *ctx;
	rspamd_dkim_sign_key_t *dkim_key;

	if (dkim_module_ctx->sign_condition_ref != -1) {
		sign = FALSE;
		L = task->cfg->lua_state;
		lua_pushcfunction (L, &rspamd_lua_traceback);
		err_idx = lua_gettop (L);

		lua_rawgeti (L, LUA_REGISTRYINDEX, dkim_module_ctx->sign_condition_ref);
		ptask = lua_newuserdata (L, sizeof (struct rspamd_task *));
		*ptask = task;
		rspamd_lua_setclass (L, "rspamd{task}", -1);

		if (lua_pcall (L, 1, 1, err_idx) != 0) {
			tb = lua_touserdata (L, -1);
			msg_err_task ("call to user extraction script failed: %v", tb);
			g_string_free (tb, TRUE);
		}
		else {
			if (lua_istable (L, -1)) {
				/*
				 * Get the following elements:
				 * - selector
				 * - domain
				 * - key
				 */
				if (!rspamd_lua_parse_table_arguments (L, -1, &err,
						"*key=S;*domain=S;*selector=S",
						&key, &domain, &selector)) {
					msg_err_task ("invalid return value from sign condition: %e",
							err);
					g_error_free (err);

					return;
				}

				dkim_key = rspamd_lru_hash_lookup (dkim_module_ctx->dkim_sign_hash,
						key, time (NULL));

				if (dkim_key == NULL) {
					dkim_key = rspamd_dkim_sign_key_load (key, &err);

					if (dkim_key == NULL) {
						msg_err_task ("cannot load dkim key %s: %e",
								key, err);
						g_error_free (err);

						return;
					}

					rspamd_lru_hash_insert (dkim_module_ctx->dkim_sign_hash,
							g_strdup (key), dkim_key,
							time (NULL), 0);
				}

				ctx = rspamd_create_dkim_sign_context (task, dkim_key,
						DKIM_CANON_RELAXED, DKIM_CANON_RELAXED,
						dkim_module_ctx->sign_headers, &err);

				if (ctx == NULL) {
					msg_err_task ("cannot create sign context: %e",
							key, err);
					g_error_free (err);

					return;
				}

				hdr = rspamd_dkim_sign (task, selector, domain, 0, 0, ctx);

				if (hdr) {
					rspamd_mempool_set_variable (task->task_pool, "dkim-signature",
							hdr, rspamd_gstring_free_hard);
				}

				sign = TRUE;
			}
			else {
				sign = FALSE;
			}
		}

		/* Result + error function */
		lua_settop (L, 0);

		if (!sign) {
			msg_debug_task ("skip signing as dkim condition callback returned"
					" false");
			return;
		}
	}
}
Exemplo n.º 2
0
gint
lua_dkim_sign_handler (lua_State *L)
{
	struct rspamd_task *task = lua_check_task (L, 1);
	luaL_argcheck (L, lua_type (L, 2) == LUA_TTABLE, 2, "'table' expected");

	GError *err = NULL;
	GString *hdr;
	const gchar *selector = NULL, *domain = NULL, *key = NULL;
	rspamd_dkim_sign_context_t *ctx;
	rspamd_dkim_sign_key_t *dkim_key;
	/*
	 * Get the following elements:
	 * - selector
	 * - domain
	 * - key
	 */
	if (!rspamd_lua_parse_table_arguments (L, 2, &err,
			"*key=S;*domain=S;*selector=S",
			&key, &domain, &selector)) {
		msg_err_task ("invalid return value from sign condition: %e",
				err);
		g_error_free (err);

		lua_pushboolean (L, FALSE);
		return 1;
	}

	if (dkim_module_ctx->dkim_sign_hash == NULL) {
		dkim_module_ctx->dkim_sign_hash = rspamd_lru_hash_new (
				128,
				g_free, /* Keys are just C-strings */
				(GDestroyNotify)rspamd_dkim_sign_key_unref);
	}

	dkim_key = rspamd_lru_hash_lookup (dkim_module_ctx->dkim_sign_hash,
			key, time (NULL));

	if (dkim_key == NULL) {
		dkim_key = rspamd_dkim_sign_key_load (key, &err);

		if (dkim_key == NULL) {
			msg_err_task ("cannot load dkim key %s: %e",
					key, err);
			g_error_free (err);

			lua_pushboolean (L, FALSE);
			return 1;
		}

		rspamd_lru_hash_insert (dkim_module_ctx->dkim_sign_hash,
				g_strdup (key), dkim_key,
				time (NULL), 0);
	}

	ctx = rspamd_create_dkim_sign_context (task, dkim_key,
			DKIM_CANON_RELAXED, DKIM_CANON_RELAXED,
			dkim_module_ctx->sign_headers, &err);

	if (ctx == NULL) {
		msg_err_task ("cannot create sign context: %e",
				err);
		g_error_free (err);

		lua_pushboolean (L, FALSE);
		return 1;
	}

	hdr = rspamd_dkim_sign (task, selector, domain, 0, 0, ctx);

	if (hdr) {
		rspamd_mempool_set_variable (task->task_pool, "dkim-signature",
				hdr, rspamd_gstring_free_hard);
		lua_pushboolean (L, TRUE);
		return 1;
	}

	lua_pushboolean (L, FALSE);
	return 1;
}
Exemplo n.º 3
0
/***
 * @function rspamd_fann.create_full(params)
 * Creates new neural network with parameters:
 * - `layers` {table/numbers}: table of layers in form: {N1, N2, N3 ... Nn} where N is number of neurons in a layer
 * - `activation_hidden` {string}: activation function type for hidden layers (`tanh` by default)
 * - `activation_output` {string}: activation function type for output layer (`tanh` by default)
 * - `sparsed` {float}: create sparsed ANN, where number is a coefficient for sparsing
 * - `learn` {string}: learning algorithm (quickprop, rprop or incremental)
 * - `randomize` {boolean}: randomize weights (true by default)
 * @return {fann} fann object
 */
static gint
lua_fann_create_full (lua_State *L)
{
#ifndef WITH_FANN
	return 0;
#else
	struct fann *f, **pfann;
	guint nlayers, *layers, i;
	const gchar *activation_hidden = NULL, *activation_output, *learn_alg = NULL;
	gdouble sparsed = 0.0;
	gboolean randomize_ann = TRUE;
	GError *err = NULL;

	if (lua_type (L, 1) == LUA_TTABLE) {
		lua_pushstring (L, "layers");
		lua_gettable (L, 1);

		if (lua_type (L, -1) != LUA_TTABLE) {
			return luaL_error (L, "bad layers attribute");
		}

		nlayers = rspamd_lua_table_size (L, -1);
		if (nlayers < 2) {
			return luaL_error (L, "bad layers attribute");
		}

		layers = g_new0 (guint, nlayers);

		for (i = 0; i < nlayers; i ++) {
			lua_rawgeti (L, -1, i + 1);
			layers[i] = luaL_checknumber (L, -1);
			lua_pop (L, 1);
		}

		lua_pop (L, 1); /* Table */

		if (!rspamd_lua_parse_table_arguments (L, 1, &err,
				"sparsed=N;randomize=B;learn=S;activation_hidden=S;activation_output=S",
				&sparsed, &randomize_ann, &learn_alg, &activation_hidden, &activation_output)) {
			g_free (layers);

			if (err) {
				gint r;

				r = luaL_error (L, "invalid arguments: %s", err->message);
				g_error_free (err);
				return r;
			}
			else {
				return luaL_error (L, "invalid arguments");
			}
		}

		if (sparsed != 0.0) {
			f = fann_create_standard_array (nlayers, layers);
		}
		else {
			f = fann_create_sparse_array (sparsed, nlayers, layers);
		}

		if (f != NULL) {
			pfann = lua_newuserdata (L, sizeof (gpointer));
			*pfann = f;
			rspamd_lua_setclass (L, "rspamd{fann}", -1);
		}
		else {
			g_free (layers);
			return luaL_error (L, "cannot create fann");
		}

		fann_set_activation_function_hidden (f,
				string_to_activation_func (activation_hidden));
		fann_set_activation_function_output (f,
				string_to_activation_func (activation_output));
		fann_set_training_algorithm (f, string_to_learn_alg (learn_alg));

		if (randomize_ann) {
			fann_randomize_weights (f, 0, 1);
		}

		g_free (layers);
	}
	else {
		return luaL_error (L, "bad arguments");
	}

	return 1;
#endif
}
Exemplo n.º 4
0
/**
 * @method rspamd_fann:train_threaded(inputs, outputs, callback, event_base, {params})
 * Trains neural network with batch of samples. Inputs and outputs should be tables of
 * equal size, each row in table should be N inputs and M outputs, e.g.
 *     {{0, 1, 1}, ...} -> {{0}, {1} ...}
 * @param {table} inputs input samples
 * @param {table} outputs output samples
 * @param {callback} function that is called when train is completed
 */
static gint
lua_fann_train_threaded (lua_State *L)
{
#ifndef WITH_FANN
	return 0;
#else
	struct fann *f = rspamd_lua_check_fann (L, 1);
	guint ninputs, noutputs, ndata, i, j;
	struct lua_fann_train_cbdata *cbdata;
	struct event_base *ev_base = lua_check_ev_base (L, 5);
	GError *err = NULL;
	const guint max_epochs_default = 1000;
	const gdouble desired_mse_default = 0.0001;

	if (f != NULL && lua_type (L, 2) == LUA_TTABLE &&
			lua_type (L, 3) == LUA_TTABLE && lua_type (L, 4) == LUA_TFUNCTION &&
			ev_base != NULL) {
		/* First check sanity, call for table.getn for that */
		ndata = rspamd_lua_table_size (L, 2);
		ninputs = fann_get_num_input (f);
		noutputs = fann_get_num_output (f);
		cbdata = g_malloc0 (sizeof (*cbdata));
		cbdata->L = L;
		cbdata->f = f;
		cbdata->train = rspamd_fann_create_train (ndata, ninputs, noutputs);
		lua_pushvalue (L, 4);
		cbdata->cbref = luaL_ref (L, LUA_REGISTRYINDEX);

		if (rspamd_socketpair (cbdata->pair, 0) == -1) {
			msg_err ("cannot open socketpair: %s", strerror (errno));
			cbdata->pair[0] = -1;
			cbdata->pair[1] = -1;
			goto err;
		}

		for (i = 0; i < ndata; i ++) {
			lua_rawgeti (L, 2, i + 1);

			if (rspamd_lua_table_size (L, -1) != ninputs) {
				msg_err ("invalid number of inputs: %d, %d expected",
						rspamd_lua_table_size (L, -1), ninputs);
				goto err;
			}

			for (j = 0; j < ninputs; j ++) {
				lua_rawgeti (L, -1, j + 1);
				cbdata->train->input[i][j] = lua_tonumber (L, -1);
				lua_pop (L, 1);
			}

			lua_pop (L, 1);
			lua_rawgeti (L, 3, i + 1);

			if (rspamd_lua_table_size (L, -1) != noutputs) {
				msg_err ("invalid number of outputs: %d, %d expected",
						rspamd_lua_table_size (L, -1), noutputs);
				goto err;
			}

			for (j = 0; j < noutputs; j++) {
				lua_rawgeti (L, -1, j + 1);
				cbdata->train->output[i][j] = lua_tonumber (L, -1);
				lua_pop (L, 1);
			}
		}

		cbdata->max_epochs = max_epochs_default;
		cbdata->desired_mse = desired_mse_default;

		if (lua_type (L, 5) == LUA_TTABLE) {
			rspamd_lua_parse_table_arguments (L, 5, NULL,
					"max_epochs=I;desired_mse=N",
					&cbdata->max_epochs, &cbdata->desired_mse);
		}

		/* Now we can call training in a separate thread */
		rspamd_socket_nonblocking (cbdata->pair[0]);
		event_set (&cbdata->io, cbdata->pair[0], EV_READ, lua_fann_thread_notify,
				cbdata);
		event_base_set (ev_base, &cbdata->io);
		/* TODO: add timeout */
		event_add (&cbdata->io, NULL);
		cbdata->t = rspamd_create_thread ("fann train", lua_fann_train_thread,
				cbdata, &err);

		if (cbdata->t == NULL) {
			msg_err ("cannot create training thread: %e", err);

			if (err) {
				g_error_free (err);
			}

			goto err;
		}
	}
	else {
		return luaL_error (L, "invalid arguments");
	}

	return 0;

err:
	if (cbdata->pair[0] != -1) {
		close (cbdata->pair[0]);
	}
	if (cbdata->pair[1] != -1) {
		close (cbdata->pair[1]);
	}

	fann_destroy_train (cbdata->train);
	luaL_unref (L, LUA_REGISTRYINDEX, cbdata->cbref);
	g_free (cbdata);
	return luaL_error (L, "invalid arguments");
#endif
}
Exemplo n.º 5
0
static gint
lua_worker_spawn_process (lua_State *L)
{
	struct rspamd_worker *w = lua_check_worker (L, 1);
	struct rspamd_lua_process_cbdata *cbdata;
	struct rspamd_abstract_worker_ctx *actx;
	struct rspamd_srv_command srv_cmd;
	const gchar *cmdline = NULL, *input = NULL;
	gsize inputlen = 0;
	pid_t pid;
	GError *err = NULL;
	gint func_cbref, cb_cbref;

	if (!rspamd_lua_parse_table_arguments (L, 2, &err,
			"func=F;exec=S;stdin=V;*on_complete=F", &func_cbref,
			&cmdline, &inputlen, &input, &cb_cbref)) {
		msg_err ("cannot get parameters list: %e", err);

		if (err) {
			g_error_free (err);
		}

		return 0;
	}

	cbdata = g_malloc0 (sizeof (*cbdata));
	cbdata->cb_cbref = cb_cbref;
	cbdata->func_cbref = func_cbref;

	if (input) {
		cbdata->out_buf = g_string_new_len (input, inputlen);
		cbdata->out_pos = 0;
	}

	if (rspamd_socketpair (cbdata->sp, TRUE) == -1) {
		msg_err ("cannot spawn socketpair: %s", strerror (errno));
		luaL_unref (L, LUA_REGISTRYINDEX, cbdata->func_cbref);
		luaL_unref (L, LUA_REGISTRYINDEX, cbdata->cb_cbref);
		g_free (cbdata);

		return 0;
	}

	actx = w->ctx;
	cbdata->wrk = w;
	cbdata->L = L;
	cbdata->ev_base = actx->ev_base;
	cbdata->sz = (guint64)-1;

	pid = fork ();

	if (pid == -1) {
		msg_err ("cannot spawn process: %s", strerror (errno));
		close (cbdata->sp[0]);
		close (cbdata->sp[1]);
		luaL_unref (L, LUA_REGISTRYINDEX, cbdata->func_cbref);
		luaL_unref (L, LUA_REGISTRYINDEX, cbdata->cb_cbref);
		g_free (cbdata);

		return 0;
	}
	else if (pid == 0) {
		/* Child */
		gint rc;
		gchar inbuf[4];

		rspamd_log_update_pid (w->cf->type, w->srv->logger);
		rc = ottery_init (w->srv->cfg->libs_ctx->ottery_cfg);

		if (rc != OTTERY_ERR_NONE) {
			msg_err ("cannot initialize PRNG: %d", rc);
			abort ();
		}
		rspamd_random_seed_fast ();
#ifdef HAVE_EVUTIL_RNG_INIT
		evutil_secure_rng_init ();
#endif

		close (cbdata->sp[0]);
		/* Here we assume that we can block on writing results */
		rspamd_socket_blocking (cbdata->sp[1]);
		event_reinit (cbdata->ev_base);
		g_hash_table_remove_all (w->signal_events);
		rspamd_worker_unblock_signals ();
		rspamd_lua_execute_lua_subprocess (L, cbdata);

		/* Wait for parent to reply and exit */
		rc = read (cbdata->sp[1], inbuf, sizeof (inbuf));

		if (memcmp (inbuf, "\0\0\0\0", 4) == 0) {
			exit (EXIT_SUCCESS);
		}
		else {
			msg_err ("got invalid reply from parent");

			exit (EXIT_FAILURE);
		}

	}

	cbdata->cpid = pid;
	cbdata->io_buf = g_string_sized_new (8);
	/* Notify main */
	memset (&srv_cmd, 0, sizeof (srv_cmd));
	srv_cmd.type = RSPAMD_SRV_ON_FORK;
	srv_cmd.cmd.on_fork.state = child_create;
	srv_cmd.cmd.on_fork.cpid = pid;
	srv_cmd.cmd.on_fork.ppid = getpid ();
	rspamd_srv_send_command (w, cbdata->ev_base, &srv_cmd, -1, NULL, NULL);

	close (cbdata->sp[1]);
	rspamd_socket_nonblocking (cbdata->sp[0]);
	/* Parent */
	rspamd_worker_set_signal_handler (SIGCHLD, w, cbdata->ev_base,
			rspamd_lua_cld_handler,
			cbdata);

	/* Add result pipe waiting */
	event_set (&cbdata->ev, cbdata->sp[0], EV_READ | EV_PERSIST,
			rspamd_lua_subprocess_io, cbdata);
	event_base_set (cbdata->ev_base, &cbdata->ev);
	/* TODO: maybe add timeout? */
	event_add (&cbdata->ev, NULL);

	return 0;
}