From d9c28dead6e8508684e385b796bdba5ae062bc9e Mon Sep 17 00:00:00 2001
From: Mike Conley <mconley@mozilla.com>
Date: Wed, 20 Jul 2016 10:03:52 -0400
Subject: [PATCH] Bug 1261842 - Make sure Marionette can handle the redirect
 and error page cases. r=automatedtester

With the initial browser defaulting to remote, there's a greater likelihood
that the DOMContentLoaded event that is handled in the "get" function will
be fired by the initial about:blank instead of the actual desired page.

get() currently works around this by ensuring that the URL of the loaded
page matches the requested one when DOMContentLoaded fires. Unfortunately,
this doesn't work for pages that redirect via their HTTP headers (and will
therefore not fire DOMContentLoaded).

This patch fixes things by adding an nsIWebProgressListener that ensures
that the requested page has started to load before paying any attention
to the DOMContentLoaded events. This handles the redirect case. We also
compare against nsIChannel.originalURI for the about: redirect case.

For neterror pages (which never open channels, and therefore are not
seen by the nsIWebProgressListener), we just check that the page that
we attempted to reach was the one that was requested.

MozReview-Commit-ID: Gbbmfwat46s

--HG--
extra : rebase_source : 1848cd67757be8780f9e50253dc0ee1131467257
---
 testing/marionette/listener.js | 69 +++++++++++++++++++++++++++++++++-
 1 file changed, 67 insertions(+), 2 deletions(-)

diff --git a/testing/marionette/listener.js b/testing/marionette/listener.js
index 62c3648776ed..6e8372ffb8ce 100644
--- a/testing/marionette/listener.js
+++ b/testing/marionette/listener.js
@@ -930,6 +930,57 @@ function pollForReadyState(msg, start, callback) {
 function get(msg) {
   let start = new Date().getTime();
   let requestedURL = new URL(msg.json.url).toString();
+  let docShell = curContainer.frame
+                             .document
+                             .defaultView
+                             .QueryInterface(Ci.nsIInterfaceRequestor)
+                             .getInterface(Ci.nsIWebNavigation)
+                             .QueryInterface(Ci.nsIDocShell);
+  let webProgress = docShell.QueryInterface(Ci.nsIInterfaceRequestor)
+                            .getInterface(Ci.nsIWebProgress);
+  let sawLoad = false;
+
+  // It's possible that a site we're being sent to will end up redirecting
+  // us before we end up on a page that fires DOMContentLoaded. We can ensure
+  // This loadListener ensures that we don't send a success signal back to
+  // the caller until we've seen the load of the requested URL attempted
+  // on this frame.
+  let loadListener = {
+    QueryInterface: XPCOMUtils.generateQI([Ci.nsIWebProgressListener,
+                                           Ci.nsISupportsWeakReference]),
+    onStateChange(webProgress, request, state, status) {
+      if (!(request instanceof Ci.nsIChannel)) {
+        return;
+      }
+
+      let isDocument = state & Ci.nsIWebProgressListener.STATE_IS_DOCUMENT;
+      let isStart = state & Ci.nsIWebProgressListener.STATE_START;
+      let loadedURL = request.URI.spec;
+      // We have to look at the originalURL because for about: pages,
+      // the loadedURL is what the about: page resolves to, and is
+      // not the one that was requested.
+      let originalURL = request.originalURI.spec;
+      let isRequestedURL = loadedURL == requestedURL ||
+                           originalURL == requestedURL;
+
+      if (isDocument && isStart && isRequestedURL) {
+        // We started loading the requested document. This document
+        // might not be the one that ends up firing DOMContentLoaded
+        // (if it, for example, redirects), but because we've started
+        // loading this URL, we know that any future DOMContentLoaded's
+        // are fair game to tell the Marionette client about.
+        sawLoad = true;
+      }
+    },
+
+    onLocationChange() {},
+    onProgressChange() {},
+    onStatusChange() {},
+    onSecurityChange() {},
+  };
+
+  webProgress.addProgressListener(loadListener,
+                                  Ci.nsIWebProgress.NOTIFY_STATE_DOCUMENT);
 
   // Prevent DOMContentLoaded events from frames from invoking this
   // code, unless the event is coming from the frame associated with
@@ -939,9 +990,22 @@ function get(msg) {
       !event.originalTarget.defaultView.frameElement ||
       event.originalTarget.defaultView.frameElement == curContainer.frame.frameElement;
 
-    let correctURL = curContainer.frame.location == requestedURL;
+    // If the page we're at fired DOMContentLoaded and appears
+    // to be the one we asked to load, then we definitely
+    // saw the load occur. We need this because for error
+    // pages, like about:neterror for unsupported protocols,
+    // we don't end up opening a channel that our
+    // WebProgressListener can monitor.
+    if (curContainer.frame.location == requestedURL) {
+      sawLoad = true;
+    }
 
-    if (correctFrame && correctURL) {
+    // We also need to make sure that the DOMContentLoaded we saw isn't
+    // for the initial about:blank of a newly created docShell.
+    let loadedNonAboutBlank = docShell.hasLoadedNonBlankURI;
+
+    if (correctFrame && sawLoad && loadedNonAboutBlank) {
+      webProgress.removeProgressListener(loadListener);
       pollForReadyState(msg, start, () => {
         removeEventListener("DOMContentLoaded", onDOMContentLoaded, false);
       });
@@ -950,6 +1014,7 @@ function get(msg) {
 
   function timerFunc() {
     removeEventListener("DOMContentLoaded", onDOMContentLoaded, false);
+    webProgress.removeProgressListener(loadListener);
     sendError(new TimeoutError("Error loading page, timed out (onDOMContentLoaded)"), msg.json.command_id);
   }
   if (msg.json.pageTimeout != null) {