diff options
Diffstat (limited to 'src/lua/lua_task.c')
| -rw-r--r-- | src/lua/lua_task.c | 188 |
1 files changed, 181 insertions, 7 deletions
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c index 355680881..0b1473b61 100644 --- a/src/lua/lua_task.c +++ b/src/lua/lua_task.c @@ -1,5 +1,5 @@ /* - * Copyright 2024 Vsevolod Stakhov + * Copyright 2025 Vsevolod Stakhov * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -1226,6 +1226,13 @@ LUA_FUNCTION_DEF(task, get_all_named_results); */ LUA_FUNCTION_DEF(task, get_dns_req); +/*** + * @method task:add_timer(timeout, callback) + * Creates a delayed execution task for the specific callback at given timeout (in seconds) + * + */ +LUA_FUNCTION_DEF(task, add_timer); + static const struct luaL_reg tasklib_f[] = { LUA_INTERFACE_DEF(task, create), LUA_INTERFACE_DEF(task, load_from_file), @@ -1353,6 +1360,7 @@ static const struct luaL_reg tasklib_m[] = { LUA_INTERFACE_DEF(task, add_named_result), LUA_INTERFACE_DEF(task, get_all_named_results), LUA_INTERFACE_DEF(task, topointer), + LUA_INTERFACE_DEF(task, add_timer), {"__tostring", rspamd_lua_class_tostring}, {NULL, NULL}}; @@ -6935,7 +6943,7 @@ lua_task_get_meta_words(lua_State *L) return luaL_error(L, "invalid arguments"); } - if (task->meta_words == NULL) { + if (!task->meta_words.a) { lua_createtable(L, 0, 0); } else { @@ -6959,7 +6967,7 @@ lua_task_get_meta_words(lua_State *L) } } - return rspamd_lua_push_words(L, task->meta_words, how); + return rspamd_lua_push_words_kvec(L, &task->meta_words, how); } return 1; @@ -7031,6 +7039,76 @@ lua_lookup_words_array(lua_State *L, return nmatched; } +static unsigned int +lua_lookup_words_kvec(lua_State *L, + int cbpos, + struct rspamd_task *task, + struct rspamd_lua_map *map, + rspamd_words_t *words) +{ + rspamd_word_t *tok; + unsigned int i, nmatched = 0; + int err_idx; + gboolean matched; + const char *key; + gsize keylen; + + if (!words || !words->a) { + return 0; + } + + for (i = 0; i < kv_size(*words); i++) { + tok = &kv_A(*words, i); + + matched = FALSE; + + if (tok->normalized.len == 0) { + continue; + } + + key = tok->normalized.begin; + keylen = tok->normalized.len; + + switch (map->type) { + case RSPAMD_LUA_MAP_SET: + case RSPAMD_LUA_MAP_HASH: + /* We know that tok->normalized is zero terminated in fact */ + if (rspamd_match_hash_map(map->data.hash, key, keylen)) { + matched = TRUE; + } + break; + case RSPAMD_LUA_MAP_REGEXP: + case RSPAMD_LUA_MAP_REGEXP_MULTIPLE: + if (rspamd_match_regexp_map_single(map->data.re_map, key, + keylen)) { + matched = TRUE; + } + break; + default: + g_assert_not_reached(); + break; + } + + if (matched) { + nmatched++; + + lua_pushcfunction(L, &rspamd_lua_traceback); + err_idx = lua_gettop(L); + lua_pushvalue(L, cbpos); /* Function */ + rspamd_lua_push_full_word(L, tok); + + if (lua_pcall(L, 1, 0, err_idx) != 0) { + msg_err_task("cannot call callback function for lookup words: %s", + lua_tostring(L, -1)); + } + + lua_settop(L, err_idx - 1); + } + } + + return nmatched; +} + static int lua_task_lookup_words(lua_State *L) { @@ -7054,13 +7132,13 @@ lua_task_lookup_words(lua_State *L) PTR_ARRAY_FOREACH(MESSAGE_FIELD(task, text_parts), i, tp) { - if (tp->utf_words) { - matches += lua_lookup_words_array(L, 3, task, map, tp->utf_words); + if (tp->utf_words.a) { + matches += lua_lookup_words_kvec(L, 3, task, map, &tp->utf_words); } } - if (task->meta_words) { - matches += lua_lookup_words_array(L, 3, task, map, task->meta_words); + if (task->meta_words.a) { + matches += lua_lookup_words_kvec(L, 3, task, map, &task->meta_words); } lua_pushinteger(L, matches); @@ -7406,6 +7484,102 @@ lua_archive_get_filename(lua_State *L) return 1; } +struct rspamd_task_timer_cbdata { + lua_State *L; + struct rspamd_task *task; + struct rspamd_symcache_dynamic_item *item; + struct rspamd_async_event *async_ev; + int cbref; + ev_timer ev; +}; + +static void +lua_timer_fin(gpointer arg) +{ + struct rspamd_task_timer_cbdata *cbdata = (struct rspamd_task_timer_cbdata *) arg; + + ev_timer_stop(cbdata->task->event_loop, &cbdata->ev); + luaL_unref(cbdata->L, LUA_REGISTRYINDEX, cbdata->cbref); +} + +static void +lua_task_timer_cb(struct ev_loop *loop, struct ev_timer *t, int events) +{ + struct rspamd_task_timer_cbdata *cbdata = (struct rspamd_task_timer_cbdata *) t->data; + lua_State *L; + bool schedule_more = false; + + L = cbdata->L; + + lua_pushcfunction(L, &rspamd_lua_traceback); + int err_idx = lua_gettop(L); + lua_rawgeti(L, LUA_REGISTRYINDEX, cbdata->cbref); + rspamd_lua_task_push(L, cbdata->task); + + if (lua_pcall(L, 1, 1, err_idx) != 0) { + msg_err("call to periodic " + "script failed: %s", + lua_tostring(L, -1)); + } + else { + if (lua_isnumber(L, -1)) { + schedule_more = true; + ev_timer_set(&cbdata->ev, lua_tonumber(L, -1), 0.0); + } + } + + if (schedule_more) { + ev_timer_again(loop, t); + } + else { + /* Cleanup */ + if (cbdata->item) { + rspamd_symcache_item_async_dec_check(cbdata->task, cbdata->item, "timer"); + cbdata->item = NULL; + } + rspamd_session_remove_event(cbdata->task->s, lua_timer_fin, cbdata); + } +} + +static int +lua_task_add_timer(lua_State *L) +{ + struct ev_loop *ev_base; + struct rspamd_task *task; + + task = lua_check_task(L, 1); + ev_base = task->event_loop; + if (!lua_isfunction(L, 3)) { + return luaL_error(L, "invalid arguments: callback expected"); + } + + if (!lua_isnumber(L, 2)) { + return luaL_error(L, "invalid arguments: timeout expected"); + } + + struct rspamd_task_timer_cbdata *cbdata = rspamd_mempool_alloc(task->task_pool, sizeof(*cbdata)); + cbdata->L = L; + lua_pushvalue(L, 3); + cbdata->ev.data = cbdata; + cbdata->cbref = luaL_ref(L, LUA_REGISTRYINDEX); + cbdata->task = task; + cbdata->item = rspamd_symcache_get_cur_item(task); + + if (cbdata->item) { + cbdata->async_ev = rspamd_session_add_event_full(task->s, lua_timer_fin, cbdata, "timer", + rspamd_symcache_dyn_item_name(cbdata->task, cbdata->item)); + rspamd_symcache_item_async_inc(task, cbdata->item, "timer"); + } + else { + cbdata->async_ev = rspamd_session_add_event(task->s, lua_timer_fin, cbdata, "timer"); + } + + ev_timer_init(&cbdata->ev, lua_task_timer_cb, lua_tonumber(L, 2), 0.0); + ev_timer_start(ev_base, &cbdata->ev); + + return 0; +} + /* Init part */ static int |