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

var unsecureEmptyURL =
  "http://example.org/browser/toolkit/components/antitracking/test/browser/empty.html";
var secureEmptyURL =
  "https://example.org/browser/toolkit/components/antitracking/test/browser/empty.html";
var secureAnotherEmptyURL =
  "https://example.com/browser/toolkit/components/antitracking/test/browser/empty.html";
var secureURL =
  "https://example.com/browser/toolkit/components/antitracking/test/browser/browser_staticPartition_HSTS.sjs";
var unsecureURL =
  "http://example.com/browser/toolkit/components/antitracking/test/browser/browser_staticPartition_HSTS.sjs";
var secureImgURL =
  "https://example.com/browser/toolkit/components/antitracking/test/browser/browser_staticPartition_HSTS.sjs?image";
var unsecureImgURL =
  "http://example.com/browser/toolkit/components/antitracking/test/browser/browser_staticPartition_HSTS.sjs?image";
var secureIncludeSubURL =
  "https://example.com/browser/toolkit/components/antitracking/test/browser/browser_staticPartition_HSTS.sjs?includeSub";
var unsecureSubEmptyURL =
  "http://test1.example.com/browser/toolkit/components/antitracking/test/browser/empty.html";
var secureSubEmptyURL =
  "https://test1.example.com/browser/toolkit/components/antitracking/test/browser/empty.html";
var unsecureNoCertSubEmptyURL =
  "http://nocert.example.com/browser/toolkit/components/antitracking/test/browser/empty.html";

function cleanupHSTS(aUseSite) {
  // Ensure to remove example.com from the HSTS list.
  let sss = Cc["@mozilla.org/ssservice;1"].getService(
    Ci.nsISiteSecurityService
  );

  for (let origin of ["example.com", "example.org"]) {
    let originAttributes = {};

    if (aUseSite) {
      originAttributes = { partitionKey: `(http,${origin})` };
    } else {
      originAttributes = { partitionKey: origin };
    }

    sss.resetState(NetUtil.newURI("http://example.com/"), originAttributes);
  }
}

function waitFor(host, type) {
  return new Promise(resolve => {
    const observer = channel => {
      if (
        channel instanceof Ci.nsIHttpChannel &&
        channel.URI.host === host &&
        channel.loadInfo.internalContentPolicyType === type
      ) {
        Services.obs.removeObserver(observer, "http-on-stop-request");
        resolve(channel.URI.spec);
      }
    };
    Services.obs.addObserver(observer, "http-on-stop-request");
  });
}

add_task(async function () {
  for (let partitionPerSite of [true, false]) {
    await SpecialPowers.pushPrefEnv({
      set: [
        ["privacy.dynamic_firstparty.use_site", partitionPerSite],
        ["security.mixed_content.upgrade_display_content", false],
      ],
    });

    let tab = (gBrowser.selectedTab = BrowserTestUtils.addTab(gBrowser));

    // Let's load the secureURL as first-party in order to activate HSTS.
    await BrowserTestUtils.loadURIString({
      browser: tab.linkedBrowser,
      uriString: secureURL,
      finalURI: secureURL,
    });

    // Let's test HSTS: unsecure -> secure.
    await BrowserTestUtils.loadURIString({
      browser: tab.linkedBrowser,
      uriString: unsecureURL,
      finalURI: secureURL,
    });
    ok(true, "unsecure -> secure, first-party works!");

    // Let's load a first-party.
    await BrowserTestUtils.loadURIString({
      browser: tab.linkedBrowser,
      uriString: unsecureEmptyURL,
      finalURI: unsecureEmptyURL,
    });

    let finalURL = waitFor(
      "example.com",
      Ci.nsIContentPolicy.TYPE_INTERNAL_IFRAME
    );

    await SpecialPowers.spawn(tab.linkedBrowser, [unsecureURL], async url => {
      let ifr = content.document.createElement("iframe");
      content.document.body.appendChild(ifr);
      ifr.src = url;
    });

    is(await finalURL, unsecureURL, "HSTS doesn't work for 3rd parties");

    gBrowser.removeCurrentTab();
    cleanupHSTS(partitionPerSite);
  }
});

add_task(async function test_subresource() {
  for (let partitionPerSite of [true, false]) {
    await SpecialPowers.pushPrefEnv({
      set: [
        ["privacy.dynamic_firstparty.use_site", partitionPerSite],
        ["security.mixed_content.upgrade_display_content", false],
      ],
    });

    let tab = (gBrowser.selectedTab = BrowserTestUtils.addTab(gBrowser));

    // Load a secure page as first party.
    await BrowserTestUtils.loadURIString({
      browser: tab.linkedBrowser,
      uriString: secureEmptyURL,
      finalURI: secureEmptyURL,
    });

    let loadPromise = waitFor(
      "example.com",
      Ci.nsIContentPolicy.TYPE_INTERNAL_IMAGE
    );

    // Load a secure subresource. HSTS won't be activated, since third
    // parties can't set HSTS.
    await SpecialPowers.spawn(tab.linkedBrowser, [secureImgURL], async url => {
      let ifr = content.document.createElement("img");
      content.document.body.appendChild(ifr);
      ifr.src = url;
    });

    // Ensure the subresource is loaded.
    await loadPromise;

    // Reload the secure page as first party.
    await BrowserTestUtils.loadURIString({
      browser: tab.linkedBrowser,
      uriString: secureEmptyURL,
      finalURI: secureEmptyURL,
    });

    let finalURL = waitFor(
      "example.com",
      Ci.nsIContentPolicy.TYPE_INTERNAL_IMAGE
    );

    // Load an unsecure subresource. It should not be upgraded to https.
    await SpecialPowers.spawn(
      tab.linkedBrowser,
      [unsecureImgURL],
      async url => {
        let ifr = content.document.createElement("img");
        content.document.body.appendChild(ifr);
        ifr.src = url;
      }
    );

    is(await finalURL, unsecureImgURL, "HSTS isn't set for 3rd parties");

    // Load the secure page with a different origin as first party.
    await BrowserTestUtils.loadURIString({
      browser: tab.linkedBrowser,
      uriString: secureAnotherEmptyURL,
      finalURI: secureAnotherEmptyURL,
    });

    finalURL = waitFor("example.com", Ci.nsIContentPolicy.TYPE_INTERNAL_IMAGE);

    // Load a unsecure subresource
    await SpecialPowers.spawn(
      tab.linkedBrowser,
      [unsecureImgURL],
      async url => {
        let ifr = content.document.createElement("img");
        content.document.body.appendChild(ifr);
        ifr.src = url;
      }
    );

    is(await finalURL, unsecureImgURL, "HSTS isn't set for 3rd parties");

    gBrowser.removeCurrentTab();
    cleanupHSTS(partitionPerSite);
  }
});

add_task(async function test_includeSubDomains() {
  for (let partitionPerSite of [true, false]) {
    await SpecialPowers.pushPrefEnv({
      set: [
        ["privacy.dynamic_firstparty.use_site", partitionPerSite],
        ["security.mixed_content.upgrade_display_content", false],
      ],
    });

    let tab = (gBrowser.selectedTab = BrowserTestUtils.addTab(gBrowser));

    // Load a secure page as first party to activate HSTS.
    await BrowserTestUtils.loadURIString({
      browser: tab.linkedBrowser,
      uriString: secureIncludeSubURL,
      finalURI: secureIncludeSubURL,
    });

    // Load a unsecure sub-domain page as first party to see if it's upgraded.
    await BrowserTestUtils.loadURIString({
      browser: tab.linkedBrowser,
      uriString: unsecureSubEmptyURL,
      finalURI: secureSubEmptyURL,
    });

    // Load a sub domain page which will trigger the cert error page.
    let certErrorLoaded = BrowserTestUtils.waitForErrorPage(tab.linkedBrowser);
    BrowserTestUtils.startLoadingURIString(
      tab.linkedBrowser,
      unsecureNoCertSubEmptyURL
    );
    await certErrorLoaded;

    // Verify the error page has the 'badStsCert' in its query string
    await SpecialPowers.spawn(tab.linkedBrowser, [], () => {
      let searchParams = new content.URLSearchParams(
        content.document.documentURI
      );

      is(
        searchParams.get("s"),
        "badStsCert",
        "The cert error page has 'badStsCert' set"
      );
    });

    gBrowser.removeCurrentTab();
    cleanupHSTS(partitionPerSite);
  }
});
