module: add protocol processing callback

7 files changed, 128 insertions, 47 deletions

diff --git a/src/knot/include/module.h b/src/knot/include/module.h
index fd7e37437..9998e48d0 100644
--- a/src/knot/include/module.h
+++ b/src/knot/include/module.h
@@ -354,7 +354,7 @@ knotd_conf_t knotd_conf_check_item(knotd_conf_check_args_t *args,
                                    const yp_name_t *item_name);
 
 /*!
- * \brief Checks if address is in at least one of given ranges.
+ * Checks if address is in at least one of given ranges.
  *
  * \param[in] range
  * \param[in] addr
@@ -505,6 +505,12 @@ int knotd_qdata_zone_rrset(const knotd_qdata_t *qdata, const knot_dname_t *zone_
                            const knot_dname_t *node_name, uint16_t type,
                            knot_rrset_t *out);
 
+/*! Transport protocol processing states. */
+typedef enum {
+	KNOTD_PROTO_STATE_PASS  = 0, /*!< Process normally. */
+	KNOTD_PROTO_STATE_BLOCK = 1, /*!< Block the packet/connection. */
+} knotd_proto_state_t;
+
 /*! General query processing states. */
 typedef enum {
 	KNOTD_STATE_NOOP  = 0, /*!< No response. */
@@ -513,7 +519,7 @@ typedef enum {
 	KNOTD_STATE_FINAL = 6, /*!< Finished and finalized (QNAME, EDNS, TSIG). */
 } knotd_state_t;
 
-/*! brief Internet query processing states. */
+/*! Internet query processing states. */
 typedef enum {
 	KNOTD_IN_STATE_BEGIN,  /*!< Begin name resolution. */
 	KNOTD_IN_STATE_NODATA, /*!< Positive result with NO data. */
@@ -527,15 +533,29 @@ typedef enum {
 
 /*! Query module processing stages. */
 typedef enum {
-	KNOTD_STAGE_BEGIN = 0,  /*!< Before query processing. */
-	KNOTD_STAGE_PREANSWER,  /*!< Before section processing. */
-	KNOTD_STAGE_ANSWER,     /*!< Answer section processing. */
-	KNOTD_STAGE_AUTHORITY,  /*!< Authority section processing. */
-	KNOTD_STAGE_ADDITIONAL, /*!< Additional section processing. */
-	KNOTD_STAGE_END,        /*!< After query processing. */
+	KNOTD_STAGE_PROTO_BEGIN = 0,  /*!< Start of transport protocol processing. */
+	KNOTD_STAGE_BEGIN,            /*!< Before query processing. */
+	KNOTD_STAGE_PREANSWER,        /*!< Before section processing. */
+	KNOTD_STAGE_ANSWER,           /*!< Answer section processing. */
+	KNOTD_STAGE_AUTHORITY,        /*!< Authority section processing. */
+	KNOTD_STAGE_ADDITIONAL,       /*!< Additional section processing. */
+	KNOTD_STAGE_END,              /*!< After query processing. */
+	KNOTD_STAGE_PROTO_END,        /*!< End of transport protocol processing. */
 } knotd_stage_t;
 
 /*!
+ * Transport protocol processing hook.
+ *
+ * \param[in] state    Current processing state.
+ * \param[in] params   Processing parameters.
+ * \param[in] mod      Module context.
+ *
+ * \return Next processing state.
+ */
+typedef knotd_proto_state_t (*knotd_mod_proto_hook_f)
+	(knotd_proto_state_t state, knotd_qdata_params_t *params, knotd_mod_t *mod);
+
+/*!
  * General processing hook.
  *
  * \param[in] state    Current processing state.
@@ -562,6 +582,17 @@ typedef knotd_in_state_t (*knotd_mod_in_hook_f)
 	(knotd_in_state_t state, knot_pkt_t *pkt, knotd_qdata_t *qdata, knotd_mod_t *mod);
 
 /*!
+ * Registers transport protocol processing module hook.
+ *
+ * \param[in] mod    Module context.
+ * \param[in] stage  Processing stage (KNOTD_STAGE_PROTO_BEGIN or KNOTD_STAGE_PROTO_END).
+ * \param[in] hook   Module hook.
+ *
+ * \return Error code, KNOT_EOK if success.
+ */
+int knotd_mod_proto_hook(knotd_mod_t *mod, knotd_stage_t stage, knotd_mod_proto_hook_f hook);
+
+/*!
  * Registers general processing module hook.
  *
  * \param[in] mod    Module context.
diff --git a/src/knot/nameserver/internet.c b/src/knot/nameserver/internet.c
index 18c310249..034fd268c 100644
--- a/src/knot/nameserver/internet.c
+++ b/src/knot/nameserver/internet.c
@@ -1,4 +1,4 @@
-/*  Copyright (C) 2022 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
+/*  Copyright (C) 2024 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
 
     This program is free software: you can redistribute it and/or modify
     it under the terms of the GNU General Public License as published by
@@ -658,7 +658,8 @@ static knot_layer_state_t answer_query(knot_pkt_t *pkt, knotd_qdata_t *qdata)
 	/* Resolve PREANSWER. */
 	if (plan != NULL) {
 		WALK_LIST(step, plan->stage[KNOTD_STAGE_PREANSWER]) {
-			SOLVE_STEP(step->process, state, step->ctx);
+			assert(step->type == QUERY_HOOK_TYPE_IN);
+			SOLVE_STEP(step->in_hook, state, step->ctx);
 		}
 	}
 
@@ -670,7 +671,8 @@ static knot_layer_state_t answer_query(knot_pkt_t *pkt, knotd_qdata_t *qdata)
 	}
 	if (plan != NULL) {
 		WALK_LIST(step, plan->stage[KNOTD_STAGE_ANSWER]) {
-			SOLVE_STEP(step->process, state, step->ctx);
+			assert(step->type == QUERY_HOOK_TYPE_IN);
+			SOLVE_STEP(step->in_hook, state, step->ctx);
 		}
 	}
 
@@ -682,7 +684,8 @@ static knot_layer_state_t answer_query(knot_pkt_t *pkt, knotd_qdata_t *qdata)
 	}
 	if (plan != NULL) {
 		WALK_LIST(step, plan->stage[KNOTD_STAGE_AUTHORITY]) {
-			SOLVE_STEP(step->process, state, step->ctx);
+			assert(step->type == QUERY_HOOK_TYPE_IN);
+			SOLVE_STEP(step->in_hook, state, step->ctx);
 		}
 	}
 
@@ -694,7 +697,8 @@ static knot_layer_state_t answer_query(knot_pkt_t *pkt, knotd_qdata_t *qdata)
 	}
 	if (plan != NULL) {
 		WALK_LIST(step, plan->stage[KNOTD_STAGE_ADDITIONAL]) {
-			SOLVE_STEP(step->process, state, step->ctx);
+			assert(step->type == QUERY_HOOK_TYPE_IN);
+			SOLVE_STEP(step->in_hook, state, step->ctx);
 		}
 	}
 
diff --git a/src/knot/nameserver/process_query.c b/src/knot/nameserver/process_query.c
index 7cb89b7fe..beced8478 100644
--- a/src/knot/nameserver/process_query.c
+++ b/src/knot/nameserver/process_query.c
@@ -553,7 +553,8 @@ static knot_layer_state_t process_query_err(knot_layer_t *ctx, knot_pkt_t *pkt)
 #define PROCESS_BEGIN(plan, step, next_state, qdata) \
 	if (plan != NULL) { \
 		WALK_LIST(step, plan->stage[KNOTD_STAGE_BEGIN]) { \
-			next_state = step->process(next_state, pkt, qdata, step->ctx); \
+			assert(step->type == QUERY_HOOK_TYPE_GENERAL); \
+			next_state = step->general_hook(next_state, pkt, qdata, step->ctx); \
 			if (next_state == KNOT_STATE_FAIL) { \
 				goto finish; \
 			} \
@@ -563,7 +564,8 @@ static knot_layer_state_t process_query_err(knot_layer_t *ctx, knot_pkt_t *pkt)
 #define PROCESS_END(plan, step, next_state, qdata) \
 	if (plan != NULL) { \
 		WALK_LIST(step, plan->stage[KNOTD_STAGE_END]) { \
-			next_state = step->process(next_state, pkt, qdata, step->ctx); \
+			assert(step->type == QUERY_HOOK_TYPE_GENERAL); \
+			next_state = step->general_hook(next_state, pkt, qdata, step->ctx); \
 			if (next_state == KNOT_STATE_FAIL) { \
 				next_state = process_query_err(ctx, pkt); \
 			} \
@@ -1005,6 +1007,30 @@ int process_query_put_rr(knot_pkt_t *pkt, knotd_qdata_t *qdata,
 	return ret;
 }
 
+knotd_proto_state_t process_query_proto(knotd_qdata_params_t *params,
+                                        const knotd_stage_t stage)
+{
+	assert(params);
+	assert(stage == KNOTD_STAGE_PROTO_BEGIN || stage == KNOTD_STAGE_PROTO_END);
+
+	knotd_proto_state_t state = KNOTD_PROTO_STATE_PASS;
+
+	rcu_read_lock();
+
+	struct query_plan *plan = conf()->query_plan;
+	if (plan != NULL) {
+		struct query_step *step;
+		WALK_LIST(step, plan->stage[stage]) {
+			assert(step->type == QUERY_HOOK_TYPE_PROTO);
+			state = step->proto_hook(state, params, step->ctx);
+		}
+	}
+
+	rcu_read_unlock();
+
+	return state;
+}
+
 /*! \brief Module implementation. */
 const knot_layer_api_t *process_query_layer(void)
 {
diff --git a/src/knot/nameserver/process_query.h b/src/knot/nameserver/process_query.h
index 42e66bebd..c005b58ef 100644
--- a/src/knot/nameserver/process_query.h
+++ b/src/knot/nameserver/process_query.h
@@ -1,4 +1,4 @@
-/*  Copyright (C) 2023 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
+/*  Copyright (C) 2024 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
 
     This program is free software: you can redistribute it and/or modify
     it under the terms of the GNU General Public License as published by
@@ -106,3 +106,14 @@ int process_query_sign_response(knot_pkt_t *pkt, knotd_qdata_t *qdata);
 int process_query_put_rr(knot_pkt_t *pkt, knotd_qdata_t *qdata,
                          const knot_rrset_t *rr, const knot_rrset_t *rrsigs,
                          uint16_t compr_hint, uint32_t flags);
+
+/*!
+ * \brief Processes all global module protocol callbacks at given stage.
+ *
+ * \param params   Query processing parameters.
+ * \param stage    Processing stage (KNOTD_STAGE_PROTO_BEGIN or KNOTD_STAGE_PROTO_END).
+ *
+ * \return Resulting state.
+ */
+knotd_proto_state_t process_query_proto(knotd_qdata_params_t *params,
+                                        const knotd_stage_t stage);
diff --git a/src/knot/nameserver/query_module.c b/src/knot/nameserver/query_module.c
index 4d50980e7..0708b1704 100644
--- a/src/knot/nameserver/query_module.c
+++ b/src/knot/nameserver/query_module.c
@@ -68,30 +68,31 @@ void query_plan_free(struct query_plan *plan)
 	free(plan);
 }
 
-static struct query_step *make_step(query_step_process_f process, void *ctx)
+int query_plan_step(struct query_plan *plan, knotd_stage_t stage,
+                    query_hook_type_t type, void *hook, void *ctx)
 {
-	struct query_step *step = calloc(1, sizeof(struct query_step));
+	struct query_step *step = calloc(1, sizeof(*step));
 	if (step == NULL) {
-		return NULL;
+		return KNOT_ENOMEM;
 	}
 
-	step->process = process;
+	step->type = type;
+	step->general_hook = hook;
 	step->ctx = ctx;
 
-	return step;
+	add_tail(&plan->stage[stage], &step->node);
+
+	return KNOT_EOK;
 }
 
-int query_plan_step(struct query_plan *plan, knotd_stage_t stage,
-                    query_step_process_f process, void *ctx)
+_public_
+int knotd_mod_proto_hook(knotd_mod_t *mod, knotd_stage_t stage, knotd_mod_proto_hook_f hook)
 {
-	struct query_step *step = make_step(process, ctx);
-	if (step == NULL) {
-		return KNOT_ENOMEM;
+	if (stage != KNOTD_STAGE_PROTO_BEGIN && stage != KNOTD_STAGE_PROTO_END) {
+		return KNOT_EINVAL;
 	}
 
-	add_tail(&plan->stage[stage], &step->node);
-
-	return KNOT_EOK;
+	return query_plan_step(mod->plan, stage, QUERY_HOOK_TYPE_PROTO, hook,  mod);
 }
 
 _public_
@@ -101,17 +102,18 @@ int knotd_mod_hook(knotd_mod_t *mod, knotd_stage_t stage, knotd_mod_hook_f hook)
 		return KNOT_EINVAL;
 	}
 
-	return query_plan_step(mod->plan, stage, hook, mod);
+	return query_plan_step(mod->plan, stage, QUERY_HOOK_TYPE_GENERAL, hook, mod);
 }
 
 _public_
 int knotd_mod_in_hook(knotd_mod_t *mod, knotd_stage_t stage, knotd_mod_in_hook_f hook)
 {
-	if (stage == KNOTD_STAGE_BEGIN || stage == KNOTD_STAGE_END) {
+	if (stage != KNOTD_STAGE_PREANSWER && stage != KNOTD_STAGE_ANSWER &&
+	    stage != KNOTD_STAGE_AUTHORITY && stage != KNOTD_STAGE_ADDITIONAL) {
 		return KNOT_EINVAL;
 	}
 
-	return query_plan_step(mod->plan, stage, hook, mod);
+	return query_plan_step(mod->plan, stage, QUERY_HOOK_TYPE_IN, hook, mod);
 }
 
 knotd_mod_t *query_module_open(conf_t *conf, server_t *server, conf_mod_id_t *mod_id,
diff --git a/src/knot/nameserver/query_module.h b/src/knot/nameserver/query_module.h
index 65957c4bc..d403b256a 100644
--- a/src/knot/nameserver/query_module.h
+++ b/src/knot/nameserver/query_module.h
@@ -1,4 +1,4 @@
-/*  Copyright (C) 2023 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
+/*  Copyright (C) 2024 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
 
     This program is free software: you can redistribute it and/or modify
     it under the terms of the GNU General Public License as published by
@@ -25,16 +25,24 @@
 #include "contrib/atomic.h"
 #include "contrib/ucw/lists.h"
 
-#define KNOTD_STAGES (KNOTD_STAGE_END + 1)
+#define KNOTD_STAGES (KNOTD_STAGE_PROTO_END + 1)
 
-typedef unsigned (*query_step_process_f)
-	(unsigned state, knot_pkt_t *pkt, knotd_qdata_t *qdata, knotd_mod_t *mod);
+typedef enum {
+	QUERY_HOOK_TYPE_PROTO,
+	QUERY_HOOK_TYPE_GENERAL,
+	QUERY_HOOK_TYPE_IN,
+} query_hook_type_t;
 
-/*! \brief Single processing step in query processing. */
+/*! \brief Single processing step in query/module processing. */
 struct query_step {
 	node_t node;
+	query_hook_type_t type;
+	union {
+		knotd_mod_proto_hook_f proto_hook;
+		knotd_mod_hook_f general_hook;
+		knotd_mod_in_hook_f in_hook;
+	};
 	void *ctx;
-	query_step_process_f process;
 };
 
 /*! Query plan represents a sequence of steps needed for query processing
@@ -53,7 +61,7 @@ void query_plan_free(struct query_plan *plan);
 
 /*! \brief Plan another step for given stage. */
 int query_plan_step(struct query_plan *plan, knotd_stage_t stage,
-                    query_step_process_f process, void *ctx);
+                    query_hook_type_t type, void *hook, void *ctx);
 
 /*! \brief Open query module identified by name. */
 knotd_mod_t *query_module_open(conf_t *conf, server_t *server, conf_mod_id_t *mod_id,
diff --git a/tests/knot/test_query_module.c b/tests/knot/test_query_module.c
index 4ab14d287..d31fa7c1d 100644
--- a/tests/knot/test_query_module.c
+++ b/tests/knot/test_query_module.c
@@ -1,4 +1,4 @@
-/*  Copyright (C) 2019 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
+/*  Copyright (C) 2024 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
 
     This program is free software: you can redistribute it and/or modify
     it under the terms of the GNU General Public License as published by
@@ -20,11 +20,10 @@
 
 #include "libknot/libknot.h"
 #include "knot/nameserver/query_module.h"
-#include "libknot/packet/pkt.h"
 
 /* Universal processing stage. */
-unsigned state_visit(unsigned state, knot_pkt_t *pkt, knotd_qdata_t *qdata,
-                     knotd_mod_t *mod)
+knotd_state_t state_visit(knotd_state_t state, knot_pkt_t *pkt, knotd_qdata_t *qdata,
+                          knotd_mod_t *mod)
 {
 	/* Visit current state */
 	bool *state_map = (bool *)mod;
@@ -49,8 +48,8 @@ int main(int argc, char *argv[])
 
 	/* Register all stage visits. */
 	int ret = KNOT_EOK;
-	for (unsigned stage = KNOTD_STAGE_BEGIN; stage < KNOTD_STAGES; ++stage) {
-		ret = query_plan_step(plan, stage, state_visit, state_map);
+	for (unsigned stage = KNOTD_STAGE_PROTO_BEGIN; stage < KNOTD_STAGES; ++stage) {
+		ret = query_plan_step(plan, stage, QUERY_HOOK_TYPE_GENERAL, state_visit, state_map);
 		if (ret != KNOT_EOK) {
 			break;
 		}
@@ -59,10 +58,10 @@ int main(int argc, char *argv[])
 
 	/* Execute the plan. */
 	int state = 0, next_state = 0;
-	for (unsigned stage = KNOTD_STAGE_BEGIN; stage < KNOTD_STAGES; ++stage) {
+	for (unsigned stage = KNOTD_STAGE_PROTO_BEGIN; stage < KNOTD_STAGES; ++stage) {
 		struct query_step *step = NULL;
 		WALK_LIST(step, plan->stage[stage]) {
-			next_state = step->process(state, NULL, NULL, step->ctx);
+			next_state = step->general_hook(state, NULL, NULL, step->ctx);
 			if (next_state != state + 1) {
 				break;
 			}