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

@@ -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}")`
+    );
+  }
 }

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", {