lua: ensure answer_clear() keeps original EDNS

Answers to EDNS requests from certain lua policies that use the
answer_clear() function would lack OPT RR and thus violate the MUST
condition in RFC6891.6.1.1.

8 files changed, 42 insertions, 8 deletions

diff --git a/NEWS b/NEWS
index 5abf1da0..12661505 100644
--- a/NEWS
+++ b/NEWS
@@ -5,6 +5,7 @@ Bugfixes
 --------
 - policy.rpz: improve logs, fix origin detection in files without $ORIGIN
 - lua: log() works again; broken in 5.4.2 (!1223)
+- policy: correctly include EDNS0 previously omitted by some actions (!1230)
 
 
 Knot Resolver 5.4.2 (2021-10-13)
diff --git a/daemon/lua/kres-gen-29.lua b/daemon/lua/kres-gen-29.lua
index 32795f2b..ded49465 100644
--- a/daemon/lua/kres-gen-29.lua
+++ b/daemon/lua/kres-gen-29.lua
@@ -378,6 +378,7 @@ int knot_pkt_put_rotate(knot_pkt_t *, uint16_t, const knot_rrset_t *, uint16_t,
 knot_pkt_t *knot_pkt_new(void *, uint16_t, knot_mm_t *);
 void knot_pkt_free(knot_pkt_t *);
 int knot_pkt_parse(knot_pkt_t *, unsigned int);
+knot_rrset_t *kr_request_ensure_edns(struct kr_request *);
 knot_pkt_t *kr_request_ensure_answer(struct kr_request *);
 struct kr_rplan *kr_resolve_plan(struct kr_request *);
 knot_mm_t *kr_resolve_pool(struct kr_request *);
diff --git a/daemon/lua/kres-gen-31.lua b/daemon/lua/kres-gen-31.lua
index f540ddf8..dc733472 100644
--- a/daemon/lua/kres-gen-31.lua
+++ b/daemon/lua/kres-gen-31.lua
@@ -378,6 +378,7 @@ int knot_pkt_put_rotate(knot_pkt_t *, uint16_t, const knot_rrset_t *, uint16_t,
 knot_pkt_t *knot_pkt_new(void *, uint16_t, knot_mm_t *);
 void knot_pkt_free(knot_pkt_t *);
 int knot_pkt_parse(knot_pkt_t *, unsigned int);
+knot_rrset_t *kr_request_ensure_edns(struct kr_request *);
 knot_pkt_t *kr_request_ensure_answer(struct kr_request *);
 struct kr_rplan *kr_resolve_plan(struct kr_request *);
 knot_mm_t *kr_resolve_pool(struct kr_request *);
diff --git a/daemon/lua/kres-gen.sh b/daemon/lua/kres-gen.sh
index eded99ea..b4244505 100755
--- a/daemon/lua/kres-gen.sh
+++ b/daemon/lua/kres-gen.sh
@@ -201,6 +201,7 @@ EOF
 ## libkres API
 ${CDEFS} ${LIBKRES} functions <<-EOF
 # Resolution request
+	kr_request_ensure_edns
 	kr_request_ensure_answer
 	kr_resolve_plan
 	kr_resolve_pool
diff --git a/daemon/lua/kres.lua b/daemon/lua/kres.lua
index ce355639..14042166 100644
--- a/daemon/lua/kres.lua
+++ b/daemon/lua/kres.lua
@@ -879,6 +879,11 @@ ffi.metatype( kr_request_t, {
 			req.vars_ref = ref
 			return var
 		end,
+		-- Ensure that answer has EDNS if needed; can't fail.
+		ensure_edns = function (req)
+			assert(ffi.istype(kr_request_t, req))
+			return C.kr_request_ensure_edns(req)
+		end,
 		-- Ensure that answer exists and return it; can't fail.
 		ensure_answer = function (req)
 			assert(ffi.istype(kr_request_t, req))
diff --git a/lib/resolve.c b/lib/resolve.c
index d9a92ec1..1ffd07ec 100644
--- a/lib/resolve.c
+++ b/lib/resolve.c
@@ -699,6 +699,27 @@ static int resolve_query(struct kr_request *request, const knot_pkt_t *packet)
 	return request->state;
 }
 
+knot_rrset_t* kr_request_ensure_edns(struct kr_request *request)
+{
+	kr_require(request && request->answer && request->qsource.packet && request->ctx);
+	knot_pkt_t* answer = request->answer;
+	bool want_edns = knot_pkt_has_edns(request->qsource.packet);
+	if (!want_edns) {
+		kr_assert(!answer->opt_rr);
+		return answer->opt_rr;
+	} else if (answer->opt_rr) {
+		return answer->opt_rr;
+	}
+
+	kr_assert(request->ctx->downstream_opt_rr);
+	answer->opt_rr = knot_rrset_copy(request->ctx->downstream_opt_rr, &answer->mm);
+	if (!answer->opt_rr)
+		return NULL;
+	if (knot_pkt_has_dnssec(request->qsource.packet))
+		knot_edns_set_do(answer->opt_rr);
+	return answer->opt_rr;
+}
+
 knot_pkt_t *kr_request_ensure_answer(struct kr_request *request)
 {
 	if (request->answer)
@@ -749,14 +770,8 @@ knot_pkt_t *kr_request_ensure_answer(struct kr_request *request)
 	}
 
 	// Prepare EDNS if required.
-	if (knot_pkt_has_edns(qs_pkt)) {
-		answer->opt_rr = knot_rrset_copy(request->ctx->downstream_opt_rr,
-						 &answer->mm);
-		if (!answer->opt_rr)
-			goto enomem; // answer is on mempool, so "leak" is OK
-		if (knot_pkt_has_dnssec(qs_pkt))
-			knot_edns_set_do(answer->opt_rr);
-	}
+	if (knot_pkt_has_edns(qs_pkt) && kr_fails_assert(kr_request_ensure_edns(request)))
+		goto enomem; // answer is on mempool, so "leak" is OK
 
 	return request->answer;
 enomem:
diff --git a/lib/resolve.h b/lib/resolve.h
index 88605b1e..efd66567 100644
--- a/lib/resolve.h
+++ b/lib/resolve.h
@@ -273,6 +273,15 @@ KR_EXPORT
 int kr_resolve_begin(struct kr_request *request, struct kr_context *ctx);
 
 /**
+ * Ensure that request->answer->opt_rr is present if query has EDNS.
+ *
+ * This function should be used after clearing a response packet to ensure its
+ * opt_rr is properly set. Returns the opt_rr (for convenience) or NULL.
+ */
+KR_EXPORT
+knot_rrset_t * kr_request_ensure_edns(struct kr_request *request);
+
+/**
  * Ensure that request->answer is usable, and return it (for convenience).
  *
  * It may return NULL, in which case it marks ->state with _FAIL and no answer will be sent.
diff --git a/modules/policy/policy.lua b/modules/policy/policy.lua
index 71d277b1..7cd8ff92 100644
--- a/modules/policy/policy.lua
+++ b/modules/policy/policy.lua
@@ -654,6 +654,7 @@ local function answer_clear(req)
 	local pkt = req:ensure_answer()
 	if pkt == nil then return nil end
 	pkt:clear_payload()
+	req:ensure_edns()
 	return pkt
 end