From 10dee6cee173bbf754ba903d1c018bcc6ddab78a Mon Sep 17 00:00:00 2001
From: TCY16 <tom@nlnetlabs.nl>
Date: Thu, 19 May 2022 14:55:26 +0200
Subject: [PATCH] create random cookies by creating and storing a random state
 in infra_cache

---
 services/cache/infra.c | 23 +++++++++++++++++------
 services/cache/infra.h |  4 +++-
 testcode/unitmain.c    |  3 ++-
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/services/cache/infra.c b/services/cache/infra.c
index e1c784540..05e2e0a82 100644
--- a/services/cache/infra.c
+++ b/services/cache/infra.c
@@ -230,7 +230,7 @@ setup_domain_limits(struct infra_cache* infra, struct config_file* cfg)
 }
 
 struct infra_cache* 
-infra_create(struct config_file* cfg)
+infra_create(struct config_file* cfg, struct ub_randstate* rnd)
 {
 	struct infra_cache* infra = (struct infra_cache*)calloc(1, 
 		sizeof(struct infra_cache));
@@ -270,6 +270,11 @@ infra_create(struct config_file* cfg)
 		infra_delete(infra);
 		return NULL;
 	}
+	if (!rnd) {
+		infra_delete(infra);
+		return NULL;
+	}
+	infra->random_state = rnd;
 	return infra;
 }
 
@@ -299,7 +304,7 @@ infra_adjust(struct infra_cache* infra, struct config_file* cfg)
 {
 	size_t maxmem;
 	if(!infra)
-		return infra_create(cfg);
+		return infra_create(cfg, ub_initstate(NULL));
 	infra->host_ttl = cfg->host_ttl;
 	infra->infra_keep_probing = cfg->infra_keep_probing;
 	infra_dp_ratelimit = cfg->ratelimit;
@@ -315,7 +320,7 @@ infra_adjust(struct infra_cache* infra, struct config_file* cfg)
 	   !slabhash_is_size(infra->client_ip_rates, cfg->ip_ratelimit_size,
 	   	cfg->ip_ratelimit_slabs)) {
 		infra_delete(infra);
-		infra = infra_create(cfg);
+		infra = infra_create(cfg, ub_initstate(NULL));
 	} else {
 		/* reapply domain limits */
 		traverse_postorder(&infra->domain_limits, domain_limit_free,
@@ -383,7 +388,14 @@ static void
 data_entry_init(struct infra_cache* infra, struct lruhash_entry* e, 
 	time_t timenow)
 {
-	uint8_t cookie[8] = {1,2,3,4,5,6,7,8};
+	int i;
+	uint8_t cookie[8] = {0,0,0,0,0,0,0,0};
+
+	for (i = 0; i < 8; i++) {
+		cookie[i] = ub_random_max(infra->random_state,
+			255); // @TODO is this correct? macro-ify?
+	}
+
 
 	struct infra_data* data = (struct infra_data*)e->data;
 	data->ttl = timenow + infra->host_ttl;
@@ -391,10 +403,9 @@ data_entry_init(struct infra_cache* infra, struct lruhash_entry* e,
 	data->edns_version = 0;
 	data->edns_lame_known = 0;
 	data->probedelay = 0;
-	// @TODO create "random" cookie
+	/* set EDNS cookie to zero, as this also sets the starting state*/
 	memset(&data->cookie, 0, sizeof(struct edns_cookie));
 	memcpy(data->cookie.data.components.client, cookie, 8);
-
 	data->isdnsseclame = 0;
 	data->rec_lame = 0;
 	data->lame_type_A = 0;
diff --git a/services/cache/infra.h b/services/cache/infra.h
index 90761eaa7..c8052eac9 100644
--- a/services/cache/infra.h
+++ b/services/cache/infra.h
@@ -162,6 +162,8 @@ struct infra_cache {
 	rbtree_type domain_limits;
 	/** hash table with query rates per client ip: ip_rate_key, ip_rate_data */
 	struct slabhash* client_ip_rates;
+        /** random state used in new entries for creating EDNS cookies (RFC9018) */
+        struct ub_randstate* random_state;
 };
 
 /** ratelimit, unless overridden by domain_limits, 0 is off */
@@ -236,7 +238,7 @@ struct rate_data {
  * @param cfg: config parameters or NULL for defaults.
  * @return: new infra cache, or NULL.
  */
-struct infra_cache* infra_create(struct config_file* cfg);
+struct infra_cache* infra_create(struct config_file* cfg, struct ub_randstate* rnd);
 
 /**
  * Delete infra cache.
diff --git a/testcode/unitmain.c b/testcode/unitmain.c
index 16aa88450..fd906df96 100644
--- a/testcode/unitmain.c
+++ b/testcode/unitmain.c
@@ -469,11 +469,12 @@ infra_test(void)
 	struct infra_key* k;
 	struct infra_data* d;
 	int init = 376;
+	struct ub_randstate* rnd = ub_initstate(NULL);
 
 	unit_show_feature("infra cache");
 	unit_assert(ipstrtoaddr("127.0.0.1", 53, &one, &onelen));
 
-	slab = infra_create(cfg);
+	slab = infra_create(cfg, rnd);
 	unit_assert( infra_host(slab, &one, onelen, zone, zonelen, now,
 		&vs, &edns_lame, &to) );
 	unit_assert( vs == 0 && to == init && edns_lame == 0 );
-- 
2.47.2