Reland "Fix a crash in WebUIMessageHandler::AllowJavascript"

This is a reland of a7c54a0f

Original change's description:
> Fix a crash in WebUIMessageHandler::AllowJavascript
>
> If WebUI's JS is sending the message when we navigate cross-site,
> then we use the new process ID from WebContents in
> WebUIImpl::CanCallJavascript and a CHECK crash will happen.
>
> This CL changes to use the old process ID from RFH in
> WebUIImpl::CanCallJavascript.
>
> Bug: 1110276
> Change-Id: Ie827060b637dbd3aa4c7ad771c0cb4862763ced2
> Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2326050
> Reviewed-by: Nasko Oskov <nasko@chromium.org>
> Commit-Queue: Miyoung Shin <myid.shin@igalia.com>
> Cr-Commit-Position: refs/heads/master@{#793816}

Bug: 1110276, 1112189, 1112238
Change-Id: I02916364f8cb3871ab4cdf17cd2827a5f1cfbb60
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2335213Reviewed-by: Nasko Oskov <nasko@chromium.org>
Reviewed-by: Tommy Li <tommycli@chromium.org>
Commit-Queue: Miyoung Shin <myid.shin@igalia.com>
Cr-Commit-Position: refs/heads/master@{#796293}

chrome/browser/ui/webui/user_actions/user_actions_ui_handler.cc

@@ -16,13 +16,23 @@ UserActionsUIHandler::UserActionsUIHandler()
 }
 
 UserActionsUIHandler::~UserActionsUIHandler() {
+  WebContentsObserver::Observe(nullptr);
   base::RemoveActionCallback(action_callback_);
 }
 
-void UserActionsUIHandler::RegisterMessages() {}
+void UserActionsUIHandler::RegisterMessages() {
+  WebContentsObserver::Observe(web_ui()->GetWebContents());
+}
+
+void UserActionsUIHandler::ReadyToCommitNavigation(
+    content::NavigationHandle* navigation_handle) {
+  AllowJavascript();
+}
 
 void UserActionsUIHandler::OnUserAction(const std::string& action,
                                         base::TimeTicks action_time) {
+  if (!IsJavascriptAllowed())
+    return;
   base::Value user_action_name(action);
   web_ui()->CallJavascriptFunctionUnsafe("userActions.observeUserAction",
                                          user_action_name);

chrome/browser/ui/webui/user_actions/user_actions_ui_handler.h

@@ -7,6 +7,7 @@
 
 #include "base/macros.h"
 #include "base/metrics/user_metrics.h"
+#include "content/public/browser/web_contents_observer.h"
 #include "content/public/browser/web_ui_message_handler.h"
 
 namespace base {
@@ -16,7 +17,8 @@ class TimeTicks;
 // UI Handler for chrome://user-actions/
 // It listens to user action notifications and passes those notifications
 // into the Javascript to update the page.
-class UserActionsUIHandler : public content::WebUIMessageHandler {
+class UserActionsUIHandler : public content::WebUIMessageHandler,
+                             public content::WebContentsObserver {
  public:
   UserActionsUIHandler();
   ~UserActionsUIHandler() override;
@@ -25,6 +27,10 @@ class UserActionsUIHandler : public content::WebUIMessageHandler {
   // Does nothing for now.
   void RegisterMessages() override;
 
+  // WebContentsObserver::
+  void ReadyToCommitNavigation(
+      content::NavigationHandle* navigation_handle) override;
+
  private:
   void OnUserAction(const std::string& action, base::TimeTicks action_time);
 

content/browser/webui/web_ui_browsertest.cc

@@ -14,9 +14,11 @@
 #include "base/run_loop.h"
 #include "base/strings/strcat.h"
 #include "base/strings/utf_string_conversions.h"
+#include "base/test/bind_test_util.h"
 #include "base/test/simple_test_tick_clock.h"
 #include "base/time/time.h"
 #include "content/browser/webui/content_web_ui_controller_factory.h"
+#include "content/browser/webui/web_ui_impl.h"
 #include "content/public/browser/child_process_security_policy.h"
 #include "content/public/browser/render_frame_host.h"
 #include "content/public/browser/render_process_host.h"
@@ -52,6 +54,10 @@ class TestWebUIMessageHandler : public WebUIMessageHandler {
         "notifyFinish",
         base::BindRepeating(&TestWebUIMessageHandler::OnNotifyFinish,
                             base::Unretained(this)));
+    web_ui()->RegisterMessageCallback(
+        "sendMessage",
+        base::BindRepeating(&TestWebUIMessageHandler::OnSendMessase,
+                            base::Unretained(this)));
   }
 
   void set_finish_closure(base::RepeatingClosure closure) {
@@ -62,6 +68,10 @@ class TestWebUIMessageHandler : public WebUIMessageHandler {
     return message_requiring_gesture_count_;
   }
 
+  void set_send_message_closure(base::OnceClosure closure) {
+    send_message_closure_ = std::move(closure);
+  }
+
  private:
   void OnMessageRequiringGesture(const base::ListValue* args) {
     ++message_requiring_gesture_count_;
@@ -72,8 +82,27 @@ class TestWebUIMessageHandler : public WebUIMessageHandler {
       finish_closure_.Run();
   }
 
+  void OnSendMessase(const base::ListValue* args) {
+    // This message will be invoked when WebContents changes the main RFH
+    // and the old main RFH is still alive during navigating from WebUI page
+    // to cross-site. WebUI message should be handled with old main RFH.
+
+    if (send_message_closure_)
+      std::move(send_message_closure_).Run();
+
+    // AllowJavascript should not have a CHECK crash.
+    AllowJavascript();
+
+    // WebUI::CallJavascriptFunctionUnsafe should be run with old main RFH.
+    web_ui()->CallJavascriptFunctionUnsafe("test");
+
+    if (finish_closure_)
+      std::move(finish_closure_).Run();
+  }
+
   int message_requiring_gesture_count_ = 0;
   base::RepeatingClosure finish_closure_;
+  base::OnceClosure send_message_closure_;
 };
 
 class WebUIRequiringGestureBrowserTest : public ContentBrowserTest {
@@ -338,6 +367,44 @@ IN_PROC_BROWSER_TEST_F(WebUIImplBrowserTest, UntrustedSchemeLoads) {
             web_contents->GetTitle());
 }
 
+// Verify that we can successfully navigate to a chrome-untrusted:// URL
+// without a crash while WebUI::Send is being performed.
+IN_PROC_BROWSER_TEST_F(WebUIImplBrowserTest, NavigateWhileWebUISend) {
+  ASSERT_TRUE(embedded_test_server()->Start());
+
+  auto* web_contents = shell()->web_contents();
+  ASSERT_TRUE(NavigateToURL(web_contents, GetWebUIURL(kChromeUIGpuHost)));
+
+  auto* test_handler = new TestWebUIMessageHandler;
+  web_contents->GetWebUI()->AddMessageHandler(base::WrapUnique(test_handler));
+
+  auto* webui = static_cast<WebUIImpl*>(web_contents->GetWebUI());
+  EXPECT_EQ(web_contents->GetMainFrame(), webui->frame_host_for_test());
+
+  test_handler->set_finish_closure(base::BindLambdaForTesting([&]() {
+    EXPECT_NE(web_contents->GetMainFrame(), webui->frame_host_for_test());
+  }));
+
+  bool received_send_message = false;
+  test_handler->set_send_message_closure(
+      base::BindLambdaForTesting([&]() { received_send_message = true; }));
+
+  base::RunLoop run_loop;
+  web_contents->GetMainFrame()->ExecuteJavaScriptForTests(
+      base::ASCIIToUTF16("onunload=function() { chrome.send('sendMessage')}"),
+      base::BindOnce([](base::OnceClosure callback,
+                        base::Value) { std::move(callback).Run(); },
+                     run_loop.QuitClosure()));
+  run_loop.Run();
+
+  RenderFrameDeletedObserver delete_observer(web_contents->GetMainFrame());
+  EXPECT_TRUE(NavigateToURL(
+      web_contents, embedded_test_server()->GetURL("/simple_page.html")));
+  delete_observer.WaitUntilDeleted();
+
+  EXPECT_TRUE(received_send_message);
+}
+
 class WebUIRequestSchemesTest : public ContentBrowserTest {
  public:
   WebUIRequestSchemesTest() {

content/browser/webui/web_ui_impl.cc

@@ -203,13 +203,11 @@ void WebUIImpl::SetController(std::unique_ptr<WebUIController> controller) {
 }
 
 bool WebUIImpl::CanCallJavascript() {
-  RenderFrameHost* frame_host = web_contents_->GetMainFrame();
-  return frame_host &&
-         (ChildProcessSecurityPolicyImpl::GetInstance()->HasWebUIBindings(
-              frame_host->GetProcess()->GetID()) ||
+  return (ChildProcessSecurityPolicyImpl::GetInstance()->HasWebUIBindings(
+              frame_host_->GetProcess()->GetID()) ||
           // It's possible to load about:blank in a Web UI renderer.
           // See http://crbug.com/42547
-          frame_host->GetLastCommittedURL().spec() == url::kAboutBlankURL);
+          frame_host_->GetLastCommittedURL().spec() == url::kAboutBlankURL);
 }
 
 void WebUIImpl::CallJavascriptFunctionUnsafe(const std::string& function_name) {
@@ -311,8 +309,7 @@ void WebUIImpl::ExecuteJavascript(const base::string16& javascript) {
   if (!CanCallJavascript())
     return;
 
-  web_contents_->GetMainFrame()->ExecuteJavaScript(javascript,
-                                                   base::NullCallback());
+  frame_host_->ExecuteJavaScript(javascript, base::NullCallback());
 }
 
 void WebUIImpl::DisallowJavascriptOnAllHandlers() {

content/browser/webui/web_ui_impl.h

@@ -93,6 +93,8 @@ class CONTENT_EXPORT WebUIImpl : public WebUI,
 
   const mojo::Remote<mojom::WebUI>& GetRemoteForTest() const { return remote_; }
 
+  RenderFrameHost* frame_host_for_test() const { return frame_host_; }
+
  private:
   class MainFrameNavigationObserver;