[headless] Added async callback option to Simple CDP protocol.

DevTools backend assumes that response and event handling is
asynchronous, and not respecting this assumption causes use after
free problems, see the associated bug.

This CL introduces async response/event callback option for
the Simple CDP client.

Drive by: Provide cleaner headless_shell shutdown routine.

(cherry picked from commit 7493091)

Bug: 1382993
Change-Id: I14b37079f1b7488777ebcf20a808b3bfb9eb44c5
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4021401
Commit-Queue: Peter Kvitek <kvitekp@chromium.org>
Reviewed-by: Andrey Kosyakov <caseq@chromium.org>
Cr-Original-Commit-Position: refs/heads/main@{#1070102}
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4027032
Reviewed-by: Dmitry Gozman <dgozman@chromium.org>
Cr-Commit-Position: refs/branch-heads/5414@{#47}
Cr-Branched-From: 4417ee5-refs/heads/main@{#1070088}

components/devtools/simple_devtools_protocol_client/simple_devtools_protocol_client.cc

@@ -7,11 +7,14 @@
 #include <algorithm>
 #include <vector>
 
+#include "base/bind.h"
 #include "base/json/json_reader.h"
 #include "base/json/json_writer.h"
 #include "base/memory/ptr_util.h"
 #include "base/strings/strcat.h"
 #include "base/strings/string_split.h"
+#include "content/public/browser/browser_task_traits.h"
+#include "content/public/browser/browser_thread.h"
 #include "content/public/browser/render_process_host.h"
 #include "content/public/browser/web_contents.h"
 #include "third_party/abseil-cpp/absl/types/optional.h"
@@ -41,6 +44,10 @@ SimpleDevToolsProtocolClient::~SimpleDevToolsProtocolClient() {
     parent_client_->sessions_.erase(session_id_);
 }
 
+void SimpleDevToolsProtocolClient::InitBrowserMainThread() {
+  browser_main_thread_ = content::GetUIThreadTaskRunner({});
+}
+
 void SimpleDevToolsProtocolClient::AttachClient(
     scoped_refptr<content::DevToolsAgentHost> agent_host) {
   DCHECK(!agent_host_);
@@ -120,7 +127,12 @@ void SimpleDevToolsProtocolClient::DispatchProtocolMessage(
     ResponseCallback callback(std::move(it->second));
     pending_response_map_.erase(it);
 
-    std::move(callback).Run(std::move(message));
+    if (browser_main_thread_) {
+      browser_main_thread_->PostTask(
+          FROM_HERE, base::BindOnce(std::move(callback), std::move(message)));
+    } else {
+      std::move(callback).Run(std::move(message));
+    }
     return;
   }
 
@@ -141,8 +153,13 @@ void SimpleDevToolsProtocolClient::DispatchProtocolMessage(
   bool first_callback = true;
   for (auto& callback : handlers) {
     if (first_callback || HasEventHandler(*event_name, callback)) {
-      callback.Run(message);
       first_callback = false;
+      if (browser_main_thread_) {
+        browser_main_thread_->PostTask(
+            FROM_HERE, base::BindOnce(std::move(callback), std::move(message)));
+      } else {
+        callback.Run(std::move(message));
+      }
     }
   }
 }

components/devtools/simple_devtools_protocol_client/simple_devtools_protocol_client.h

@@ -13,6 +13,7 @@
 #include "base/containers/span.h"
 #include "base/memory/raw_ptr.h"
 #include "base/memory/scoped_refptr.h"
+#include "base/task/sequenced_task_runner.h"
 #include "base/values.h"
 #include "content/public/browser/devtools_agent_host.h"
 
@@ -32,6 +33,8 @@ class SimpleDevToolsProtocolClient : public content::DevToolsAgentHostClient {
 
   ~SimpleDevToolsProtocolClient() override;
 
+  void InitBrowserMainThread();
+
   void AttachClient(scoped_refptr<content::DevToolsAgentHost> agent_host);
   void DetachClient();
 
@@ -71,6 +74,7 @@ class SimpleDevToolsProtocolClient : public content::DevToolsAgentHostClient {
                        const EventCallback& event_callback);
 
   const std::string session_id_;
+  scoped_refptr<base::SequencedTaskRunner> browser_main_thread_;
   base::raw_ptr<SimpleDevToolsProtocolClient> parent_client_ = nullptr;
   base::flat_map<std::string, SimpleDevToolsProtocolClient*> sessions_;
 

headless/app/headless_shell.cc

@@ -242,6 +242,8 @@ void HeadlessShell::OnBrowserStart(HeadlessBrowser* browser) {
   }
 #endif
 
+  devtools_client_.InitBrowserMainThread();
+
   file_task_runner_ = base::ThreadPool::CreateSequencedTaskRunner(
       {base::MayBlock(), base::TaskPriority::BEST_EFFORT});
 
@@ -307,17 +309,18 @@ void HeadlessShell::ShutdownSoon() {
   if (shutdown_pending_)
     return;
   shutdown_pending_ = true;
+
   DCHECK(browser_);
-  if (web_contents_)
-    web_contents_->Close();
-  DCHECK(!web_contents_);
   browser_->BrowserMainThread()->PostTask(
       FROM_HERE,
       base::BindOnce(&HeadlessShell::Shutdown, weak_factory_.GetWeakPtr()));
 }
 
 void HeadlessShell::Shutdown() {
+  if (web_contents_)
+    web_contents_->Close();
   DCHECK(!web_contents_);
+
   browser_->Shutdown();
 }