From 45e38b3df3e5aee9c93de0fe7a78b20811e65bae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Vavru=C5=A1a?= <mvavrusa@cloudflare.com>
Date: Thu, 5 Apr 2018 22:43:57 -0700
Subject: [PATCH] check per-query flags instead of global options, getter for
 NS name

Checking query flags instead of global context option allows setting
overrides on individual queries. The effect is the same as query flags
start by copying request flags which start by copying context options.
---
 daemon/lua/kres-gen.lua |  2 ++
 daemon/lua/kres-gen.sh  |  2 ++
 daemon/lua/kres.lua     |  9 +++++++++
 lib/layer/iterate.h     |  1 +
 lib/nsrep.c             | 11 +++++------
 lib/nsrep.h             |  1 -
 lib/resolve.c           | 10 ++++++----
 lib/zonecut.c           |  4 ++--
 8 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/daemon/lua/kres-gen.lua b/daemon/lua/kres-gen.lua
index 911e2e090..885076360 100644
--- a/daemon/lua/kres-gen.lua
+++ b/daemon/lua/kres-gen.lua
@@ -274,6 +274,7 @@ int kr_rplan_pop(struct kr_rplan *, struct kr_query *);
 struct kr_query *kr_rplan_resolved(struct kr_rplan *);
 int kr_nsrep_set(struct kr_query *, size_t, const struct sockaddr *);
 uint32_t kr_rand_uint(uint32_t);
+int kr_make_query(struct kr_query *query, knot_pkt_t *pkt);
 void kr_pkt_make_auth_header(knot_pkt_t *);
 int kr_pkt_put(knot_pkt_t *, const knot_dname_t *, uint32_t, uint16_t, uint16_t, const uint8_t *, uint16_t);
 int kr_pkt_recycle(knot_pkt_t *);
@@ -294,6 +295,7 @@ void kr_qflags_clear(struct kr_qflags *, struct kr_qflags);
 int kr_zonecut_add(struct kr_zonecut *, const knot_dname_t *, const knot_rdata_t *);
 _Bool kr_zonecut_is_empty(struct kr_zonecut *);
 void kr_zonecut_set(struct kr_zonecut *, const knot_dname_t *);
+const knot_dname_t *kr_zonecut_find_nsname(struct kr_zonecut *);
 uint64_t kr_now();
 knot_rrset_t *kr_ta_get(map_t *, const knot_dname_t *);
 int kr_ta_add(map_t *, const knot_dname_t *, uint16_t, uint32_t, const uint8_t *, uint16_t);
diff --git a/daemon/lua/kres-gen.sh b/daemon/lua/kres-gen.sh
index fa2f11461..5752f2e03 100755
--- a/daemon/lua/kres-gen.sh
+++ b/daemon/lua/kres-gen.sh
@@ -142,6 +142,7 @@ EOF
 	kr_nsrep_set
 # Utils
 	kr_rand_uint
+	kr_make_query
 	kr_pkt_make_auth_header
 	kr_pkt_put
 	kr_pkt_recycle
@@ -162,6 +163,7 @@ EOF
 	kr_zonecut_add
 	kr_zonecut_is_empty
 	kr_zonecut_set
+	kr_zonecut_find_nsname
 	kr_now
 # Trust anchors
 	kr_ta_get
diff --git a/daemon/lua/kres.lua b/daemon/lua/kres.lua
index 6e9cf8da4..65ecad9e4 100644
--- a/daemon/lua/kres.lua
+++ b/daemon/lua/kres.lua
@@ -627,10 +627,19 @@ ffi.metatype( knot_pkt_t, {
 local kr_query_t = ffi.typeof('struct kr_query')
 ffi.metatype( kr_query_t, {
 	__index = {
+		-- Return query domain name
 		name = function(qry)
 			assert(ffi.istype(kr_query_t, qry))
 			return dname2wire(qry.sname)
 		end,
+		-- Write this query into packet
+		write = function(qry, pkt)
+			assert(ffi.istype(kr_query_t, qry))
+			assert(ffi.istype(knot_pkt_t, pkt))
+			local ret = C.kr_make_query(qry, pkt)
+			if ret ~= 0 then return nil, knot_error_t(ret) end
+			return true
+		end,
 	},
 })
 -- Metatype for request
diff --git a/lib/layer/iterate.h b/lib/layer/iterate.h
index 914c97237..a9395bfc1 100644
--- a/lib/layer/iterate.h
+++ b/lib/layer/iterate.h
@@ -32,4 +32,5 @@ enum {
 int kr_response_classify(const knot_pkt_t *pkt);
 
 /** Make next iterative query. */
+KR_EXPORT
 int kr_make_query(struct kr_query *query, knot_pkt_t *pkt);
diff --git a/lib/nsrep.c b/lib/nsrep.c
index 0420bdae0..9a6824c4b 100644
--- a/lib/nsrep.c
+++ b/lib/nsrep.c
@@ -83,10 +83,9 @@ static void update_nsrep_set(struct kr_nsrep *ns, const knot_dname_t *name, uint
 #undef ADDR_SET
 
 static unsigned eval_addr_set(const pack_t *addr_set, struct kr_context *ctx,
-			      unsigned score, uint8_t *addr[])
+			      struct kr_qflags opts, unsigned score, uint8_t *addr[])
 {
 	kr_nsrep_rtt_lru_t *rtt_cache = ctx->cache_rtt;
-	struct kr_qflags opts = ctx->options;
 	kr_nsrep_rtt_lru_entry_t *rtt_cache_entry_ptr[KR_NSREP_MAXADDR] = { NULL, };
 	assert (KR_NSREP_MAXADDR >= 2);
 	unsigned rtt_cache_entry_score[KR_NSREP_MAXADDR] = { score, KR_NS_MAX_SCORE + 1, };
@@ -237,16 +236,16 @@ static int eval_nsrep(const knot_dname_t *owner, const pack_t *addr_set, struct
 			if (reputation & KR_NS_NOIP4) {
 				score = KR_NS_UNKNOWN;
 				/* Try to start with clean slate */
-				if (!(ctx->options.NO_IPV6)) {
+				if (!(qry->flags.NO_IPV6)) {
 					reputation &= ~KR_NS_NOIP6;
 				}
-				if (!(ctx->options.NO_IPV4)) {
+				if (!(qry->flags.NO_IPV4)) {
 					reputation &= ~KR_NS_NOIP4;
 				}
 			}
 		}
 	} else {
-		score = eval_addr_set(addr_set, ctx, score, addr_choice);
+		score = eval_addr_set(addr_set, ctx, qry->flags, score, addr_choice);
 	}
 
 	/* Probabilistic bee foraging strategy (naive).
@@ -372,7 +371,7 @@ int kr_nsrep_elect_addr(struct kr_query *qry, struct kr_context *ctx)
 	}
 	/* Evaluate addr list */
 	uint8_t *addr_choice[KR_NSREP_MAXADDR] = { NULL, };
-	unsigned score = eval_addr_set(addr_set, ctx, ns->score, addr_choice);
+	unsigned score = eval_addr_set(addr_set, ctx, qry->flags, ns->score, addr_choice);
 	update_nsrep_set(ns, ns->name, addr_choice, score);
 	return kr_ok();
 }
diff --git a/lib/nsrep.h b/lib/nsrep.h
index 36aea3a64..f1b3865d2 100644
--- a/lib/nsrep.h
+++ b/lib/nsrep.h
@@ -184,4 +184,3 @@ int kr_nsrep_copy_set(struct kr_nsrep *dst, const struct kr_nsrep *src);
  */
 KR_EXPORT
 int kr_nsrep_sort(struct kr_nsrep *ns, kr_nsrep_rtt_lru_t *rtt_cache);
-
diff --git a/lib/resolve.c b/lib/resolve.c
index eed3fd902..f93e30cf3 100644
--- a/lib/resolve.c
+++ b/lib/resolve.c
@@ -1370,10 +1370,12 @@ int kr_resolve_produce(struct kr_request *request, struct sockaddr **dst, int *t
 		/* Caller is interested in always tracking a zone cut, even if the answer is cached
 		 * this is normally not required, and incurrs another cache lookups for cached answer. */
 		if (qry->flags.ALWAYS_CUT) {
-			switch(zone_cut_check(request, qry, packet)) {
-			case KR_STATE_FAIL: return KR_STATE_FAIL;
-			case KR_STATE_DONE: return KR_STATE_PRODUCE;
-			default: break;
+			if (!(qry->flags.STUB)) {
+				switch(zone_cut_check(request, qry, packet)) {
+				case KR_STATE_FAIL: return KR_STATE_FAIL;
+				case KR_STATE_DONE: return KR_STATE_PRODUCE;
+				default: break;
+				}
 			}
 		}
 		/* Resolve current query and produce dependent or finish */
diff --git a/lib/zonecut.c b/lib/zonecut.c
index 73218d036..120465ba4 100644
--- a/lib/zonecut.c
+++ b/lib/zonecut.c
@@ -346,10 +346,10 @@ static int fetch_ns(struct kr_context *ctx, struct kr_zonecut *cut,
 		unsigned *cached = lru_get_try(ctx->cache_rep,
 				(const char *)ns_name, knot_dname_size(ns_name));
 		unsigned reputation = (cached) ? *cached : 0;
-		if (!(reputation & KR_NS_NOIP4) && !(ctx->options.NO_IPV4)) {
+		if (!(reputation & KR_NS_NOIP4) && !(qry->flags.NO_IPV4)) {
 			fetch_addr(cut, &ctx->cache, ns_name, KNOT_RRTYPE_A, qry);
 		}
-		if (!(reputation & KR_NS_NOIP6) && !(ctx->options.NO_IPV6)) {
+		if (!(reputation & KR_NS_NOIP6) && !(qry->flags.NO_IPV6)) {
 			fetch_addr(cut,  &ctx->cache, ns_name, KNOT_RRTYPE_AAAA, qry);
 		}
 	}
-- 
GitLab