Add HeadlessTabSocket

headless/lib/renderer/headless_content_renderer_client.cc

Index: headless/lib/renderer/headless_content_renderer_client.cc
diff --git a/headless/lib/renderer/headless_content_renderer_client.cc b/headless/lib/renderer/headless_content_renderer_client.cc
index c4c263e62a728378c6b8c298ffc0ae40b55c1dc7..b627f692601aabfc5bfe377a2261f2de20a02c89 100644
--- a/headless/lib/renderer/headless_content_renderer_client.cc
+++ b/headless/lib/renderer/headless_content_renderer_client.cc
@@ -5,6 +5,9 @@
 #include "headless/lib/renderer/headless_content_renderer_client.h"
 
 #include "base/memory/ptr_util.h"
+#include "base/strings/utf_string_conversions.h"
+#include "content/public/common/bindings_policy.h"
+#include "content/public/renderer/render_frame.h"
 #include "printing/features/features.h"
 
 #if BUILDFLAG(ENABLE_BASIC_PRINTING)
@@ -26,4 +29,69 @@ void HeadlessContentRendererClient::RenderFrameCreated(
 #endif
 }
 
+void HeadlessContentRendererClient::RunScriptsAtDocumentStart(
+    content::RenderFrame* render_frame) {
+  if (!(render_frame->GetEnabledBindings() &
+        content::BindingsPolicy::BINDINGS_POLICY_HEADLESS)) {
+    return;
+  }
+
+  render_frame->ExecuteJavaScript(base::UTF8ToUTF16(R"(
+      (function() {
+      let tabSocket = null;
+      let messagesToSend = [];
+
+      function getNextEmbedderMessage() {
+        tabSocket.awaitNextMessageFromEmbedder().then(
+            function(result) {
+              window.TabSocket.onmessage(new CustomEvent(
+                  "onmessage", {detail : {message : result.message}}));
+              getNextEmbedderMessage();
+            });
+      };
+
+      window.TabSocket = new Proxy(
+        {
+          send: function(message) {
+            if (tabSocket) {
+              tabSocket.sendMessageToEmbedder(message);
+            } else {
+              messagesToSend.push(message);
+            }
+          },
+          onmessage: null
+        },
+        {
+          set: function(target, prop, value, receiver) {
+            if (prop !== "onmessage")
+              return false;
+            target[prop] = value;
+            if (tabSocket)
+              getNextEmbedderMessage();
+            return true;
+          }
+        });
+
+      // Note define() defines a module in the mojo module dependency
+      // system. While we don't expose our module, the callback below only
+      // fires after the requested modules have been loaded.
+      define([
+          'headless/lib/tab_socket.mojom',
+          'content/public/renderer/frame_interfaces',
+          ], function(tabSocketMojom, frameInterfaces) {
+        tabSocket = new tabSocketMojom.TabSocketPtr(
+            frameInterfaces.getInterface(tabSocketMojom.TabSocket.name));
+        // Send any messages that may have been created before the dependency
+        // was resolved.
+        for (var i = 0; i < messagesToSend.length; i++) {
+          tabSocket.sendMessageToEmbedder(messagesToSend[i]);
+        }
+        messagesToSend = [];
+        if (window.TabSocket.onmessage)
+          getNextEmbedderMessage();
+      });
+
+      })(); )"));
+}
+
 }  // namespace headless