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