[Reland] Allow headless TabSocket in isolated worlds & remove obsolete logic

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 b627f692601aabfc5bfe377a2261f2de20a02c89..fddf9ecc8ef85dc1e7577829e39a780420162b85 100644
--- a/headless/lib/renderer/headless_content_renderer_client.cc
+++ b/headless/lib/renderer/headless_content_renderer_client.cc
@@ -7,8 +7,16 @@
 #include "base/memory/ptr_util.h"
 #include "base/strings/utf_string_conversions.h"
 #include "content/public/common/bindings_policy.h"
+#include "content/public/common/isolated_world_ids.h"
 #include "content/public/renderer/render_frame.h"
+#include "content/public/renderer/render_frame_observer.h"
+#include "gin/handle.h"
+#include "gin/object_template_builder.h"
+#include "gin/wrappable.h"
+#include "headless/lib/tab_socket.mojom.h"
 #include "printing/features/features.h"
+#include "services/service_manager/public/cpp/interface_provider.h"
+#include "third_party/WebKit/public/web/WebKit.h"
 
 #if BUILDFLAG(ENABLE_BASIC_PRINTING)
 #include "components/printing/renderer/print_web_view_helper.h"
@@ -21,77 +29,162 @@ HeadlessContentRendererClient::HeadlessContentRendererClient() {}
 
 HeadlessContentRendererClient::~HeadlessContentRendererClient() {}
 
+// Keep in sync with DOMWrapperWorld::WorldId.
+enum {
+  kDevToolsFirstIsolatedWorldId = (1 << 29) + 1,

[jochen (gone - plz use gerrit), 2017/05/15 05:41:56]

can you pull those IDs into a separate header, so you can just include it here?

[alex clarke (OOO till 29th), 2017/05/15 09:50:57]

It would be nice if we could share the same header from inside blink too.  Not
sure where such a header would live however.  For now I've just added it in the
same directory as this file.

+  kDevToolsLastIsolatedWorldId = kDevToolsFirstIsolatedWorldId + 100,
+};
+
+v8::Local<v8::String> ASCIIStringToV8String(v8::Isolate* isolate,

[jochen (gone - plz use gerrit), 2017/05/15 05:41:56]

that's gin::Converter<std::string>::ToV8()

[alex clarke (OOO till 29th), 2017/05/15 09:50:57]

Done.

+                                            const std::string& s) {
+  DCHECK(base::IsStringASCII(s));
+  return v8::String::NewFromUtf8(isolate, s.data(), v8::NewStringType::kNormal,
+                                 s.size())
+      .ToLocalChecked();
+}
+
+class HeadlessIsolatedWorldTabSocketBindings

[Sami, 2017/05/12 17:18:02]

s/Isolated// since this is also used for the main world?

[alex clarke (OOO till 29th), 2017/05/15 09:50:57]

Good point, this was originally isolated only but it morphed into this.

+    : public gin::Wrappable<HeadlessIsolatedWorldTabSocketBindings>,
+      public content::RenderFrameObserver {
+ public:
+  explicit HeadlessIsolatedWorldTabSocketBindings(
+      content::RenderFrame* render_frame)
+      : content::RenderFrameObserver(render_frame) {}
+  ~HeadlessIsolatedWorldTabSocketBindings() override {}
+
+  // content::RenderFrameObserver implementation:
+  void DidCreateScriptContext(v8::Local<v8::Context> context,
+                              int world_id) override {
+    if (world_id == content::ISOLATED_WORLD_ID_GLOBAL) {
+      // For the main world only inject TabSocket if
+      // BINDINGS_POLICY_HEADLESS_MAIN_WORLD is set.
+      if (!(render_frame()->GetEnabledBindings() &
+            content::BindingsPolicy::BINDINGS_POLICY_HEADLESS_MAIN_WORLD)) {
+        return;
+      }
+    } else {
+      // For the isolated worlds only inject TabSocket if
+      // BINDINGS_POLICY_HEADLESS_ISOLATED_WORLD is set and the world id falls
+      // within the range reserved for DevTools created isolated worlds.
+      if (!(render_frame()->GetEnabledBindings() &
+            content::BindingsPolicy::BINDINGS_POLICY_HEADLESS_ISOLATED_WORLD)) {
+        return;
+      }
+      if (world_id < kDevToolsFirstIsolatedWorldId ||
+          world_id > kDevToolsLastIsolatedWorldId) {
+        return;
+      }
+    }
+
+    InitializeTabSocketBindings(context);
+  }
+
+  void OnDestruct() override {}
+
+  // gin::WrappableBase implementation:
+  gin::ObjectTemplateBuilder GetObjectTemplateBuilder(
+      v8::Isolate* isolate) override {
+    return gin::Wrappable<HeadlessIsolatedWorldTabSocketBindings>::
+        GetObjectTemplateBuilder(isolate)
+            .SetMethod("send",
+                       &HeadlessIsolatedWorldTabSocketBindings::SendImpl)
+            .SetProperty("onmessage",
+                         &HeadlessIsolatedWorldTabSocketBindings::GetOnMessage,
+                         &HeadlessIsolatedWorldTabSocketBindings::SetOnMessage);
+  }
+
+  static gin::WrapperInfo kWrapperInfo;
+
+ private:
+  void SendImpl(const std::string& message) {
+    EnsureTabSocketPtr()->SendMessageToEmbedder(message);
+  }
+
+  v8::Local<v8::Value> GetOnMessage() { return GetOnMessageCallback(); }
+
+  void SetOnMessage(v8::Local<v8::Function> callback) {
+    GetOnMessageCallback_.Reset(blink::MainThreadIsolate(), callback);
+
+    EnsureTabSocketPtr()->AwaitNextMessageFromEmbedder(base::Bind(
+        &HeadlessIsolatedWorldTabSocketBindings::OnNextMessageFromEmbedder,
+        base::Unretained(this)));
+  }
+
+  void OnNextMessageFromEmbedder(const std::string& message) {
+    CallOnMessageCallback(message);
+  }
+
+  v8::Local<v8::Value> CallOnMessageCallback(const std::string& message) {
+    v8::Isolate* isolate = blink::MainThreadIsolate();
+    v8::HandleScope handle_scope(isolate);
+    v8::Local<v8::Context> context = GetContext();
+    v8::Context::Scope context_scope(context);
+    v8::Local<v8::Value> argv[] = {
+        ASCIIStringToV8String(isolate, message),

[Sami, 2017/05/12 17:18:02]

Hmm, is the message really guaranteed to ASCII?

[alex clarke (OOO till 29th), 2017/05/15 09:50:57]

That's the api contract.

+    };
+    v8::TryCatch try_catch(isolate);
+    v8::MicrotasksScope microtasks(isolate,
+                                   v8::MicrotasksScope::kDoNotRunMicrotasks);
+
+    v8::MaybeLocal<v8::Value> maybe_result = GetOnMessageCallback()->Call(
+        context, context->Global(), arraysize(argv), argv);
+
+    v8::Local<v8::Value> result;
+    if (!maybe_result.ToLocal(&result))
+      return v8::Local<v8::Value>();
+    return result;
+  }
+
+  void InitializeTabSocketBindings(v8::Local<v8::Context> context) {
+    // Add TabSocket bindings to the DevTools created isolated world.
+    v8::Isolate* isolate = blink::MainThreadIsolate();
+    v8::HandleScope handle_scope(isolate);
+    if (context.IsEmpty())
+      return;
+
+    v8::Context::Scope context_scope(context);
+    gin::Handle<HeadlessIsolatedWorldTabSocketBindings> bindings =
+        gin::CreateHandle(isolate, this);
+    if (bindings.IsEmpty())
+      return;
+
+    v8::Local<v8::Object> global = context->Global();
+    global->Set(gin::StringToV8(isolate, "TabSocket"), bindings.ToV8());
+    context_.Reset(blink::MainThreadIsolate(), context);

[jochen (gone - plz use gerrit), 2017/05/15 05:41:56]

this will create a memory leak, as the context now can never die, and so can
this object that lives in the context never die

[alex clarke (OOO till 29th), 2017/05/15 09:50:57]

I based this on what mojo does (see gin::ContextHolder & MojoMainRunner).  It
seems they release the UniquePersistents in WillReleaseScriptContext and I've
done the same, hopefully that works.  

+  }
+
+  headless::TabSocketPtr& EnsureTabSocketPtr() {
+    if (!tab_socket_ptr_.is_bound()) {
+      render_frame()->GetRemoteInterfaces()->GetInterface(
+          mojo::MakeRequest(&tab_socket_ptr_));
+    }
+    return tab_socket_ptr_;
+  }
+
+  v8::Local<v8::Context> GetContext() {
+    return v8::Local<v8::Context>::New(blink::MainThreadIsolate(), context_);
+  }
+
+  v8::Local<v8::Function> GetOnMessageCallback() {
+    return v8::Local<v8::Function>::New(blink::MainThreadIsolate(),
+                                        GetOnMessageCallback_);
+  }
+
+  headless::TabSocketPtr tab_socket_ptr_;
+  v8::UniquePersistent<v8::Context> context_;
+  v8::UniquePersistent<v8::Function> GetOnMessageCallback_;
+};
+
+gin::WrapperInfo HeadlessIsolatedWorldTabSocketBindings::kWrapperInfo = {
+    gin::kEmbedderNativeGin};
+
 void HeadlessContentRendererClient::RenderFrameCreated(
     content::RenderFrame* render_frame) {
 #if BUILDFLAG(ENABLE_BASIC_PRINTING)
   new printing::PrintWebViewHelper(
       render_frame, base::MakeUnique<HeadlessPrintWebViewHelperDelegate>());
 #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();
-      });
-
-      })(); )"));
+  new HeadlessIsolatedWorldTabSocketBindings(render_frame);
 }
 
 }  // namespace headless