noudp: add configuration option for filtering rate

src/knot/modules/noudp/noudp.c

-/*  Copyright (C) 2017 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
+/*  Copyright (C) 2019 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
@@ -16,6 +16,18 @@
 
 #include "knot/include/module.h"
 
+#define MOD_UDP_ALLOW_RATE	"\x0e""udp-allow-rate"
+
+const yp_item_t noudp_conf[] = {
+	{ MOD_UDP_ALLOW_RATE, YP_TINT, YP_VINT = { 0, UINT32_MAX, 0 } },
+	{ NULL }
+};
+
+typedef struct {
+	uint32_t rate;
+	uint32_t *counters;
+} noudp_ctx_t;
+
 static bool is_udp(knotd_qdata_t *qdata)
 {
 	return qdata->params->flags & KNOTD_QUERY_FLAG_LIMIT_SIZE;
@@ -25,6 +37,11 @@ static knotd_state_t noudp_begin(knotd_state_t state, knot_pkt_t *pkt,
                                  knotd_qdata_t *qdata, knotd_mod_t *mod)
 {
 	if (is_udp(qdata)) {
+		noudp_ctx_t *ctx = knotd_mod_ctx(mod);
+		if (ctx->rate > 0 && ++ctx->counters[qdata->params->thread_id] >= ctx->rate) {
+			ctx->counters[qdata->params->thread_id] = 0;
+			return state;
+		}
 		knot_wire_set_tc(pkt->wire);
 		return KNOTD_STATE_DONE;
 	}
@@ -34,10 +51,34 @@ static knotd_state_t noudp_begin(knotd_state_t state, knot_pkt_t *pkt,
 
 int noudp_load(knotd_mod_t *mod)
 {
-	knotd_mod_hook(mod, KNOTD_STAGE_BEGIN, noudp_begin);
+	noudp_ctx_t *ctx = calloc(1, sizeof(noudp_ctx_t));
+	if (ctx == NULL) {
+		return KNOT_ENOMEM;
+	}
+
+	knotd_conf_t conf = knotd_conf_mod(mod, MOD_UDP_ALLOW_RATE);
+	ctx->rate = conf.single.integer;
+	if (ctx->rate > 0) {
+		knotd_conf_t udp = knotd_conf_env(mod, KNOTD_CONF_ENV_WORKERS_UDP);
+		size_t udp_workers = udp.single.integer;
+		ctx->counters = calloc(udp_workers, sizeof(uint32_t));
+		if (ctx->counters == NULL) {
+			free(ctx);
+			return KNOT_ENOMEM;
+		}
+	}
+
+	knotd_mod_ctx_set(mod, ctx);
 
-	return KNOT_EOK;
+	return knotd_mod_hook(mod, KNOTD_STAGE_BEGIN, noudp_begin);
+}
+
+void noudp_unload(knotd_mod_t *mod)
+{
+	noudp_ctx_t *ctx = knotd_mod_ctx(mod);
+	free(ctx->counters);
+	free(ctx);
 }
 
 KNOTD_MOD_API(noudp, KNOTD_MOD_FLAG_SCOPE_ANY | KNOTD_MOD_FLAG_OPT_CONF,
-              noudp_load, NULL, NULL, NULL);
+              noudp_load, noudp_unload, noudp_conf, NULL);

src/knot/modules/noudp/noudp.rst

@@ -3,18 +3,47 @@
 ``noudp`` — No UDP response
 ===========================
 
-The module sends empty truncated response to any UDP query. TCP queries are
+The module sends empty truncated response to a UDP query. TCP queries are
 not affected.
 
 Example
 -------
 
-To enable this module globally, you need to add something like the following
-to the configuration file::
+To enable this module for all configured zones and every UDP query::
 
     template:
       - id: default
         global-module: mod-noudp
 
+Or with specified UDP allow rate::
+
+    mod-noudp:
+      - id: sometimes
+        udp-allow-rate: 1000  # Don't truncate every 1000th UDP query
+
+    template:
+      - id: default
+        module: mod-noudp/sometimes
+
+Module reference
+----------------
+
+::
+
+  mod-noudp:
+   - id: STR
+     udp-allow-rate: INT
+
+.. _mod-noudp_udp-allow-rate:
+
+udp-allow-rate
+..............
+
+Specifies how many UDP queries will pass the filter. Value 0 means none.
+A non-zero value means every N\ :sup:`th` UDP query will pass the filter.
+
 .. NOTE::
-   This module is not configurable.
+   The rate value is associated with one UDP worker. If more UDP workers are
+   configured, the specified value may not be obvious to clients.
+
+*Default:* 0