From dfa242407d14ff95cbc13f2e4df535b6a2cec5ca Mon Sep 17 00:00:00 2001
From: Nihanth Subramanya <nhnt11@gmail.com>
Date: Thu, 4 Jun 2020 19:02:33 +0000
Subject: [PATCH] Bug 1631609 - Test provider steering heuristic.
 r=valentin,johannh

Depends on D76618

Differential Revision: https://phabricator.services.mozilla.com/D76619
---
 .../doh-rollout/test/browser/browser.ini      |  1 +
 .../test/browser/browser_providerSteering.js  | 77 +++++++++++++++++++
 .../doh-rollout/test/browser/head.js          |  9 ++-
 3 files changed, 86 insertions(+), 1 deletion(-)
 create mode 100644 browser/extensions/doh-rollout/test/browser/browser_providerSteering.js

diff --git a/browser/extensions/doh-rollout/test/browser/browser.ini b/browser/extensions/doh-rollout/test/browser/browser.ini
index d045eeafc9ee..b270fa59e13e 100644
--- a/browser/extensions/doh-rollout/test/browser/browser.ini
+++ b/browser/extensions/doh-rollout/test/browser/browser.ini
@@ -8,6 +8,7 @@ skip-if = debug # Bug 1548006 - reloading an addon upsets the leak checker.
 [browser_doorhanger_newProfile.js]
 [browser_policyOverride.js]
 skip-if = (!debug && bits == 64) #Bug 1605297
+[browser_providerSteering.js]
 [browser_rollback.js]
 skip-if = asan && !debug && os == 'linux' && bits == 64 && os_version == '18.04' # Bug 1613994 for linux1804
 [browser_trrSelect.js]
diff --git a/browser/extensions/doh-rollout/test/browser/browser_providerSteering.js b/browser/extensions/doh-rollout/test/browser/browser_providerSteering.js
new file mode 100644
index 000000000000..73cb802c3f9a
--- /dev/null
+++ b/browser/extensions/doh-rollout/test/browser/browser_providerSteering.js
@@ -0,0 +1,77 @@
+"use strict";
+
+const TEST_DOMAIN = "doh.test";
+
+add_task(setup);
+
+add_task(async function testProviderSteering() {
+  setPassingHeuristics();
+  Preferences.set(prefs.DOH_ENABLED_PREF, true);
+  await BrowserTestUtils.waitForCondition(() => {
+    return Preferences.get(prefs.DOH_SELF_ENABLED_PREF);
+  });
+  is(Preferences.get(prefs.DOH_SELF_ENABLED_PREF), true, "Breadcrumb saved.");
+  await checkHeuristicsTelemetry("enable_doh", "startup");
+
+  let providerTestcases = [
+    {
+      name: "provider1",
+      canonicalName: "foo.provider1.com",
+      uri: "https://foo.provider1.com/query",
+    },
+    {
+      name: "provider2",
+      canonicalName: "bar.provider2.com",
+      uri: "https://bar.provider2.com/query",
+    },
+  ];
+  Preferences.set(
+    prefs.DOH_PROVIDER_STEERING_LIST_PREF,
+    JSON.stringify(providerTestcases)
+  );
+
+  for (let { name, canonicalName, uri } of providerTestcases) {
+    gDNSOverride.addIPOverride(TEST_DOMAIN, "9.9.9.9");
+    gDNSOverride.setCnameOverride(TEST_DOMAIN, canonicalName);
+    let trrURIChanged = TestUtils.topicObserved(
+      "network:trr-uri-changed",
+      () => {
+        // We need this check because this topic is observed once immediately
+        // after the network change when the URI is reset, and then when the
+        // provider steering heuristic runs and sets it to our uri.
+        return gDNSService.currentTrrURI == uri;
+      }
+    );
+    simulateNetworkChange();
+    await trrURIChanged;
+    is(gDNSService.currentTrrURI, uri, "TRR URI set to provider endpoint");
+    await checkHeuristicsTelemetry("enable_doh", "netchange", name);
+    gDNSOverride.clearHostOverride(TEST_DOMAIN);
+  }
+
+  let trrURIChanged = TestUtils.topicObserved("network:trr-uri-changed");
+  simulateNetworkChange();
+  await trrURIChanged;
+  is(
+    gDNSService.currentTrrURI,
+    "https://dummytrr.com/query",
+    "TRR URI set to auto-selected"
+  );
+  await checkHeuristicsTelemetry("enable_doh", "netchange");
+
+  // Set enterprise roots enabled and ensure provider steering is disabled.
+  Preferences.set("security.enterprise_roots.enabled", true);
+  gDNSOverride.setCnameOverride(
+    TEST_DOMAIN,
+    providerTestcases[0].canonicalName
+  );
+  simulateNetworkChange();
+  await checkHeuristicsTelemetry("disable_doh", "netchange");
+  is(
+    gDNSService.currentTrrURI,
+    "https://dummytrr.com/query",
+    "TRR URI set to auto-selected"
+  );
+  Preferences.reset("security.enterprise_roots.enabled");
+  gDNSOverride.clearHostOverride(TEST_DOMAIN);
+});
diff --git a/browser/extensions/doh-rollout/test/browser/head.js b/browser/extensions/doh-rollout/test/browser/head.js
index ccd12b5ef67d..d4be714d6d1b 100644
--- a/browser/extensions/doh-rollout/test/browser/head.js
+++ b/browser/extensions/doh-rollout/test/browser/head.js
@@ -52,6 +52,8 @@ const prefs = {
   DOH_TRR_SELECT_DRY_RUN_RESULT_PREF:
     "doh-rollout.trr-selection.dry-run-result",
   DOH_PROVIDER_STEERING_PREF: "doh-rollout.provider-steering.enabled",
+  DOH_PROVIDER_STEERING_LIST_PREF:
+    "doh-rollout.provider-steering.provider-list",
   PROFILE_CREATION_THRESHOLD_PREF: "doh-rollout.profileCreationThreshold",
 };
 
@@ -149,7 +151,11 @@ function ensureNoTRRSelectionTelemetry() {
   is(events.length, 0, "Found no trrselect events.");
 }
 
-async function checkHeuristicsTelemetry(decision, evaluateReason) {
+async function checkHeuristicsTelemetry(
+  decision,
+  evaluateReason,
+  steeredProvider = ""
+) {
   let events;
   await BrowserTestUtils.waitForCondition(() => {
     events = Services.telemetry.snapshotEvents(
@@ -165,6 +171,7 @@ async function checkHeuristicsTelemetry(decision, evaluateReason) {
   if (evaluateReason) {
     is(events[0][5].evaluateReason, evaluateReason, "Got the expected reason.");
   }
+  is(events[0][5].steeredProvider, steeredProvider, "Got expected provider.");
 
   // After checking the event, clear all telemetry. Since we check for a single
   // event above, this ensures all heuristics events are intentional and tested.