diff options
Diffstat (limited to 'src/lib-dict/dict-iter-lua.c')
| -rw-r--r-- | src/lib-dict/dict-iter-lua.c | 193 |
1 files changed, 193 insertions, 0 deletions
diff --git a/src/lib-dict/dict-iter-lua.c b/src/lib-dict/dict-iter-lua.c new file mode 100644 index 0000000..780de03 --- /dev/null +++ b/src/lib-dict/dict-iter-lua.c @@ -0,0 +1,193 @@ +/* Copyright (c) 2021 Dovecot authors, see the included COPYING file */ + +#include "lib.h" +#include "array.h" +#include "dict.h" +#include "dlua-script-private.h" +#include "dict-lua-private.h" +#include "dlua-wrapper.h" + +struct lua_dict_iter { + pool_t pool; + struct dict_iterate_context *iter; + ARRAY(int) refs; + int error_ref; + + lua_State *L; + bool yielded:1; +}; + +static void lua_dict_iter_unref(struct lua_dict_iter *iter) +{ + const char *error; + + /* deinit iteration if it hasn't been done yet */ + if (dict_iterate_deinit(&iter->iter, &error) < 0) { + e_error(dlua_script_from_state(iter->L)->event, + "Dict iteration failed: %s", error); + } + + pool_unref(&iter->pool); +} + +DLUA_WRAP_C_DATA(dict_iter, struct lua_dict_iter, lua_dict_iter_unref, NULL); + +static int lua_dict_iterate_step(lua_State *L); + +/* resume after a yield */ +static int lua_dict_iterate_step_continue(lua_State *L, + int status ATTR_UNUSED, + lua_KContext ctx ATTR_UNUSED) +{ + return lua_dict_iterate_step(L); +} + +static void lua_dict_iterate_more(struct lua_dict_iter *iter); + +/* + * Iteration step function + * + * Takes two args (a userdata state, and previous value) and returns the + * next value. + */ +static int lua_dict_iterate_step(lua_State *L) +{ + struct lua_dict_iter *iter; + const int *refs; + unsigned nrefs; + + DLUA_REQUIRE_ARGS(L, 2); + + iter = xlua_dict_iter_getptr(L, 1, NULL); + iter->yielded = FALSE; + + lua_dict_iterate_more(iter); + + if (iter->iter != NULL) { + /* iteration didn't end yet - yield */ + return lua_dict_iterate_step_continue(L, + lua_yieldk(L, 0, 0, lua_dict_iterate_step_continue), 0); + } + + /* dict iteration ended - return first key-value pair */ + refs = array_get(&iter->refs, &nrefs); + i_assert(nrefs % 2 == 0); + + if (nrefs == 0) { + if (iter->error_ref != 0) { + /* dict iteration generated an error - raise it now */ + lua_rawgeti(L, LUA_REGISTRYINDEX, iter->error_ref); + luaL_unref(L, LUA_REGISTRYINDEX, iter->error_ref); + return lua_error(L); + } + + return 0; /* return nil */ + } + + /* get the key & value from the registry */ + lua_rawgeti(L, LUA_REGISTRYINDEX, refs[0]); + lua_rawgeti(L, LUA_REGISTRYINDEX, refs[1]); + luaL_unref(L, LUA_REGISTRYINDEX, refs[0]); + luaL_unref(L, LUA_REGISTRYINDEX, refs[1]); + + array_delete(&iter->refs, 0, 2); + + return 2; +} + +static void lua_dict_iterate_more(struct lua_dict_iter *iter) +{ + const char *key, *const *values; + lua_State *L = iter->L; + const char *error; + + if (iter->iter == NULL) + return; /* done iterating the dict */ + + while (dict_iterate_values(iter->iter, &key, &values)) { + int ref; + + /* stash key */ + lua_pushstring(L, key); + ref = luaL_ref(L, LUA_REGISTRYINDEX); + array_push_back(&iter->refs, &ref); + + /* stash values */ + lua_newtable(L); + for (unsigned int i = 0; values[i] != NULL; i++) { + lua_pushstring(L, values[i]); + lua_seti(L, -2, i + 1); + } + ref = luaL_ref(L, LUA_REGISTRYINDEX); + array_push_back(&iter->refs, &ref); + } + + if (dict_iterate_has_more(iter->iter)) + return; + + if (dict_iterate_deinit(&iter->iter, &error) < 0) { + lua_pushstring(L, error); + iter->error_ref = luaL_ref(L, LUA_REGISTRYINDEX); + } +} + +/* dict iter callback */ +static void lua_dict_iterate_callback(struct lua_dict_iter *iter) +{ + if (iter->yielded) + return; + iter->yielded = TRUE; + dlua_pcall_yieldable_resume(iter->L, 1); +} + +/* + * Iterate a dict at key [-(3|4),+2,e] + * + * Args: + * 1) userdata: sturct dict *dict + * 2) string: key + * 3) integer: flags + * 4*) string: username + * + * Returns: + * Returns a iteration step function and dict iter userdata. + * Username will be NULL if not provided in args. + */ +int lua_dict_iterate(lua_State *L) +{ + enum dict_iterate_flags flags; + struct lua_dict_iter *iter; + struct dict *dict; + const char *path, *username = NULL; + pool_t pool; + + DLUA_REQUIRE_ARGS_IN(L, 3, 4); + + dict = dlua_check_dict(L, 1); + path = luaL_checkstring(L, 2); + flags = luaL_checkinteger(L, 3); + if (lua_gettop(L) >= 4) + username = luaL_checkstring(L, 4); + lua_dict_check_key_prefix(L, path, username); + + struct dict_op_settings set = { + .username = username, + }; + + /* set up iteration */ + pool = pool_alloconly_create("lua dict iter", 128); + iter = p_new(pool, struct lua_dict_iter, 1); + iter->pool = pool; + iter->iter = dict_iterate_init(dict, &set, path, flags | + DICT_ITERATE_FLAG_ASYNC); + p_array_init(&iter->refs, iter->pool, 32); + iter->L = L; + + dict_iterate_set_async_callback(iter->iter, + lua_dict_iterate_callback, iter); + + /* push return values: func, state */ + lua_pushcfunction(L, lua_dict_iterate_step); + xlua_pushdict_iter(L, iter, FALSE); + return 2; +} |