1 files changed, 304 insertions, 0 deletions
diff --git a/daemon/ffimodule.c b/daemon/ffimodule.c
new file mode 100644
index 0000000..4206b4c
--- /dev/null
+++ b/daemon/ffimodule.c
@@ -0,0 +1,304 @@
+/*  Copyright (C) CZ.NIC, z.s.p.o. <knot-resolver@labs.nic.cz>
+ *  SPDX-License-Identifier: GPL-3.0-or-later
+ */
+
+#include <uv.h>
+#include <lua.h>
+#include <lauxlib.h>
+
+#include "daemon/bindings/impl.h"
+#include "daemon/engine.h"
+#include "daemon/ffimodule.h"
+#include "daemon/worker.h"
+#include "lib/module.h"
+#include "lib/layer.h"
+
+/** @internal Slots for layer callbacks.
+  * Each slot ID corresponds to Lua reference in module API. */
+enum {
+	SLOT_begin = 0,
+	SLOT_reset,
+	SLOT_finish,
+	SLOT_consume,
+	SLOT_produce,
+	SLOT_checkout,
+	SLOT_answer_finalize,
+	SLOT_count /* dummy, must be the last */
+};
+
+/** Lua registry indices for functions that wrap layer callbacks (shared by all lua modules). */
+static int l_ffi_wrap_slots[SLOT_count] = { 0 };
+
+/** @internal Continue with coroutine. */
+static void l_ffi_resume_cb(uv_idle_t *check)
+{
+	lua_State *L = check->data;
+	int status = lua_resume(L, 0);
+	if (status != LUA_YIELD) {
+		uv_idle_stop(check); /* Stop coroutine */
+		uv_close((uv_handle_t *)check, (uv_close_cb)free);
+	}
+	lua_pop(L, lua_gettop(L));
+}
+
+/** @internal Schedule deferred continuation. */
+static int l_ffi_defer(lua_State *L)
+{
+	uv_idle_t *check = malloc(sizeof(*check));
+	if (!check) {
+		return kr_error(ENOMEM);
+	}
+	uv_idle_init(uv_default_loop(), check);
+	check->data = L;
+	return uv_idle_start(check, l_ffi_resume_cb);
+}
+
+/** Common part of calling modname.(de)init in lua.
+ * The function to call should be on top of the stack and it gets popped. */
+static int l_ffi_modcb(lua_State *L, struct kr_module *module)
+{
+	if (lua_isnil(L, -1)) {
+		lua_pop(L, 1); /* .(de)init == nil, maybe even the module table doesn't exist */
+		return kr_ok();
+	}
+	lua_getglobal(L, "modules_ffi_wrap_modcb");
+	lua_insert(L, -2); /* swap with .(de)init */
+	lua_pushpointer(L, module);
+	if (lua_pcall(L, 2, 0, 0) == 0)
+		return kr_ok();
+	kr_log_error(SYSTEM, "error: %s\n", lua_tostring(L, -1));
+	lua_pop(L, 1);
+	return kr_error(1);
+}
+
+static int l_ffi_deinit(struct kr_module *module)
+{
+	/* Call .deinit(), if it exists. */
+	lua_State *L = the_worker->engine->L;
+	lua_getglobal(L, module->name);
+	lua_getfield(L, -1, "deinit");
+	const int ret = l_ffi_modcb(L, module);
+	lua_pop(L, 1); /* the module's table */
+
+	const kr_layer_api_t *api = module->layer;
+	if (!api) {
+		return ret;
+	}
+	/* Unregister layer callback references from registry. */
+	for (int si = 0; si < SLOT_count; ++si) {
+		if (api->cb_slots[si] > 0) {
+			luaL_unref(L, LUA_REGISTRYINDEX, api->cb_slots[si]);
+		}
+	}
+	free_const(api);
+	return ret;
+}
+
+kr_layer_t kr_layer_t_static;
+
+/** @internal Helper for calling a layer Lua function by e.g. SLOT_begin. */
+static int l_ffi_call_layer(kr_layer_t *ctx, int slot_ix)
+{
+	const int wrap_slot = l_ffi_wrap_slots[slot_ix];
+	const int cb_slot = ctx->api->cb_slots[slot_ix];
+	kr_require(wrap_slot > 0 && cb_slot > 0);
+	lua_State *L = the_worker->engine->L;
+	lua_rawgeti(L, LUA_REGISTRYINDEX, wrap_slot);
+	lua_rawgeti(L, LUA_REGISTRYINDEX, cb_slot);
+	/* We pass the content of *ctx via a global structure to avoid
+	 * lua (full) userdata, as that's relatively expensive (GC-allocated).
+	 * Performance: copying isn't ideal, but it's not visible in profiles. */
+	memcpy(&kr_layer_t_static, ctx, sizeof(*ctx));
+
+	int ret = lua_pcall(L, 1, 1, 0);
+	/* Handle result of the pcall.
+	 * Default state: ctx->req->state seems safer than ctx->state,
+	 * in case the pcall touched req->state. */
+	int state = ctx->req->state;
+	if (ret) { /* Exception or another lua problem. */
+		state = KR_STATE_FAIL;
+		kr_log_error(SYSTEM, "error: %s\n", lua_tostring(L, -1));
+
+	} else if (lua_isnumber(L, -1)) { /* Explicitly returned state. */
+		state = lua_tointeger(L, -1);
+		if (!kr_state_consistent(state)) {
+			kr_log_error(SYSTEM, "error: nonsense state returned from lua module layer: %d\n",
+					state);
+			state = KR_STATE_FAIL;
+		}
+
+	} else if (lua_isnil(L, -1)) { /* Don't change state. */
+
+	} else if (kr_fails_assert(!lua_isthread(L, -1))) { /* Continuations */
+		/* TODO: unused, possibly in a bad shape.  Meant KR_STATE_YIELD? */
+		if (l_ffi_defer(lua_tothread(L, -1)) != 0)
+			state = KR_STATE_FAIL;
+
+	} else { /* Nonsense returned. */
+		state = KR_STATE_FAIL;
+		kr_log_error(SYSTEM, "error: nonsense returned from lua module layer: %s\n",
+				lua_tostring(L, -1));
+		/* Unfortunately we can't easily get name of the module/function here. */
+	}
+	lua_pop(L, 1);
+	return state;
+}
+
+static int l_ffi_layer_begin(kr_layer_t *ctx)
+{
+	return l_ffi_call_layer(ctx, SLOT_begin);
+}
+
+static int l_ffi_layer_reset(kr_layer_t *ctx)
+{
+	return l_ffi_call_layer(ctx, SLOT_reset);
+}
+
+static int l_ffi_layer_finish(kr_layer_t *ctx)
+{
+	ctx->pkt = ctx->req->answer;
+	return l_ffi_call_layer(ctx, SLOT_finish);
+}
+
+static int l_ffi_layer_consume(kr_layer_t *ctx, knot_pkt_t *pkt)
+{
+	if (ctx->state & KR_STATE_FAIL) {
+		return ctx->state; /* Already failed, skip */
+	}
+	ctx->pkt = pkt;
+	return l_ffi_call_layer(ctx, SLOT_consume);
+}
+
+static int l_ffi_layer_produce(kr_layer_t *ctx, knot_pkt_t *pkt)
+{
+	if (ctx->state & KR_STATE_FAIL) {
+		return ctx->state; /* Already failed, skip */
+	}
+	ctx->pkt = pkt;
+	return l_ffi_call_layer(ctx, SLOT_produce);
+}
+
+static int l_ffi_layer_checkout(kr_layer_t *ctx, knot_pkt_t *pkt,
+				struct sockaddr *dst, int type)
+{
+	if (ctx->state & KR_STATE_FAIL) {
+		return ctx->state; /* Already failed, skip */
+	}
+	ctx->pkt = pkt;
+	ctx->dst = dst;
+	ctx->is_stream = (type == SOCK_STREAM);
+	return l_ffi_call_layer(ctx, SLOT_checkout);
+}
+
+static int l_ffi_layer_answer_finalize(kr_layer_t *ctx)
+{
+	return l_ffi_call_layer(ctx, SLOT_answer_finalize);
+}
+
+int ffimodule_init(lua_State *L)
+{
+	/* Wrappers defined in ./lua/sandbox.lua */
+	/* for API: (int state, kr_request_t *req) */
+	lua_getglobal(L, "modules_ffi_layer_wrap1");
+	const int wrap1 = luaL_ref(L, LUA_REGISTRYINDEX);
+	/* for API: (int state, kr_request_t *req, knot_pkt_t *) */
+	lua_getglobal(L, "modules_ffi_layer_wrap2");
+	const int wrap2 = luaL_ref(L, LUA_REGISTRYINDEX);
+	lua_getglobal(L, "modules_ffi_layer_wrap_checkout");
+	const int wrap_checkout = luaL_ref(L, LUA_REGISTRYINDEX);
+	if (wrap1 == LUA_REFNIL || wrap2 == LUA_REFNIL || wrap_checkout == LUA_REFNIL) {
+		return kr_error(ENOENT);
+	}
+
+	const int slots[SLOT_count] = {
+		[SLOT_begin]   = wrap1,
+		[SLOT_reset]   = wrap1,
+		[SLOT_finish]  = wrap2,
+		[SLOT_consume] = wrap2,
+		[SLOT_produce] = wrap2,
+		[SLOT_checkout] = wrap_checkout,
+		[SLOT_answer_finalize] = wrap1,
+	};
+	memcpy(l_ffi_wrap_slots, slots, sizeof(l_ffi_wrap_slots));
+	return kr_ok();
+}
+void ffimodule_deinit(lua_State *L)
+{
+	/* Unref each wrapper function from lua.
+	 * It's probably useless, as we're about to destroy lua_State, but... */
+	const int wrapsIndices[] = {
+		SLOT_begin,
+		SLOT_consume,
+		SLOT_checkout,
+	};
+	for (int i = 0; i < sizeof(wrapsIndices) / sizeof(wrapsIndices[0]); ++i) {
+		luaL_unref(L, LUA_REGISTRYINDEX, l_ffi_wrap_slots[wrapsIndices[i]]);
+	}
+}
+
+/** @internal Conditionally register layer trampoline
+  * @warning Expects 'module.layer' to be on top of Lua stack. */
+#define LAYER_REGISTER(L, api, name) do { \
+	int *cb_slot = (api)->cb_slots + SLOT_ ## name; \
+	lua_getfield((L), -1, #name); \
+	if (!lua_isnil((L), -1)) { \
+		(api)->name = l_ffi_layer_ ## name; \
+		*cb_slot = luaL_ref((L), LUA_REGISTRYINDEX); \
+	} else { \
+		lua_pop((L), 1); \
+	} \
+} while(0)
+
+/** @internal Create C layer api wrapper. */
+static kr_layer_api_t *l_ffi_layer_create(lua_State *L, struct kr_module *module)
+{
+	/* Fabricate layer API wrapping the Lua functions
+	 * reserve slots after it for references to Lua callbacks. */
+	const size_t api_length = offsetof(kr_layer_api_t, cb_slots)
+				+ (SLOT_count * sizeof(module->layer->cb_slots[0]));
+	kr_layer_api_t *api = calloc(1, api_length);
+	if (api) {
+		LAYER_REGISTER(L, api, begin);
+		LAYER_REGISTER(L, api, finish);
+		LAYER_REGISTER(L, api, consume);
+		LAYER_REGISTER(L, api, produce);
+		LAYER_REGISTER(L, api, checkout);
+		LAYER_REGISTER(L, api, answer_finalize);
+		LAYER_REGISTER(L, api, reset);
+	}
+	return api;
+}
+
+#undef LAYER_REGISTER
+
+int ffimodule_register_lua(struct engine *engine, struct kr_module *module, const char *name)
+{
+	/* Register module in Lua */
+	lua_State *L = engine->L;
+	lua_getglobal(L, "require");
+	lua_pushfstring(L, "kres_modules.%s", name);
+	if (lua_pcall(L, 1, LUA_MULTRET, 0) != 0) {
+		kr_log_error(SYSTEM, "error: %s\n", lua_tostring(L, -1));
+		lua_pop(L, 1);
+		return kr_error(ENOENT);
+	}
+	lua_setglobal(L, name);
+	lua_getglobal(L, name);
+
+	/* Create FFI module with trampolined functions. */
+	memset(module, 0, sizeof(*module));
+	module->name = strdup(name);
+	module->deinit = &l_ffi_deinit;
+	/* Bake layer API if defined in module */
+	lua_getfield(L, -1, "layer");
+	if (!lua_isnil(L, -1)) {
+		module->layer = l_ffi_layer_create(L, module);
+	}
+	lua_pop(L, 1); /* .layer table */
+
+	/* Now call .init(), if it exists. */
+	lua_getfield(L, -1, "init");
+	const int ret = l_ffi_modcb(L, module);
+	lua_pop(L, 1); /* the module's table */
+	return ret;
+}