Adds support for headless chrome embedder mojo services

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 96c90ae99dadcb266392a12af636e1b76b2faabc..c3845589241aab37e85b084ac94f4ea8cede9529 100644
--- a/headless/lib/browser/headless_web_contents_impl.cc
+++ b/headless/lib/browser/headless_web_contents_impl.cc
@@ -5,8 +5,10 @@
 #include "headless/lib/browser/headless_web_contents_impl.h"
 
 #include "base/bind.h"
+#include "base/json/json_writer.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/weak_ptr.h"
+#include "base/strings/utf_string_conversions.h"
 #include "base/trace_event/trace_event.h"
 #include "content/public/browser/devtools_agent_host.h"
 #include "content/public/browser/navigation_handle.h"
@@ -16,7 +18,6 @@
 #include "content/public/browser/render_widget_host_view.h"
 #include "content/public/browser/web_contents.h"
 #include "content/public/browser/web_contents_delegate.h"
-#include "content/public/browser/web_contents_observer.h"
 #include "content/public/common/bindings_policy.h"
 #include "content/public/common/service_registry.h"
 #include "content/public/renderer/render_frame.h"
@@ -79,6 +80,8 @@ std::unique_ptr<HeadlessWebContentsImpl> HeadlessWebContentsImpl::Create(
       base::WrapUnique(new HeadlessWebContentsImpl(
           content::WebContents::Create(create_params), browser));
 
+  headless_web_contents->mojo_services_ = std::move(builder->mojo_services_);
+  headless_web_contents->mojo_bindings_ = std::move(builder->mojo_bindings_);
   headless_web_contents->InitializeScreen(parent_window, builder->window_size_);
   if (!headless_web_contents->OpenURL(builder->initial_url_))
     return nullptr;
@@ -113,7 +116,8 @@ void HeadlessWebContentsImpl::InitializeScreen(aura::Window* parent_window,
 HeadlessWebContentsImpl::HeadlessWebContentsImpl(
     content::WebContents* web_contents,
     HeadlessBrowserImpl* browser)
-    : web_contents_delegate_(new HeadlessWebContentsImpl::Delegate(browser)),
+    : content::WebContentsObserver(web_contents),
+      web_contents_delegate_(new HeadlessWebContentsImpl::Delegate(browser)),
       web_contents_(web_contents),
       browser_(browser) {
   web_contents_->SetDelegate(web_contents_delegate_.get());
@@ -123,6 +127,43 @@ HeadlessWebContentsImpl::~HeadlessWebContentsImpl() {
   web_contents_->Close();
 }
 
+void HeadlessWebContentsImpl::RenderFrameCreated(
+    content::RenderFrameHost* render_frame_host) {
+  if (!mojo_bindings_.empty()) {
+    render_frame_host->GetRenderViewHost()->AllowBindings(
+        content::BINDINGS_POLICY_MOJO);
+  }
+
+  content::ServiceRegistry* service_registry =
+      render_frame_host->GetServiceRegistry();
+
+  for (const MojoService& service : mojo_services_) {
+    service_registry->AddService(service.service_name, service.service_factory,
+                                 browser_->BrowserMainThread());
+  }
+}
+
+void HeadlessWebContentsImpl::DocumentOnLoadCompletedInMainFrame() {
+  // Install mojo bindings by executing the bindings js.  Note this has to be
+  // done here since this is the earliest callback the browser has after the
+  // mojo module service is created at document start.
+  content::RenderFrameHost* render_frame_host = web_contents_->GetMainFrame();
+  base::ListValue module_list;
+  for (const MojoBindings& binding : mojo_bindings_) {
+    // Note RenderFrameHostImpl is unaware of headless chrome and we have to
+    // work around it's limits on when it's safe to call JS by using the
+    // ForTests version.
+    // TODO(alexclarke): Consider fixing this.
+    render_frame_host->ExecuteJavaScriptForTests(
+        base::UTF8ToUTF16(binding.js_bindings));
+    module_list.AppendString(binding.mojom_name);
+  }
+  std::string json_module_list;
+  CHECK(base::JSONWriter::Write(module_list, &json_module_list));
+  render_frame_host->ExecuteJavaScriptForTests(base::UTF8ToUTF16(
+      "mojo.resolvePromisesForServices_(" + json_module_list + ");"));
+}
+
 bool HeadlessWebContentsImpl::OpenURL(const GURL& url) {
   if (!url.is_valid())
     return false;
@@ -196,8 +237,41 @@ HeadlessWebContents::Builder& HeadlessWebContents::Builder::SetBrowserContext(
   return *this;
 }
 
+HeadlessWebContents::Builder& HeadlessWebContents::Builder::AddMojoService(
+    const std::string& service_name,
+    const base::Callback<void(mojo::ScopedMessagePipeHandle)>&
+        service_factory) {
+  mojo_services_.emplace_back(service_name, service_factory);
+  return *this;
+}
+
+HeadlessWebContents::Builder& HeadlessWebContents::Builder::AddJsMojoBindings(
+    const std::string& mojom_name,
+    const std::string& js_bindings) {
+  mojo_bindings_.emplace_back(mojom_name, js_bindings);
+  return *this;
+}
+
 HeadlessWebContents* HeadlessWebContents::Builder::Build() {
   return browser_->CreateWebContents(this);
 }
 
+HeadlessWebContents::Builder::MojoService::MojoService() {}
+
+HeadlessWebContents::Builder::MojoService::MojoService(
+    const std::string& service_name,
+    const base::Callback<void(mojo::ScopedMessagePipeHandle)>& service_factory)
+    : service_name(service_name), service_factory(service_factory) {}
+
+HeadlessWebContents::Builder::MojoService::~MojoService() {}
+
+HeadlessWebContents::Builder::MojoBindings::MojoBindings() {}
+
+HeadlessWebContents::Builder::MojoBindings::MojoBindings(
+    const std::string& mojom_name,
+    const std::string& js_bindings)
+    : mojom_name(mojom_name), js_bindings(js_bindings) {}
+
+HeadlessWebContents::Builder::MojoBindings::~MojoBindings() {}
+
 }  // namespace headless