Add HeadlessTabSocket

headless/lib/browser/headless_web_contents_impl.cc

Index: headless/lib/browser/headless_web_contents_impl.cc
diff --git a/headless/lib/browser/headless_web_contents_impl.cc b/headless/lib/browser/headless_web_contents_impl.cc
index a967c3c0f9bbd69a55b133e0ef6b1af9ea8f412c..ad43eb9b12e01dc9497c381538041e3892fdeeb7 100644
--- a/headless/lib/browser/headless_web_contents_impl.cc
+++ b/headless/lib/browser/headless_web_contents_impl.cc
@@ -31,6 +31,7 @@
 #include "headless/lib/browser/headless_browser_impl.h"
 #include "headless/lib/browser/headless_browser_main_parts.h"
 #include "headless/lib/browser/headless_devtools_client_impl.h"
+#include "headless/lib/browser/headless_tab_socket_impl.h"
 #include "printing/features/features.h"
 #include "services/service_manager/public/cpp/interface_registry.h"
 
@@ -145,6 +146,14 @@ std::unique_ptr<HeadlessWebContentsImpl> HeadlessWebContentsImpl::Create(
           content::WebContents::Create(create_params),
           builder->browser_context_));
 
+  if (builder->create_tab_socket_) {
+    headless_web_contents->headless_tab_socket_ =
+        base::MakeUnique<HeadlessTabSocketImpl>();
+    builder->AddMojoService(base::Bind(
+        &HeadlessTabSocketImpl::CreateMojoService,
+        base::Unretained(headless_web_contents->headless_tab_socket_.get())));
+  }
+
   headless_web_contents->mojo_services_ = std::move(builder->mojo_services_);
   headless_web_contents->InitializeScreen(builder->window_size_);
   if (!headless_web_contents->OpenURL(builder->initial_url_))
@@ -292,6 +301,10 @@ HeadlessBrowserContextImpl* HeadlessWebContentsImpl::browser_context() const {
   return browser_context_;
 }
 
+HeadlessTabSocket* HeadlessWebContentsImpl::GetHeadlessTabSocket() const {
+  return headless_tab_socket_.get();
+}
+
 HeadlessWebContents::Builder::Builder(
     HeadlessBrowserContextImpl* browser_context)
     : browser_context_(browser_context),
@@ -321,6 +334,12 @@ HeadlessWebContents::Builder& HeadlessWebContents::Builder::AddMojoService(
   return *this;
 }
 
+HeadlessWebContents::Builder& HeadlessWebContents::Builder::CreateTabSocket(
+    bool create_tab_socket) {
+  create_tab_socket_ = create_tab_socket;
+  return *this;
+}
+
 HeadlessWebContents* HeadlessWebContents::Builder::Build() {
   return browser_context_->CreateWebContents(this);
 }