From 726f30aa648c2c0e36f91252e193f00a0187b046 Mon Sep 17 00:00:00 2001
From: Valentin Gosu <valentin.gosu@gmail.com>
Date: Wed, 7 Apr 2021 07:29:54 +0000
Subject: [PATCH] Bug 1700076 - Add ability to count number of DoH requests
 that reach the TRRServer r=necko-reviewers,dragana

Differential Revision: https://phabricator.services.mozilla.com/D109589
---
 netwerk/test/unit/head_trr.js              | 30 +++++++++++++++++++---
 netwerk/test/unit/test_trr_confirmation.js |  1 +
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/netwerk/test/unit/head_trr.js b/netwerk/test/unit/head_trr.js
index 1ec6e36d8d24..ad346c481622 100644
--- a/netwerk/test/unit/head_trr.js
+++ b/netwerk/test/unit/head_trr.js
@@ -247,6 +247,10 @@ class TRRServerCode {
     // value: array [answer1, answer2]
     global.dns_query_answers = {};
 
+    // key: domain
+    // value: a map containing {key: type, value: number of requests}
+    global.dns_query_counts = {};
+
     global.http2 = require("http2");
     global.server = global.http2.createSecureServer(options, global.handler);
 
@@ -257,6 +261,13 @@ class TRRServerCode {
 
     return global.server.address().port;
   }
+
+  static getRequestCount(domain, type) {
+    if (!global.dns_query_counts[domain]) {
+      return 0;
+    }
+    return global.dns_query_counts[domain][type] || 0;
+  }
 }
 
 /// This is the default handler for /dns-query
@@ -291,10 +302,15 @@ function trrQueryHandler(req, resp, url) {
 
   function processRequest(req, resp, payload) {
     let dnsQuery = global.dnsPacket.decode(payload);
-    let response =
-      global.dns_query_answers[
-        `${dnsQuery.questions[0].name}/${dnsQuery.questions[0].type}`
-      ] || {};
+    let domain = dnsQuery.questions[0].name;
+    let type = dnsQuery.questions[0].type;
+    let response = global.dns_query_answers[`${domain}/${type}`] || {};
+
+    if (!global.dns_query_counts[domain]) {
+      global.dns_query_counts[domain] = {};
+    }
+    global.dns_query_counts[domain][type] =
+      global.dns_query_counts[domain][type] + 1 || 1;
 
     let flags = global.dnsPacket.RECURSION_DESIRED;
     flags |= response.flags || 0;
@@ -399,4 +415,10 @@ class TRRServer {
     )}`;
     return this.execute(text);
   }
+
+  async requestCount(domain, type) {
+    return this.execute(
+      `TRRServerCode.getRequestCount("${domain}", "${type}")`
+    );
+  }
 }
diff --git a/netwerk/test/unit/test_trr_confirmation.js b/netwerk/test/unit/test_trr_confirmation.js
index 3c1c42a4ce27..3aaf878ec3f7 100644
--- a/netwerk/test/unit/test_trr_confirmation.js
+++ b/netwerk/test/unit/test_trr_confirmation.js
@@ -108,6 +108,7 @@ add_task(async function confirm_ok() {
     "Should be CONFIRM_TRYING_OK"
   );
   await new TRRDNSListener("example.com", { expectedAnswer: "1.2.3.4" });
+  equal(await trrServer.requestCount("example.com", "A"), 1);
   await waitForConfirmationState(CONFIRM_OK, 1000);
 
   await trrServer.registerDoHAnswers("confirm.example.com", "NS", {