Bug 766210 - Delay updating nsPrefetchNode's reference to redirect channel until we know redirect is approved, r=jduell

uriloader/prefetch/nsPrefetchService.cpp

@@ -220,12 +220,13 @@ nsPrefetchNode::CancelChannel(nsresult error)
 // nsPrefetchNode::nsISupports
 //-----------------------------------------------------------------------------
 
-NS_IMPL_ISUPPORTS5(nsPrefetchNode,
+NS_IMPL_ISUPPORTS6(nsPrefetchNode,
                    nsIDOMLoadStatus,
                    nsIRequestObserver,
                    nsIStreamListener,
                    nsIInterfaceRequestor,
-                   nsIChannelEventSink)
+                   nsIChannelEventSink,
+                   nsIRedirectResultListener)
 
 //-----------------------------------------------------------------------------
 // nsPrefetchNode::nsIStreamListener
@@ -325,6 +326,12 @@ nsPrefetchNode::GetInterface(const nsIID &aIID, void **aResult)
         return NS_OK;
     }
 
+    if (aIID.Equals(NS_GET_IID(nsIRedirectResultListener))) {
+        NS_ADDREF_THIS();
+        *aResult = static_cast<nsIRedirectResultListener *>(this);
+        return NS_OK;
+    }
+
     return NS_ERROR_NO_INTERFACE;
 }
 
@@ -361,12 +368,28 @@ nsPrefetchNode::AsyncOnChannelRedirect(nsIChannel *aOldChannel,
                                   NS_LITERAL_CSTRING("prefetch"),
                                   false);
 
-    mChannel = aNewChannel;
+    // Assign to mChannel after we get notification about success of the
+    // redirect in OnRedirectResult.
+    mRedirectChannel = aNewChannel;
 
     callback->OnRedirectVerifyCallback(NS_OK);
     return NS_OK;
 }
 
+//-----------------------------------------------------------------------------
+// nsPrefetchNode::nsIRedirectResultListener
+//-----------------------------------------------------------------------------
+
+NS_IMETHODIMP
+nsPrefetchNode::OnRedirectResult(bool proceeding)
+{
+    if (proceeding && mRedirectChannel)
+        mChannel = mRedirectChannel;
+
+    mRedirectChannel = nsnull;
+
+    return NS_OK;
+}
 
 //-----------------------------------------------------------------------------
 // nsPrefetchService <public>

uriloader/prefetch/nsPrefetchService.h

@@ -9,6 +9,7 @@
 #include "nsIObserver.h"
 #include "nsIInterfaceRequestor.h"
 #include "nsIChannelEventSink.h"
+#include "nsIRedirectResultListener.h"
 #include "nsIWebProgressListener.h"
 #include "nsIStreamListener.h"
 #include "nsIChannel.h"
@@ -85,6 +86,7 @@ class nsPrefetchNode MOZ_FINAL : public nsIDOMLoadStatus
                                , public nsIStreamListener
                                , public nsIInterfaceRequestor
                                , public nsIChannelEventSink
+                               , public nsIRedirectResultListener
 {
 public:
     NS_DECL_ISUPPORTS
@@ -93,6 +95,7 @@ public:
     NS_DECL_NSISTREAMLISTENER
     NS_DECL_NSIINTERFACEREQUESTOR
     NS_DECL_NSICHANNELEVENTSINK
+    NS_DECL_NSIREDIRECTRESULTLISTENER
 
     nsPrefetchNode(nsPrefetchService *aPrefetchService,
                    nsIURI *aURI,
@@ -112,6 +115,7 @@ public:
 private:
     nsRefPtr<nsPrefetchService> mService;
     nsCOMPtr<nsIChannel>        mChannel;
+    nsCOMPtr<nsIChannel>        mRedirectChannel;
     PRUint16                    mState;
     PRInt32                     mBytesRead;
 };