Stop injection when injections are invalidated

extensions/renderer/script_injection_manager.cc

Index: extensions/renderer/script_injection_manager.cc
diff --git a/extensions/renderer/script_injection_manager.cc b/extensions/renderer/script_injection_manager.cc
index f6743c12b0f5de6697bede30d7bc930820059229..1e06e5eedd7eb0be04ea856da33eca4d8081d10b 100644
--- a/extensions/renderer/script_injection_manager.cc
+++ b/extensions/renderer/script_injection_manager.cc
@@ -4,6 +4,7 @@
 
 #include "extensions/renderer/script_injection_manager.h"
 
+#include <deque>
 #include <memory>
 #include <utility>
 
@@ -103,6 +104,50 @@ class ScriptInjectionManager::RFOHelper : public content::RenderFrameObserver {
   base::WeakPtrFactory<RFOHelper> weak_factory_;
 };
 
+// Keeps track of a ScriptInjection while it is injected. The caller must make
+// sure that the ScriptInjection is not destructed when the Invalidate* methods
+// are called.
+class ScriptInjectionManager::ScriptInjectionWatcher {

[Devlin, 2016/07/11 23:30:55]

I'm still not quite happy with this implementation.  It seems a little like it's
more complicated than it should be (I'm worried we're going to forget to put one
of these somewhere, etc), while also not fully addressing the concerns (e.g.
what if a ScriptInjection blocks, and finishes asynchronously? The injection is
stored in running_injections_, but it doesn't look like we'll update the
extension or frame for those).

I wonder if there's a way that we can put a little more of this logic into the
ScriptInjection/InjectionHost.  For one thing, we now have the
RendererExtensionRegistry(), which we could check for the presence of the
extension.  Then rather than having these updates, we can just check
injection_host_->IsValid() at set times throughout the injection process. 
Unfortunately, that doesn't answer the question of frames.

I'll think on this a bit more and circle back, but if you had any thoughts,
happy to hear them. :)

[robwu, 2016/07/12 08:59:36]

We only need to be careful when InjectJs() is called.
Currently, when a script is *not* being injected, it is stored in
|pending_injections_| or |running_injections_|. If it *is* about to be injected,
there will always be a ScriptInjectionWatcher on the stack, which was designed
to be resilient against unexpected reentrancy (with respect to the same
frame/injection).
Scripts injection should be aborted in the following cases:

1. Extension unloads.
2. Extension reloads.
3. Frame unloads (includes frame reloads).

Using RendererExtensionRegistry seems to only cover case 1, not case 2.
I only have half-baked thoughts for now. The issue is that InjectionHost and
ScriptInjection can outlive each other, so there needs to be some mediator that
outlives both (or some token that is reset whenever either of the invalidation
events occur). I found this role in the UserScriptManager (and the new watcher).

+ public:
+  ScriptInjectionWatcher(ScriptInjection* injection,
+                         ScriptInjectionManager* manager)
+    : injection_(injection), manager_(manager), host_id_(injection->host_id()) {
+    manager_->injection_watchers_.push_back(this);
+  }
+
+  ~ScriptInjectionWatcher() {
+    auto& watchers = manager_->injection_watchers_;
+    for (auto it = watchers.begin(); it != watchers.end(); ++it) {
+      if (*it == this) {
+        watchers.erase(it);
+        break;
+      }
+    }
+  }
+
+  // Mark the injection as invalid because the extension disappeared.
+  // This method can safely be called multiple times.
+  void InvalidateHost() {
+    injection_->OnHostRemoved();
+  }
+
+  // Mark the injection as invalid because the frame was invalidated.
+  // This method can safely be called multiple times.
+  void InvalidateIfFrameEquals(content::RenderFrame* frame) {
+    if (injection_->render_frame() == frame)
+      injection_->invalidate_render_frame();
+  }
+
+  const HostID& host_id() const { return host_id_; }
+
+ private:
+  ScriptInjection* injection_;  // Not owned, must outlive us.
+  ScriptInjectionManager* manager_;  // Owns us.
+  HostID host_id_;
+
+  DISALLOW_COPY_AND_ASSIGN(ScriptInjectionWatcher);
+};
+
 ScriptInjectionManager::RFOHelper::RFOHelper(content::RenderFrame* render_frame,
                                              ScriptInjectionManager* manager)
     : content::RenderFrameObserver(render_frame),
@@ -267,6 +312,7 @@ ScriptInjectionManager::~ScriptInjectionManager() {
     injection->invalidate_render_frame();
   for (const auto& injection : running_injections_)
     injection->invalidate_render_frame();
+  DCHECK(injection_watchers_.empty());
 }
 
 void ScriptInjectionManager::OnRenderFrameCreated(
@@ -285,6 +331,11 @@ void ScriptInjectionManager::OnExtensionUnloaded(
       ++iter;
     }
   }
+
+  for (auto watcher : injection_watchers_) {
+    if (watcher->host_id().id() == extension_id)
+      watcher->InvalidateHost();
+  }
 }
 
 void ScriptInjectionManager::OnInjectionFinished(
@@ -303,10 +354,17 @@ void ScriptInjectionManager::OnUserScriptsUpdated(
     const std::vector<UserScript*>& scripts) {
   for (auto iter = pending_injections_.begin();
        iter != pending_injections_.end();) {
-    if (changed_hosts.count((*iter)->host_id()) > 0)
+    if (changed_hosts.count((*iter)->host_id()) > 0) {
+      (*iter)->OnHostRemoved();
       iter = pending_injections_.erase(iter);
-    else
+    } else {
       ++iter;
+    }
+  }
+
+  for (auto watcher : injection_watchers_) {
+    if (changed_hosts.count(watcher->host_id()) > 0)
+      watcher->InvalidateHost();
   }
 }
 
@@ -320,10 +378,6 @@ void ScriptInjectionManager::RemoveObserver(RFOHelper* helper) {
 }
 
 void ScriptInjectionManager::InvalidateForFrame(content::RenderFrame* frame) {
-  // If the frame invalidated is the frame being injected into, we need to
-  // note it.
-  active_injection_frames_.erase(frame);
-
   for (auto iter = pending_injections_.begin();
        iter != pending_injections_.end();) {
     if ((*iter)->render_frame() == frame)
@@ -332,6 +386,10 @@ void ScriptInjectionManager::InvalidateForFrame(content::RenderFrame* frame) {
       ++iter;
   }
 
+  for (auto watcher : injection_watchers_) {
+    watcher->InvalidateIfFrameEquals(frame);
+  }
+
   frame_statuses_.erase(frame);
 }
 
@@ -390,22 +448,30 @@ void ScriptInjectionManager::InjectScripts(
   user_script_set_manager_->GetAllInjections(&frame_injections, frame, tab_id,
                                              run_location);
 
-  // Note that we are running in |frame|.
-  active_injection_frames_.insert(frame);
+  // At this point the ScriptInjections in frame_injections are not reference
+  // anywhere else. Set up ScriptInjectionWatchers so that the injections can be
+  // invalidated if any of the scripts destroys the frame or injection host.
+  std::deque<std::unique_ptr<ScriptInjectionWatcher>> scoped_watchers;
+  for (auto& iter : frame_injections) {
+    scoped_watchers.push_back(
+        base::MakeUnique<ScriptInjectionWatcher>(iter.get(), this));
+  }
 
   ScriptsRunInfo scripts_run_info(frame, run_location);
   for (auto iter = frame_injections.begin(); iter != frame_injections.end();) {
     // It's possible for the frame to be invalidated in the course of injection
     // (if a script removes its own frame, for example). If this happens, abort.
-    if (!active_injection_frames_.count(frame))
+    if (!(*iter)->render_frame())
       break;
     std::unique_ptr<ScriptInjection> injection(std::move(*iter));
     iter = frame_injections.erase(iter);
     TryToInject(std::move(injection), run_location, &scripts_run_info);
-  }
 
-  // We are done running in the frame.
-  active_injection_frames_.erase(frame);
+    // Destroy the watcher because the ScriptInjection* may be invalid, and
+    // running the next iteration of the loop may trigger observers that touch
+    // the active watchers.
+    scoped_watchers.pop_front();
+  }
 
   scripts_run_info.LogRun(activity_logging_enabled_);
 }
@@ -431,6 +497,7 @@ void ScriptInjectionManager::TryToInject(
       running_injections_.push_back(std::move(injection));
       break;
     case ScriptInjection::INJECTION_FINISHED:
+    case ScriptInjection::INJECTION_CANCELED:
       break;
   }
 }
@@ -460,6 +527,7 @@ void ScriptInjectionManager::HandleExecuteCode(
       iter == frame_statuses_.end() ? UserScript::UNDEFINED : iter->second;
 
   ScriptsRunInfo scripts_run_info(render_frame, run_location);
+  ScriptInjectionWatcher watcher(injection.get(), this);
   TryToInject(std::move(injection), run_location, &scripts_run_info);
 }
 
@@ -474,6 +542,7 @@ void ScriptInjectionManager::HandleExecuteDeclarativeScript(
           script_id, render_frame, tab_id, url, extension_id);
   if (injection.get()) {
     ScriptsRunInfo scripts_run_info(render_frame, UserScript::BROWSER_DRIVEN);
+    ScriptInjectionWatcher watcher(injection.get(), this);
     // TODO(markdittmer): Use return value of TryToInject for error handling.
     TryToInject(std::move(injection), UserScript::BROWSER_DRIVEN,
                 &scripts_run_info);
@@ -502,6 +571,7 @@ void ScriptInjectionManager::HandlePermitScriptInjection(int64_t request_id) {
 
   ScriptsRunInfo scripts_run_info(injection->render_frame(),
                                   UserScript::RUN_DEFERRED);
+  ScriptInjectionWatcher watcher(injection.get(), this);
   ScriptInjection::InjectionResult res = injection->OnPermissionGranted(
       &scripts_run_info);
   if (res == ScriptInjection::INJECTION_BLOCKED)