From 94fce43ed918cd5a4af60f0fce2cd746dce49c19 Mon Sep 17 00:00:00 2001
From: Cheng Zhao <zcbenz@gmail.com>
Date: Thu, 15 Mar 2018 15:16:30 +0900
Subject: [PATCH] Emit document-start for the correct env

This fixes the crash in RunScriptsAtDocumentStart when "affinity" option
is specified. Previously we were assuming only one main frame exists in
the renderer process, but the "affinity" option breaks this option.

There is also a bug that "node::Environment::GetCurrent" does not return
nullptr for context without a env in it, I'm not sure whether it is a
bug of Node or V8.
---
 atom/renderer/atom_renderer_client.cc | 40 +++++++++++++++++----------
 atom/renderer/atom_renderer_client.h  | 17 ++++++++++++
 2 files changed, 43 insertions(+), 14 deletions(-)

diff --git a/atom/renderer/atom_renderer_client.cc b/atom/renderer/atom_renderer_client.cc
index 5bcf7e0f7104..2ba4e9dc6596 100644
--- a/atom/renderer/atom_renderer_client.cc
+++ b/atom/renderer/atom_renderer_client.cc
@@ -64,21 +64,19 @@ void AtomRendererClient::RenderViewCreated(content::RenderView* render_view) {
 void AtomRendererClient::RunScriptsAtDocumentStart(
     content::RenderFrame* render_frame) {
   // Inform the document start pharse.
-  node::Environment* env = node_bindings_->uv_env();
-  if (env) {
-    v8::HandleScope handle_scope(env->isolate());
+  v8::HandleScope handle_scope(v8::Isolate::GetCurrent());
+  node::Environment* env = GetEnvironment(render_frame);
+  if (env)
     mate::EmitEvent(env->isolate(), env->process_object(), "document-start");
-  }
 }
 
 void AtomRendererClient::RunScriptsAtDocumentEnd(
     content::RenderFrame* render_frame) {
   // Inform the document end pharse.
-  node::Environment* env = node_bindings_->uv_env();
-  if (env) {
-    v8::HandleScope handle_scope(env->isolate());
+  v8::HandleScope handle_scope(v8::Isolate::GetCurrent());
+  node::Environment* env = GetEnvironment(render_frame);
+  if (env)
     mate::EmitEvent(env->isolate(), env->process_object(), "document-end");
-  }
 }
 
 void AtomRendererClient::DidCreateScriptContext(
@@ -88,6 +86,8 @@ void AtomRendererClient::DidCreateScriptContext(
   if (!render_frame->IsMainFrame() && !IsDevToolsExtension(render_frame))
     return;
 
+  injected_frames_.insert(render_frame);
+
   // Prepare the node bindings.
   if (!node_integration_initialized_) {
     node_integration_initialized_ = true;
@@ -102,6 +102,7 @@ void AtomRendererClient::DidCreateScriptContext(
 
   // Setup node environment for each window.
   node::Environment* env = node_bindings_->CreateEnvironment(context);
+  environments_.insert(env);
 
   // Add Electron extended APIs.
   atom_bindings_->BindTo(env->isolate(), env->process_object());
@@ -121,14 +122,14 @@ void AtomRendererClient::DidCreateScriptContext(
 
 void AtomRendererClient::WillReleaseScriptContext(
     v8::Handle<v8::Context> context, content::RenderFrame* render_frame) {
-  // Only allow node integration for the main frame, unless it is a devtools
-  // extension page.
-  if (!render_frame->IsMainFrame() && !IsDevToolsExtension(render_frame))
-    return;
+  injected_frames_.erase(render_frame);
 
   node::Environment* env = node::Environment::GetCurrent(context);
-  if (env)
-    mate::EmitEvent(env->isolate(), env->process_object(), "exit");
+  if (environments_.find(env) == environments_.end())
+    return;
+  environments_.erase(env);
+
+  mate::EmitEvent(env->isolate(), env->process_object(), "exit");
 
   // The main frame may be replaced.
   if (env == node_bindings_->uv_env())
@@ -209,5 +210,16 @@ void AtomRendererClient::SetupMainWorldOverrides(
   ignore_result(func->Call(context, v8::Null(isolate), 1, args));
 }
 
+node::Environment* AtomRendererClient::GetEnvironment(
+    content::RenderFrame* render_frame) const {
+  if (injected_frames_.find(render_frame) == injected_frames_.end())
+    return nullptr;
+  v8::HandleScope handle_scope(v8::Isolate::GetCurrent());
+  node::Environment* env = node::Environment::GetCurrent(
+      render_frame->GetWebFrame()->MainWorldScriptContext());
+  if (environments_.find(env) == environments_.end())
+    return nullptr;
+  return env;
+}
 
 }  // namespace atom
diff --git a/atom/renderer/atom_renderer_client.h b/atom/renderer/atom_renderer_client.h
index 0c411d2d35d5..40e7205c3b50 100644
--- a/atom/renderer/atom_renderer_client.h
+++ b/atom/renderer/atom_renderer_client.h
@@ -5,11 +5,16 @@
 #ifndef ATOM_RENDERER_ATOM_RENDERER_CLIENT_H_
 #define ATOM_RENDERER_ATOM_RENDERER_CLIENT_H_
 
+#include <set>
 #include <string>
 #include <vector>
 
 #include "atom/renderer/renderer_client_base.h"
 
+namespace node {
+class Environment;
+}
+
 namespace atom {
 
 class AtomBindings;
@@ -54,12 +59,24 @@ class AtomRendererClient : public RendererClientBase {
   void WillDestroyWorkerContextOnWorkerThread(
       v8::Local<v8::Context> context) override;
 
+  node::Environment* GetEnvironment(content::RenderFrame* frame) const;
+
   // Whether the node integration has been initialized.
   bool node_integration_initialized_;
 
   std::unique_ptr<NodeBindings> node_bindings_;
   std::unique_ptr<AtomBindings> atom_bindings_;
 
+  // The node::Environment::GetCurrent API does not return nullptr when it
+  // is called for a context without node::Environment, so we have to keep
+  // a book of the environments created.
+  std::set<node::Environment*> environments_;
+
+  // Getting main script context from web frame would lazily initializes
+  // its script context. Doing so in a web page without scripts would trigger
+  // assertion, so we have to keep a book of injected web frames.
+  std::set<content::RenderFrame*> injected_frames_;
+
   DISALLOW_COPY_AND_ASSIGN(AtomRendererClient);
 };