/* Any copyright is dedicated to the Public Domain.
   http://creativecommons.org/publicdomain/zero/1.0/
*/

"use strict";

const { ContentTaskUtils } = ChromeUtils.import(
  "resource://testing-common/ContentTaskUtils.jsm"
);
const { TelemetryUtils } = ChromeUtils.import(
  "resource://gre/modules/TelemetryUtils.jsm"
);
const { ObjectUtils } = ChromeUtils.import(
  "resource://gre/modules/ObjectUtils.jsm"
);
const { PermissionTestUtils } = ChromeUtils.import(
  "resource://testing-common/PermissionTestUtils.jsm"
);

const HC_PERMISSION = "hc_telemetry";

async function waitForProcessesEvents(
  aProcesses,
  aAdditionalCondition = data => true
) {
  await ContentTaskUtils.waitForCondition(() => {
    const events = Services.telemetry.snapshotEvents(
      Ci.nsITelemetry.DATASET_PRERELEASE_CHANNELS
    );
    return (
      aProcesses.every(p => Object.keys(events).includes(p)) &&
      aAdditionalCondition(events)
    );
  });
}

/**
 * Wait for a specific event to appear in the given process data.
 * @param {String} aProcess the name of the process we expect the event to appear.
 * @param {Array} aEventData the event data to look for.
 * @return {Promise} Resolved when the event is found or rejected if the search
 *         times out.
 */
async function waitForEvent(aProcess, aEventData) {
  await waitForProcessesEvents([aProcess], events => {
    let processEvents = events[aProcess].map(e => e.slice(1));
    if (processEvents.length == 0) {
      return false;
    }

    return processEvents.find(e => ObjectUtils.deepEqual(e, aEventData));
  });
}

/**
 * Remove the trailing null/undefined from an event definition.
 * This is useful for comparing the sample events (that might
 * contain null/undefined) to the data from the snapshot (which might
 * filter them).
 */
function removeTrailingInvalidEntry(aEvent) {
  while (
    aEvent[aEvent.length - 1] === undefined ||
    aEvent[aEvent.length - 1] === null
  ) {
    aEvent.pop();
  }
  return aEvent;
}

add_task(async function test_setup() {
  // Make sure the newly spawned content processes will have extended Telemetry and
  // hybrid content telemetry enabled.
  await SpecialPowers.pushPrefEnv({
    set: [
      [TelemetryUtils.Preferences.OverridePreRelease, true],
      [TelemetryUtils.Preferences.HybridContentEnabled, true],
      [TelemetryUtils.Preferences.LogLevel, "Trace"],
    ],
  });
  // And take care of the already initialized one as well.
  let canRecordExtended = Services.telemetry.canRecordExtended;
  Services.telemetry.canRecordExtended = true;
  registerCleanupFunction(
    () => (Services.telemetry.canRecordExtended = canRecordExtended)
  );
});

add_task(async function test_untrusted_http_origin() {
  Services.telemetry.clearEvents();

  // Install a custom handler that intercepts hybrid content telemetry messages
  // and makes the test fail. We don't expect any message from non secure contexts.
  const messageName = "HybridContentTelemetry:onTelemetryMessage";
  let makeTestFail = () => ok(false, `Received an unexpected ${messageName}.`);
  Services.mm.addMessageListener(messageName, makeTestFail);

  // Try to use the API on a non-secure host.
  const testHost = "http://example.org";
  let testHttpUri = Services.io.newURI(testHost);
  PermissionTestUtils.add(
    testHttpUri,
    HC_PERMISSION,
    Services.perms.ALLOW_ACTION
  );
  let url = getRootDirectory(gTestPath) + "hybrid_content.html";
  url = url.replace("chrome://mochitests/content", testHost);
  let newTab = await BrowserTestUtils.openNewForegroundTab(gBrowser, url);

  // Try to use the API. Also record a content event from outside HCT: we'll
  // use this event to know when we can stop waiting for the hybrid content data.
  const TEST_CONTENT_EVENT = ["telemetry.test", "main_and_content", "object1"];
  Services.telemetry.setEventRecordingEnabled("telemetry.test", true);
  await ContentTask.spawn(
    newTab.linkedBrowser,
    [TEST_CONTENT_EVENT],
    ([testContentEvent]) => {
      // Call the hybrid content telemetry API.
      let contentWin = Cu.waiveXrays(content);
      contentWin.testRegisterEvents(testContentEvent[0], JSON.stringify({}));
      // Record from the usual telemetry API a "canary" event.
      Services.telemetry.recordEvent(...testContentEvent);
    }
  );

  // Let's support both e10s/non-e10s testing.
  const processName = Services.appinfo.browserTabsRemoteAutostart
    ? "content"
    : "parent";
  await waitForEvent(processName, TEST_CONTENT_EVENT);

  // This is needed otherwise the test will fail due to missing test passes.
  ok(true, "The untrusted HTTP page was not able to use the API.");

  // Finally clean up the listener.
  BrowserTestUtils.removeTab(newTab);
  PermissionTestUtils.remove(testHttpUri, HC_PERMISSION);
  Services.mm.removeMessageListener(messageName, makeTestFail);
  Services.telemetry.setEventRecordingEnabled("telemetry.test", false);
});

add_task(async function test_secure_non_whitelisted_origin() {
  Services.telemetry.clearEvents();

  // Install a custom handler that intercepts hybrid content telemetry messages
  // and makes the test fail. We don't expect any message from non whitelisted pages.
  const messageName = "HybridContentTelemetry:onTelemetryMessage";
  let makeTestFail = () => ok(false, `Received an unexpected ${messageName}.`);
  Services.mm.addMessageListener(messageName, makeTestFail);

  // Try to use the API on a secure host but don't give the page enough privileges.
  const testHost = "https://example.org";
  let url = getRootDirectory(gTestPath) + "hybrid_content.html";
  url = url.replace("chrome://mochitests/content", testHost);
  let newTab = await BrowserTestUtils.openNewForegroundTab(gBrowser, url);

  // Try to use the API. Also record a content event from outside HCT: we'll
  // use this event to know when we can stop waiting for the hybrid content data.
  const TEST_CONTENT_EVENT = ["telemetry.test", "main_and_content", "object1"];
  Services.telemetry.setEventRecordingEnabled("telemetry.test", true);
  await ContentTask.spawn(
    newTab.linkedBrowser,
    [TEST_CONTENT_EVENT],
    ([testContentEvent]) => {
      // Call the hybrid content telemetry API.
      let contentWin = Cu.waiveXrays(content);
      contentWin.testRegisterEvents(testContentEvent[0], JSON.stringify({}));
      // Record from the usual telemetry API a "canary" event.
      Services.telemetry.recordEvent(...testContentEvent);
    }
  );

  // Let's support both e10s/non-e10s testing.
  const processName = Services.appinfo.browserTabsRemoteAutostart
    ? "content"
    : "parent";
  await waitForEvent(processName, TEST_CONTENT_EVENT);

  // This is needed otherwise the test will fail due to missing test passes.
  ok(true, "The HTTPS page without permission was not able to use the API.");

  // Finally clean up the listener.
  BrowserTestUtils.removeTab(newTab);
  Services.mm.removeMessageListener(messageName, makeTestFail);
  Services.telemetry.setEventRecordingEnabled("telemetry.test", false);
});

add_task(async function test_trusted_disabled_hybrid_telemetry() {
  Services.telemetry.clearEvents();

  // This test requires hybrid content telemetry to be disabled.
  await SpecialPowers.pushPrefEnv({
    set: [[TelemetryUtils.Preferences.HybridContentEnabled, false]],
  });

  // Install a custom handler that intercepts hybrid content telemetry messages
  // and makes the test fail. We don't expect any message when the API is disabled.
  const messageName = "HybridContentTelemetry:onTelemetryMessage";
  let makeTestFail = () => ok(false, `Received an unexpected ${messageName}.`);
  Services.mm.addMessageListener(messageName, makeTestFail);

  // Try to use the API on a secure host.
  const testHost = "https://example.org";
  let testHttpsUri = Services.io.newURI(testHost);
  PermissionTestUtils.add(
    testHttpsUri,
    HC_PERMISSION,
    Services.perms.ALLOW_ACTION
  );
  let url = getRootDirectory(gTestPath) + "hybrid_content.html";
  url = url.replace("chrome://mochitests/content", testHost);
  let newTab = await BrowserTestUtils.openNewForegroundTab(gBrowser, url);

  // Try to use the API. Also record a content event from outside HCT: we'll
  // use this event to know when we can stop waiting for the hybrid content data.
  const TEST_CONTENT_EVENT = ["telemetry.test", "main_and_content", "object1"];
  Services.telemetry.setEventRecordingEnabled("telemetry.test", true);
  await ContentTask.spawn(
    newTab.linkedBrowser,
    [TEST_CONTENT_EVENT],
    ([testContentEvent]) => {
      // Call the hybrid content telemetry API.
      let contentWin = Cu.waiveXrays(content);
      contentWin.testRegisterEvents(testContentEvent[0], JSON.stringify({}));
      // Record from the usual telemetry API a "canary" event.
      Services.telemetry.recordEvent(...testContentEvent);
    }
  );

  // Let's support both e10s/non-e10s testing.
  const processName = Services.appinfo.browserTabsRemoteAutostart
    ? "content"
    : "parent";
  await waitForEvent(processName, TEST_CONTENT_EVENT);

  // This is needed otherwise the test will fail due to missing test passes.
  ok(true, "There were no unintended hybrid content API usages.");

  // Finally clean up the listener.
  await SpecialPowers.popPrefEnv();
  BrowserTestUtils.removeTab(newTab);
  PermissionTestUtils.remove(testHttpsUri, HC_PERMISSION);
  Services.mm.removeMessageListener(messageName, makeTestFail);
  Services.telemetry.setEventRecordingEnabled("telemetry.test", false);
});

add_task(async function test_hybrid_content_with_iframe() {
  Services.telemetry.clearEvents();

  // Open a trusted page that can use in the HCT in a new tab.
  const testOuterPageHost = "https://example.com";
  let testHttpsUri = Services.io.newURI(testOuterPageHost);
  PermissionTestUtils.add(
    testHttpsUri,
    HC_PERMISSION,
    Services.perms.ALLOW_ACTION
  );
  let url = getRootDirectory(gTestPath) + "hybrid_content.html";
  let outerUrl = url.replace("chrome://mochitests/content", testOuterPageHost);
  let newTab = await BrowserTestUtils.openNewForegroundTab(gBrowser, outerUrl);

  // Install a custom handler that intercepts hybrid content telemetry messages
  // and makes the test fail. This needs to be done after the tab is opened.
  const messageName = "HybridContentTelemetry:onTelemetryMessage";
  let makeTestFail = () => ok(false, `Received an unexpected ${messageName}.`);
  Services.mm.addMessageListener(messageName, makeTestFail);

  // Enable recording the canary event.
  const TEST_CONTENT_EVENT = ["telemetry.test", "main_and_content", "object1"];
  Services.telemetry.setEventRecordingEnabled("telemetry.test", true);

  // Add an iframe to the test page. The URI in the iframe should not be able
  // to use HCT to the the missing privileges.
  const testHost = "https://example.org";
  let iframeUrl = url.replace("chrome://mochitests/content", testHost);
  await ContentTask.spawn(
    newTab.linkedBrowser,
    [iframeUrl, TEST_CONTENT_EVENT],
    async function([iframeUrl, testContentEvent]) {
      let doc = content.document;
      let iframe = doc.createElement("iframe");
      let promiseIframeLoaded = ContentTaskUtils.waitForEvent(
        iframe,
        "load",
        false
      );
      iframe.src = iframeUrl;
      doc.body.insertBefore(iframe, doc.body.firstChild);
      await promiseIframeLoaded;

      // Call the hybrid content telemetry API.
      let contentWin = Cu.waiveXrays(iframe.contentWindow);
      contentWin.testRegisterEvents(testContentEvent[0], JSON.stringify({}));

      // Record from the usual telemetry API a "canary" event.
      Services.telemetry.recordEvent(...testContentEvent);
    }
  );

  // Let's support both e10s/non-e10s testing.
  const processName = Services.appinfo.browserTabsRemoteAutostart
    ? "content"
    : "parent";
  await waitForEvent(processName, TEST_CONTENT_EVENT);

  // This is needed otherwise the test will fail due to missing test passes.
  ok(
    true,
    "There were no unintended hybrid content API usages from the iframe."
  );

  // Cleanup permissions and remove the tab.
  BrowserTestUtils.removeTab(newTab);
  Services.mm.removeMessageListener(messageName, makeTestFail);
  PermissionTestUtils.remove(testHttpsUri, HC_PERMISSION);
  Services.telemetry.setEventRecordingEnabled("telemetry.test", false);
});

add_task(async function test_hybrid_content_recording() {
  const testHost = "https://example.org";
  const TEST_EVENT_CATEGORY = "telemetry.test.hct";
  const RECORDED_TEST_EVENTS = [
    [TEST_EVENT_CATEGORY, "test1", "object1"],
    [
      TEST_EVENT_CATEGORY,
      "test2",
      "object1",
      null,
      { key1: "foo", key2: "bar" },
    ],
    [TEST_EVENT_CATEGORY, "test2", "object1", "some value"],
    [TEST_EVENT_CATEGORY, "test1", "object1", null, null],
    [TEST_EVENT_CATEGORY, "test1", "object1", "", null],
  ];
  const NON_RECORDED_TEST_EVENTS = [
    [TEST_EVENT_CATEGORY, "unknown", "unknown"],
  ];

  Services.telemetry.clearEvents();

  // Give the test host enough privileges to use the API and open the test page.
  let testHttpsUri = Services.io.newURI(testHost);
  PermissionTestUtils.add(
    testHttpsUri,
    HC_PERMISSION,
    Services.perms.ALLOW_ACTION
  );
  let url = getRootDirectory(gTestPath) + "hybrid_content.html";
  url = url.replace("chrome://mochitests/content", testHost);
  let newTab = await BrowserTestUtils.openNewForegroundTab(gBrowser, url);

  // Register some events and record them in Telemetry.
  await ContentTask.spawn(
    newTab.linkedBrowser,
    [TEST_EVENT_CATEGORY, RECORDED_TEST_EVENTS, NON_RECORDED_TEST_EVENTS],
    ([eventCategory, recordedTestEvents, nonRecordedTestEvents]) => {
      let contentWin = Cu.waiveXrays(content);

      // If we tried to call contentWin.Mozilla.ContentTelemetry.* functions
      // and pass non-string parameters, |waiveXrays| would complain and not
      // let us access them. To work around this, we generate test functions
      // in the test HTML file and unwrap the passed JSON blob there.
      contentWin.testRegisterEvents(
        eventCategory,
        JSON.stringify({
          // Event with only required fields.
          test1: {
            methods: ["test1"],
            objects: ["object1"],
          },
          // Event with extra_keys.
          test2: {
            methods: ["test2", "test2b"],
            objects: ["object1"],
            extra_keys: ["key1", "key2"],
          },
        })
      );

      // Record some valid events.
      recordedTestEvents.forEach(e =>
        contentWin.testRecordEvents(JSON.stringify(e))
      );

      // Test recording an unknown event. The Telemetry API itself is supposed to throw,
      // but we catch that in hybrid content telemetry and log an error message.
      nonRecordedTestEvents.forEach(e =>
        contentWin.testRecordEvents(JSON.stringify(e))
      );
    }
  );

  // Wait for the data to be in the snapshot, then get the Telemetry data.
  await waitForProcessesEvents(["dynamic"]);
  let snapshot = Services.telemetry.snapshotEvents(
    Ci.nsITelemetry.DATASET_PRERELEASE_CHANNELS
  );

  // Check that the dynamically register events made it to the snapshot.
  ok(
    "dynamic" in snapshot,
    "The snapshot must contain the 'dynamic' process section"
  );
  let dynamicEvents = snapshot.dynamic.map(e => e.slice(1));
  is(
    dynamicEvents.length,
    RECORDED_TEST_EVENTS.length,
    "Should match expected event count."
  );
  for (let i = 0; i < RECORDED_TEST_EVENTS.length; ++i) {
    SimpleTest.isDeeply(
      dynamicEvents[i],
      removeTrailingInvalidEntry(RECORDED_TEST_EVENTS[i]),
      "Should have recorded the expected event."
    );
  }

  // Cleanup permissions and remove the tab.
  BrowserTestUtils.removeTab(newTab);
  PermissionTestUtils.remove(testHttpsUri, HC_PERMISSION);
});

add_task(async function test_can_upload() {
  const testHost = "https://example.org";

  await SpecialPowers.pushPrefEnv({
    set: [[TelemetryUtils.Preferences.FhrUploadEnabled, true]],
  });

  // Give the test host enough privileges to use the API and open the test page.
  let testHttpsUri = Services.io.newURI(testHost);
  PermissionTestUtils.add(
    testHttpsUri,
    HC_PERMISSION,
    Services.perms.ALLOW_ACTION
  );
  let url = getRootDirectory(gTestPath) + "hybrid_content.html";
  url = url.replace("chrome://mochitests/content", testHost);
  let newTab = await BrowserTestUtils.openNewForegroundTab(gBrowser, url);

  // Check that CanUpload reports the correct value.
  await ContentTask.spawn(newTab.linkedBrowser, {}, async function() {
    let contentWin = Cu.waiveXrays(content);

    await contentWin.Mozilla.ContentTelemetry.initPromise();

    // We don't need to pass any parameter, we can safely call Mozilla.ContentTelemetry.
    let canUpload = contentWin.Mozilla.ContentTelemetry.canUpload();
    ok(
      canUpload,
      "CanUpload must report 'true' if the preference has that value."
    );
  });

  // Flip the pref and check again.
  await SpecialPowers.pushPrefEnv({
    set: [[TelemetryUtils.Preferences.FhrUploadEnabled, false]],
  });
  await ContentTask.spawn(newTab.linkedBrowser, {}, async function() {
    let contentWin = Cu.waiveXrays(content);
    await contentWin.Mozilla.ContentTelemetry.initPromise();
    let canUpload = contentWin.Mozilla.ContentTelemetry.canUpload();
    ok(
      !canUpload,
      "CanUpload must report 'false' if the preference has that value."
    );
  });

  // Cleanup permissions and remove the tab.
  BrowserTestUtils.removeTab(newTab);
  PermissionTestUtils.remove(testHttpsUri, HC_PERMISSION);
});

add_task(async function test_init_rejects() {
  await SpecialPowers.pushPrefEnv({
    set: [[TelemetryUtils.Preferences.FhrUploadEnabled, true]],
  });

  // Give the test host no privilege: init() will throw.
  const testHostNoPrivilege = "https://example.org";
  let url = getRootDirectory(gTestPath) + "hybrid_content.html";
  url = url.replace("chrome://mochitests/content", testHostNoPrivilege);
  let newTab = await BrowserTestUtils.openNewForegroundTab(gBrowser, url);

  // Check that CanUpload reports false if we threw. We don't block on spawning
  // the content task, as we want to spawn another tab that loads an URI with HCT
  // privileges in order to make sure that init works ok there. We want both pages
  // to load in parallel.
  let page1Promise = ContentTask.spawn(
    newTab.linkedBrowser,
    {},
    async function() {
      let contentWin = Cu.waiveXrays(content);

      await Assert.rejects(
        contentWin.Mozilla.ContentTelemetry.initPromise(),
        /Origin not trusted/,
        "The init promise must reject if the page has no HCT permission."
      );

      // We don't need to pass any parameter, we can safely call Mozilla.ContentTelemetry.
      let canUpload = contentWin.Mozilla.ContentTelemetry.canUpload();
      ok(!canUpload, "CanUpload must report 'false' if init failed.");
    }
  );

  // Give the test host HCT privileges and test that init doesn't throw.
  const testHostPrivileges = "https://example.com";
  let testUrlPrivileges = Services.io.newURI(testHostPrivileges);
  PermissionTestUtils.add(
    testUrlPrivileges,
    HC_PERMISSION,
    Services.perms.ALLOW_ACTION
  );
  let urlWithPrivs = getRootDirectory(gTestPath) + "hybrid_content.html";
  urlWithPrivs = urlWithPrivs.replace(
    "chrome://mochitests/content",
    testHostPrivileges
  );
  let otherTab = await BrowserTestUtils.openNewForegroundTab(
    gBrowser,
    urlWithPrivs
  );

  // Check that CanUpload reports the correct value.
  let page2Promise = ContentTask.spawn(
    otherTab.linkedBrowser,
    {},
    async function() {
      let contentWin = Cu.waiveXrays(content);

      await contentWin.Mozilla.ContentTelemetry.initPromise();

      // We don't need to pass any parameter, we can safely call Mozilla.ContentTelemetry.
      let canUpload = contentWin.Mozilla.ContentTelemetry.canUpload();
      ok(
        canUpload,
        "CanUpload must report the expected value if init succeeded."
      );
    }
  );

  await Promise.all([page1Promise, page2Promise]);

  // Cleanup permissions and remove the tab.
  BrowserTestUtils.removeTab(newTab);
  BrowserTestUtils.removeTab(otherTab);
  PermissionTestUtils.remove(testUrlPrivileges, HC_PERMISSION);
});