From: Marek Vavruša <mvavrusa@cloudflare.com>
Date: Mon, 27 Nov 2017 22:03:44 +0000 (-0800)
Subject: daemon: added callback for resolve() query initialisation
X-Git-Tag: v2.0.0~16^2~25
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=bb9c2e89bdf8654a2d9dfa08fa74b386d9d1b862;p=thirdparty%2Fknot-resolver.git

daemon: added callback for resolve() query initialisation

this allows for changing request settings or performing other actions
just after the request is created.
---

diff --git a/daemon/README.rst b/daemon/README.rst
index 1d891e6c4..b2f8c4637 100644
--- a/daemon/README.rst
+++ b/daemon/README.rst
@@ -450,15 +450,28 @@ Environment
       > user('root')
       Operation not permitted
 
-.. function:: resolve(qname, qtype[, qclass = kres.class.IN, options = 0, callback = nil])
+.. function:: resolve(name, type[, class = kres.class.IN, options = 0, finish = nil, begin = nil])
 
-   :param string qname: Query name (e.g. 'com.')
-   :param number qtype: Query type (e.g. ``kres.type.NS``)
-   :param number qclass: Query class *(optional)* (e.g. ``kres.class.IN``)
+   :param string name: Query name (e.g. 'com.')
+   :param number type: Query type (e.g. ``kres.type.NS``)
+   :param number class: Query class *(optional)* (e.g. ``kres.class.IN``)
    :param number options: Resolution options (see query flags)
-   :param function callback: Callback to be executed when resolution completes (e.g. `function cb (pkt, req) end`). The callback gets a packet containing the final answer and doesn't have to return anything.
+   :param function finish: Callback to be executed when resolution completes (e.g. `function cb (pkt, req) end`). The callback gets a packet containing the final answer and doesn't have to return anything.
+   :param function begin: Callback to be executed when the request is created.
    :return: boolean
 
+   The function can also be executed with a table of arguments instead. This is useful if you'd like to skip some arguments, for example:
+
+   .. code-block:: lua
+
+      resolve {
+         name = 'example.com',
+         type = kres.type.AAAA,
+         finish = function ()
+             print('done')
+         end,
+      }
+
    Example:
 
    .. code-block:: lua
diff --git a/daemon/bindings.c b/daemon/bindings.c
index 140ec7df3..a98b277f7 100644
--- a/daemon/bindings.c
+++ b/daemon/bindings.c
@@ -1244,22 +1244,40 @@ static int wrk_resolve(lua_State *L)
 	if (options->DNSSEC_WANT) {
 		knot_edns_set_do(pkt->opt_rr);
 	}
+
 	if (options->DNSSEC_CD) {
 		knot_wire_set_cd(pkt->wire);
 	}
 
+	/* Create task and start with a first question */
+	struct qr_task *task = worker_resolve_start(worker, pkt, *options);
+	if (!task) {
+		knot_rrset_free(&pkt->opt_rr, NULL);
+		knot_pkt_free(&pkt);
+		lua_pushstring(L, "couldn't create a resolution request");
+		lua_error(L);
+	}
+
+	/* Store completion callback in registry */
 	if (lua_isfunction(L, 5)) {
-		/* Store callback in registry */
 		lua_pushvalue(L, 5);
 		int cb = luaL_ref(L, LUA_REGISTRYINDEX);
-		ret = worker_resolve(worker, pkt, *options, resolve_callback, (void *) (intptr_t)cb);
-	} else {
-		ret = worker_resolve(worker, pkt, *options, NULL, NULL);
+		task->on_complete = resolve_callback;
+		task->baton = (void *) (intptr_t)cb;
 	}
-	
+
+	/* Add initialisation callback */
+	if (lua_isfunction(L, 6)) {
+		lua_pushvalue(L, 6);
+		lua_pushlightuserdata(L, &task->req);
+		(void) execute_callback(L, 1);
+	}
+
+	/* Start execution */
+	int ret = worker_resolve_exec(task, pkt);
+	lua_pushboolean(L, ret == 0);
 	knot_rrset_free(&pkt->opt_rr, NULL);
 	knot_pkt_free(&pkt);
-	lua_pushboolean(L, ret == 0);
 	return 1;
 }
 
diff --git a/daemon/lua/sandbox.lua b/daemon/lua/sandbox.lua
index 402db091f..aa4b318e1 100644
--- a/daemon/lua/sandbox.lua
+++ b/daemon/lua/sandbox.lua
@@ -23,16 +23,27 @@ end
 
 -- Resolver bindings
 kres = require('kres')
-trust_anchors = require('trust_anchors')
 if rawget(kres, 'str2dname') ~= nil then
 	todname = kres.str2dname
 end
 
--- Compat. wrapper for query flags.
-worker.resolve = function (p1, p2, p3, options, p5)
+-- Compatibility wrapper for query flags.
+worker.resolve = function (qname, qtype, qclass, options, finish, begin)
+	-- Alternatively use named arguments
+	if type(qname) == 'table' then
+		local t = qname
+		qname = t.name
+		qtype = t.type or kres.type.A
+		qclass = t.class or kres.class.IN
+		options = t.options
+		finish = t.finish
+		begin = t.begin
+	end
+	-- Translate options and resolve
 	options = kres.mk_qflags(options)
-	return worker.resolve_unwrapped (p1, p2, p3, options, p5)
+	return worker.resolve_unwrapped(qname, qtype, qclass, options, finish, begin)
 end
+
 resolve = worker.resolve
 
 -- Shorthand for aggregated per-worker information
@@ -172,7 +183,7 @@ end
 
 -- Make sandboxed environment
 local function make_sandbox(defined)
-	local __protected = { modules = true, cache = true, net = true, trust_anchors = true }
+	local __protected = { worker = true, env = true, modules = true, cache = true, net = true, trust_anchors = true }
 
 	-- Compute and export the list of top-level names (hidden otherwise)
 	local nl = ""
@@ -203,6 +214,7 @@ else -- Lua 5.2+
 end
 
 -- Load embedded modules
+trust_anchors = require('trust_anchors')
 modules.load('ta_signal_query')
 modules.load('priming')
 modules.load('detect_time_skew')
diff --git a/daemon/worker.c b/daemon/worker.c
index 3b8f0959a..fb8115c8b 100644
--- a/daemon/worker.c
+++ b/daemon/worker.c
@@ -57,7 +57,7 @@ struct req
 #define qr_task_ref(task) \
 	do { ++(task)->refs; } while(0)
 #define qr_task_unref(task) \
-	do { if (--(task)->refs == 0) { qr_task_free(task); } } while (0)
+	do { if (task && --(task)->refs == 0) { qr_task_free(task); } } while (0)
 #define qr_valid_handle(task, checked) \
 	(!uv_is_closing((checked)) || (task)->source.handle == (checked))
 
@@ -280,6 +280,7 @@ static struct qr_task *qr_task_create(struct worker_ctx *worker, uv_handle_t *ha
 	task->source.handle = handle;
 	task->timeout = NULL;
 	task->on_complete = NULL;
+	task->baton = NULL;
 	task->req.qsource.key = NULL;
 	task->req.qsource.addr = NULL;
 	task->req.qsource.dst_addr = NULL;
@@ -1037,33 +1038,55 @@ int worker_process_tcp(struct worker_ctx *worker, uv_stream_t *handle, const uin
 	return submitted;
 }
 
-int worker_resolve(struct worker_ctx *worker, knot_pkt_t *query, struct kr_qflags options,
-		   worker_cb_t on_complete, void *baton)
+struct qr_task *worker_resolve_start(struct worker_ctx *worker, knot_pkt_t *query, struct kr_qflags options)
 {
 	if (!worker || !query) {
-		return kr_error(EINVAL);
+		return NULL;
 	}
 
-	/* Create task */
 	struct qr_task *task = qr_task_create(worker, NULL, NULL);
 	if (!task) {
-		return kr_error(ENOMEM);
+		return NULL;
 	}
-	task->baton = baton;
-	task->on_complete = on_complete;
-	/* Start task */
+
 	int ret = qr_task_start(task, query);
+	if (ret != 0) {
+		qr_task_unref(task);
+		return NULL;
+	}
 
 	/* Set options late, as qr_task_start() -> kr_resolve_begin() rewrite it. */
 	kr_qflags_set(&task->req.options, options);
+	return task;
+}
 
-	if (ret != 0) {
-		qr_task_unref(task);
-		return ret;
+int worker_resolve_exec(struct qr_task *task, knot_pkt_t *query)
+{
+	if (!task) {
+		return kr_error(EINVAL);
 	}
 	return qr_task_step(task, NULL, query);
 }
 
+int worker_resolve(struct worker_ctx *worker, knot_pkt_t *query, struct kr_qflags options,
+		   worker_cb_t on_complete, void *baton)
+{
+	if (!worker || !query) {
+		return kr_error(EINVAL);
+	}
+
+	/* Create task */
+	struct qr_task *task = worker_resolve_start(worker, query, options);
+	if (!task) {
+		return kr_error(ENOMEM);
+	}
+
+	/* Install completion handler */
+	task->baton = baton;
+	task->on_complete = on_complete;
+	return worker_resolve_exec(task, query);
+}
+
 /** Reserve worker buffers */
 static int worker_reserve(struct worker_ctx *worker, size_t ring_maxlen)
 {
diff --git a/daemon/worker.h b/daemon/worker.h
index 55357e88f..0d774b489 100644
--- a/daemon/worker.h
+++ b/daemon/worker.h
@@ -23,6 +23,7 @@
 
 /** Worker state (opaque). */
 struct worker_ctx;
+struct qr_task;
 /** Worker callback */
 typedef void (*worker_cb_t)(struct worker_ctx *worker, struct kr_request *req, void *baton);
 
@@ -52,6 +53,21 @@ int worker_process_tcp(struct worker_ctx *worker, uv_stream_t *handle,
  */
 int worker_end_tcp(struct worker_ctx *worker, uv_handle_t *handle);
 
+/**
+ * Start query resolution with given query.
+ *
+ * @return task or NULL
+ */
+struct qr_task *worker_resolve_start(struct worker_ctx *worker, knot_pkt_t *query, struct kr_qflags options);
+
+/**
+ * Execute a request with given query.
+ * It expects task to be created with \fn worker_resolve_start.
+ *
+ * @return 0 or an error code
+ */
+int worker_resolve_exec(struct qr_task *task, knot_pkt_t *query);
+
 /**
  * Schedule query for resolution.
  *
diff --git a/modules/predict/predict.lua b/modules/predict/predict.lua
index 5160b03fe..7dd38754c 100644
--- a/modules/predict/predict.lua
+++ b/modules/predict/predict.lua
@@ -34,7 +34,7 @@ function predict.drain()
 	local deleted = 0
 	for key, _ in pairs(predict.queue) do
 		local qtype, qname = key:match('(%S*)%s(.*)')
-		worker.resolve(qname, kres.type[qtype], kres.class.IN, 'NO_CACHE')
+		resolve(qname, kres.type[qtype], kres.class.IN, {'NO_CACHE'})
 		predict.queue[key] = nil
 		deleted = deleted + 1
 		-- Resolve smaller batches at a time
@@ -177,7 +177,7 @@ predict.layer = {
 		for i = 0, (tonumber(qrys.len) - 1) do -- size_t doesn't work for some reason
 			local qry = qrys.at[i]
 			if qry.flags.EXPIRING == true then
-				worker.resolve(kres.dname2str(qry.sname), qry.stype, qry.sclass, 'NO_CACHE')
+				resolve(kres.dname2str(qry.sname), qry.stype, qry.sclass, {'NO_CACHE'})
 			end
 		end
 	end
diff --git a/modules/predict/predict_test.lua b/modules/predict/predict_test.lua
index 1f6d241ad..fe8a8192c 100644
--- a/modules/predict/predict_test.lua
+++ b/modules/predict/predict_test.lua
@@ -5,7 +5,7 @@ modules = { 'predict' }
 local resolve_count = 0
 local current_epoch = 0
 
-worker.resolve = function ()
+resolve = function ()
 	resolve_count = resolve_count + 1
 end