Pass the webstate to be replaced to the pre-render service

There is only 1 pre-render service per BrowserState which makes
the last created BVC the delegate for the the pre-render service
so when there are multiple windows, the last one created will provide
the webState to be replaced to the service and the service will use that
to get the navigation history which will be wrong in cases where the
prerendered webstate is not on the last created window.

As a work around pass the webstate to be replaced to the service

Bug: 1139286
Change-Id: I5a8651ba1c52ecf9e0a9d2ea0a3a18c1e0afdfc5
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2486961Reviewed-by: Mark Cogan <marq@chromium.org>
Commit-Queue: Mohammad Refaat <mrefaat@chromium.org>
Cr-Commit-Position: refs/heads/master@{#820297}

ios/chrome/browser/prerender/fake_prerender_service.h

@@ -27,6 +27,7 @@ class FakePrerenderService : public PrerenderService {
   void StartPrerender(const GURL& url,
                       const web::Referrer& referrer,
                       ui::PageTransition transition,
+                      web::WebState* web_state_to_replace,
                       bool immediately) override;
   bool MaybeLoadPrerenderedURL(const GURL& url,
                                ui::PageTransition transition,

ios/chrome/browser/prerender/fake_prerender_service.mm

@@ -18,6 +18,7 @@ void FakePrerenderService::SetDelegate(id<PreloadControllerDelegate> delegate) {
 void FakePrerenderService::StartPrerender(const GURL& url,
                                           const web::Referrer& referrer,
                                           ui::PageTransition transition,
+                                          web::WebState* web_state_to_replace,
                                           bool immediately) {
   preload_url_ = url;
 }

ios/chrome/browser/prerender/preload_controller.h

@@ -48,9 +48,10 @@ class WebState;
 // Prerenders the given |url| with the given |transition|.  Normally, prerender
 // requests are fulfilled after a short delay, to prevent unnecessary prerenders
 // while the user is typing.  If |immediately| is YES, this method starts
-// prerendering immediately, with no delay.  |immediately| should be set to YES
-// only when there is a very high confidence that the user will navigate to the
-// given |url|.
+// prerendering immediately, with no delay. |currentWebState| is used to create
+// a new WebState for the prerender with the same session. |immediately| should
+// be set to YES only when there is a very high confidence that the user will
+// navigate to the given |url|.
 //
 // If there is already an existing request for |url|, this method does nothing
 // and does not reset the delay timer.  If there is an existing request for a
@@ -59,6 +60,7 @@ class WebState;
 - (void)prerenderURL:(const GURL&)url
             referrer:(const web::Referrer&)referrer
           transition:(ui::PageTransition)transition
+     currentWebState:(web::WebState*)currentWebState
          immediately:(BOOL)immediately;
 
 // Cancels any outstanding prerender requests and destroys any prerendered Tabs.

ios/chrome/browser/prerender/preload_controller.mm

@@ -159,6 +159,7 @@ const base::Feature kPreloadDelayWebStateReset{
                                  PreloadCancelling> {
   std::unique_ptr<web::WebStateDelegateBridge> _webStateDelegate;
   std::unique_ptr<web::WebStateObserverBridge> _webStateObserver;
+  std::unique_ptr<web::WebStateObserverBridge> _webStateToReplaceObserver;
   std::unique_ptr<PrefObserverBridge> _observerBridge;
   std::unique_ptr<ConnectionTypeObserverBridge> _connectionTypeObserver;
   std::unique_ptr<web::WebStatePolicyDeciderBridge> _policyDeciderBridge;
@@ -174,6 +175,11 @@ const base::Feature kPreloadDelayWebStateReset{
 
   // The dialog presenter.
   std::unique_ptr<web::JavaScriptDialogPresenter> _dialogPresenter;
+
+  // A weak pointer to the webState that will be replaced with the prerendered
+  // one. This is needed by |startPrerender| to build the new webstate with the
+  // same sessions.
+  web::WebState* _webStateToReplace;
 }
 
 // The ChromeBrowserState passed on initialization.
@@ -245,6 +251,8 @@ const base::Feature kPreloadDelayWebStateReset{
         net::NetworkChangeNotifier::GetConnectionType());
     _webStateDelegate = std::make_unique<web::WebStateDelegateBridge>(self);
     _webStateObserver = std::make_unique<web::WebStateObserverBridge>(self);
+    _webStateToReplaceObserver =
+        std::make_unique<web::WebStateObserverBridge>(self);
     _observerBridge = std::make_unique<PrefObserverBridge>(self);
     _prefChangeRegistrar.Init(_browserState->GetPrefs());
     _observerBridge->ObserveChangesForPreference(
@@ -255,7 +263,7 @@ const base::Feature kPreloadDelayWebStateReset{
       _connectionTypeObserver =
           std::make_unique<ConnectionTypeObserverBridge>(self);
     }
-
+    _webStateToReplace = nullptr;
     [[NSNotificationCenter defaultCenter]
         addObserver:self
            selector:@selector(didReceiveMemoryWarning)
@@ -325,6 +333,7 @@ const base::Feature kPreloadDelayWebStateReset{
 - (void)prerenderURL:(const GURL&)url
             referrer:(const web::Referrer&)referrer
           transition:(ui::PageTransition)transition
+     currentWebState:(web::WebState*)currentWebState
          immediately:(BOOL)immediately {
   // TODO(crbug.com/754050): If CanPrerenderURL() returns false, should we
   // cancel any scheduled prerender requests?
@@ -340,6 +349,12 @@ const base::Feature kPreloadDelayWebStateReset{
   }
 
   [self removeScheduledPrerenderRequests];
+  _webStateToReplace = currentWebState;
+  // Observing the |_webStateToReplace| to make sure that if it's destructed
+  // the pre-rendering will be canceled.
+  if (_webStateToReplace) {
+    _webStateToReplace->AddObserver(_webStateToReplaceObserver.get());
+  }
   _scheduledRequest =
       std::make_unique<PrerenderRequest>(url, transition, referrer);
 
@@ -452,6 +467,9 @@ const base::Feature kPreloadDelayWebStateReset{
 
 - (void)webState:(web::WebState*)webState
     didFinishNavigation:(web::NavigationContext*)navigation {
+  // the |_webStateToReplace| is observed for destruction event only.
+  if (_webStateToReplace == webState)
+    return;
   DCHECK_EQ(webState, _webState.get());
   if ([self shouldCancelPreloadForMimeType:webState->GetContentsMimeType()])
     [self schedulePrerenderCancel];
@@ -459,6 +477,10 @@ const base::Feature kPreloadDelayWebStateReset{
 
 - (void)webState:(web::WebState*)webState
     didLoadPageWithSuccess:(BOOL)loadSuccess {
+  // the |_webStateToReplace| is observed for destruction event only.
+  if (_webStateToReplace == webState)
+    return;
+
   DCHECK_EQ(webState, _webState.get());
   // The load should have been cancelled when the navigation finishes, but this
   // makes sure that we didn't miss one.
@@ -469,6 +491,14 @@ const base::Feature kPreloadDelayWebStateReset{
   }
 }
 
+- (void)webStateDestroyed:(web::WebState*)webState {
+  if (_webState.get() == webState)
+    return;
+  DCHECK_EQ(webState, _webStateToReplace);
+  // There is no way to create a pre-rendered webState without existing webState
+  // web state to replace, So cancel the prerender.
+  [self schedulePrerenderCancel];
+}
 #pragma mark - CRWWebStatePolicyDecider
 
 - (WebStatePolicyDecider::PolicyDecision)
@@ -568,6 +598,10 @@ const base::Feature kPreloadDelayWebStateReset{
 - (void)removeScheduledPrerenderRequests {
   [NSObject cancelPreviousPerformRequestsWithTarget:self];
   _scheduledRequest = nullptr;
+  if (_webStateToReplace) {
+    _webStateToReplace->RemoveObserver(_webStateToReplaceObserver.get());
+  }
+  _webStateToReplace = nullptr;
 }
 
 #pragma mark - Prerender Helpers
@@ -577,17 +611,27 @@ const base::Feature kPreloadDelayWebStateReset{
   [self destroyPreviewContents];
   self.prerenderedURL = self.scheduledURL;
   std::unique_ptr<PrerenderRequest> request = std::move(_scheduledRequest);
+  // No need to observer the destruction of the |_webStateToReplace| anymore
+  // as it will be used here.
+  if (_webStateToReplace) {
+    _webStateToReplace->RemoveObserver(_webStateToReplaceObserver.get());
+  }
+
+  // TODO(crbug.com/1140583): The correct way is to always get the
+  // webStateToReplace from the delegate. however this is not possible because
+  // there is only one delegate per browser state.
+  if (!_webStateToReplace)
+    _webStateToReplace = [self.delegate webStateToReplace];
 
-  web::WebState* webStateToReplace = [self.delegate webStateToReplace];
-  if (!self.prerenderedURL.is_valid() || !webStateToReplace) {
+  if (!self.prerenderedURL.is_valid() || !_webStateToReplace) {
     [self destroyPreviewContents];
     return;
   }
 
   web::WebState::CreateParams createParams(self.browserState);
   _webState = web::WebState::CreateWithStorageSession(
-      createParams, webStateToReplace->BuildSessionStorage());
-
+      createParams, _webStateToReplace->BuildSessionStorage());
+  _webStateToReplace = nullptr;
   // Add the preload controller as a policyDecider before other tab helpers, so
   // that it can block the navigation if needed before other policy deciders
   // execute thier side effects (eg. AppLauncherTabHelper launching app).

ios/chrome/browser/prerender/preload_controller_unittest.mm

@@ -118,6 +118,7 @@ TEST_F(PreloadControllerTest, DontPreloadNonWebURLs) {
   [controller_ prerenderURL:GURL()
                    referrer:kReferrer
                  transition:kTransition
+            currentWebState:nil
                 immediately:YES];
   EXPECT_FALSE([controller_ releasePrerenderContents]);
 
@@ -126,6 +127,7 @@ TEST_F(PreloadControllerTest, DontPreloadNonWebURLs) {
   [controller_ prerenderURL:GURL("chrome://newtab")
                    referrer:kReferrer
                  transition:kTransition
+            currentWebState:nil
                 immediately:YES];
   EXPECT_FALSE([controller_ releasePrerenderContents]);
 
@@ -134,6 +136,7 @@ TEST_F(PreloadControllerTest, DontPreloadNonWebURLs) {
   [controller_ prerenderURL:GURL("about:flags")
                    referrer:kReferrer
                  transition:kTransition
+            currentWebState:nil
                 immediately:YES];
   EXPECT_FALSE([controller_ releasePrerenderContents]);
 }

ios/chrome/browser/prerender/prerender_service.h

@@ -29,9 +29,14 @@ class PrerenderService : public KeyedService {
   // Prerenders the given |url| with the given |transition|.  Normally,
   // prerender requests are fulfilled after a short delay, to prevent
   // unnecessary prerenders while the user is typing.  If |immediately| is YES,
-  // this method starts prerendering immediately, with no delay.  |immediately|
-  // should be set to YES only when there is a very high confidence that the
-  // user will navigate to the given |url|.
+  // this method starts prerendering immediately, with no delay.
+  // |web_state_to_replace| is provided so that the new prerendered web state
+  // can have the same session data.  |immediately| should be set to YES only
+  // when there is a very high confidence that the user will navigate to the
+  // given |url|.
+  // TODO(crbug.com/1140583): passing |web_state_to_replace| is a workaround for
+  // not having prerender service per browser, remove it once prerenderService
+  // is a browser agent.
   //
   // If there is already an existing request for |url|, this method does nothing
   // and does not reset the delay timer.  If there is an existing request for a
@@ -40,6 +45,7 @@ class PrerenderService : public KeyedService {
   virtual void StartPrerender(const GURL& url,
                               const web::Referrer& referrer,
                               ui::PageTransition transition,
+                              web::WebState* web_state_to_replace,
                               bool immediately) = 0;
 
   // If |url| is prerendered, loads the prerendered web state into

ios/chrome/browser/prerender/prerender_service_impl.h

@@ -26,6 +26,7 @@ class PrerenderServiceImpl : public PrerenderService {
   void StartPrerender(const GURL& url,
                       const web::Referrer& referrer,
                       ui::PageTransition transition,
+                      web::WebState* web_state_to_replace,
                       bool immediately) override;
   bool MaybeLoadPrerenderedURL(const GURL& url,
                                ui::PageTransition transition,

ios/chrome/browser/prerender/prerender_service_impl.mm

@@ -38,10 +38,12 @@ void PrerenderServiceImpl::SetDelegate(id<PreloadControllerDelegate> delegate) {
 void PrerenderServiceImpl::StartPrerender(const GURL& url,
                                           const web::Referrer& referrer,
                                           ui::PageTransition transition,
+                                          web::WebState* web_state_to_replace,
                                           bool immediately) {
   [controller_ prerenderURL:url
                    referrer:referrer
                  transition:transition
+            currentWebState:web_state_to_replace
                 immediately:immediately];
 }
 

ios/chrome/browser/safe_browsing/safe_browsing_tab_helper_unittest.mm

@@ -799,7 +799,7 @@ TEST_P(SafeBrowsingTabHelperTest, UnsafeMainFrameRequestCancelsPrerendering) {
       PrerenderServiceFactory::GetForBrowserState(
           ChromeBrowserState::FromBrowserState(web_state_.GetBrowserState()));
   prerender_service->StartPrerender(unsafe_url, web::Referrer(),
-                                    ui::PAGE_TRANSITION_LINK,
+                                    ui::PAGE_TRANSITION_LINK, &web_state_,
                                     /*immediately=*/true);
 
   EXPECT_TRUE(ShouldAllowRequestUrl(unsafe_url).ShouldAllowNavigation());
@@ -830,7 +830,7 @@ TEST_P(SafeBrowsingTabHelperTest, UnsafeSubframeRequestCancelsPrerendering) {
       PrerenderServiceFactory::GetForBrowserState(
           ChromeBrowserState::FromBrowserState(web_state_.GetBrowserState()));
   prerender_service->StartPrerender(main_frame_item->GetURL(), web::Referrer(),
-                                    ui::PAGE_TRANSITION_LINK,
+                                    ui::PAGE_TRANSITION_LINK, &web_state_,
                                     /*immediately=*/true);
 
   EXPECT_TRUE(ShouldAllowRequestUrl(unsafe_url, /*for_main_frame=*/false)
@@ -864,7 +864,7 @@ TEST_P(SafeBrowsingTabHelperTest,
       PrerenderServiceFactory::GetForBrowserState(
           ChromeBrowserState::FromBrowserState(web_state_.GetBrowserState()));
   prerender_service->StartPrerender(safe_url, web::Referrer(),
-                                    ui::PAGE_TRANSITION_LINK,
+                                    ui::PAGE_TRANSITION_LINK, &web_state_,
                                     /*immediately=*/true);
 
   EXPECT_TRUE(ShouldAllowRequestUrl(safe_url).ShouldAllowNavigation());
@@ -891,7 +891,7 @@ TEST_P(SafeBrowsingTabHelperTest,
       PrerenderServiceFactory::GetForBrowserState(
           ChromeBrowserState::FromBrowserState(web_state_.GetBrowserState()));
   prerender_service->StartPrerender(main_frame_item->GetURL(), web::Referrer(),
-                                    ui::PAGE_TRANSITION_LINK,
+                                    ui::PAGE_TRANSITION_LINK, &web_state_,
                                     /*immediately=*/true);
 
   EXPECT_TRUE(ShouldAllowRequestUrl(safe_url, /*for_main_frame=*/false)

ios/chrome/browser/ui/omnibox/chrome_omnibox_client_ios.mm

@@ -159,7 +159,7 @@ void ChromeOmniboxClientIOS::OnResultChanged(
     ui::PageTransition transition = ui::PageTransitionFromInt(
         match.transition | ui::PAGE_TRANSITION_FROM_ADDRESS_BAR);
     service->StartPrerender(match.destination_url, web::Referrer(), transition,
-                            is_inline_autocomplete);
+                            controller_->GetWebState(), is_inline_autocomplete);
   } else {
     service->CancelPrerender();
   }