[iOS] Tweaks to the session restore IPC callback

This CL makes sure that the message sent at the end of the restoration
is only added if a session restoration is in progress. It is also making
sure to remove the script once the restoration is complete.

Fixed: 1127521
Change-Id: Ie70be6a9233c540e6a4c6a05531eb1f5b1b1e206
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2463824Reviewed-by: Eugene But <eugenebut@chromium.org>
Reviewed-by: Justin Cohen <justincohen@chromium.org>
Commit-Queue: Gauthier Ambard <gambard@chromium.org>
Cr-Commit-Position: refs/heads/master@{#824079}

ios/web/web_state/ui/crw_web_controller.mm

@@ -447,14 +447,18 @@ typedef void (^ViewportStateCompletion)(const web::PageViewportState*);
                            name:kScriptMessageName
                         webView:_webView];
 
-    // TODO(crbug.com/1127521) Consider consolidating session restore script
-    // logic into a different place.
-    [messageRouter
-        setScriptMessageHandler:^(WKScriptMessage* message) {
-          [weakSelf didReceiveSessionRestoreScriptMessage:message];
-        }
-                           name:kSessionRestoreScriptMessageName
-                        webView:_webView];
+    if (self.webStateImpl->GetNavigationManager()
+            ->IsRestoreSessionInProgress()) {
+      // The session restoration script needs to use IPC to notify the app of
+      // the last step of the session restoration. See the restore_session.html
+      // file or crbug.com/1127521.
+      [messageRouter
+          setScriptMessageHandler:^(WKScriptMessage* message) {
+            [weakSelf didReceiveSessionRestoreScriptMessage:message];
+          }
+                             name:kSessionRestoreScriptMessageName
+                          webView:_webView];
+    }
 
     _webView.allowsBackForwardNavigationGestures =
         _allowsBackForwardNavigationGestures;
@@ -1121,8 +1125,6 @@ typedef void (^ViewportStateCompletion)(const web::PageViewportState*);
   }
 }
 
-// TODO(crbug.com/1127521) Consider consolidating session restore script
-// logic into a different place.
 - (void)didReceiveSessionRestoreScriptMessage:(WKScriptMessage*)message {
   if ([message.name isEqualToString:kSessionRestoreScriptMessageName] &&
       [message.body[@"offset"] isKindOfClass:[NSNumber class]]) {
@@ -1132,6 +1134,13 @@ typedef void (^ViewportStateCompletion)(const web::PageViewportState*);
     // Don't use |_jsInjector| -executeJavaScript here, as it relies on
     // |windowID| being injected before window.onload starts.
     web::ExecuteJavaScript(self.webView, method, nil);
+
+    // Removes the script as it is no longer needed.
+    CRWWKScriptMessageRouter* messageRouter =
+        [self webViewConfigurationProvider].GetScriptMessageRouter();
+    [messageRouter
+        removeScriptMessageHandlerForName:kSessionRestoreScriptMessageName
+                                  webView:_webView];
   } else {
     DLOG(WARNING) << "Invalid session restore JS message name.";
   }