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..b952b23a66709267dbe0bc1a5bc05cbd4520156a 100644
--- a/extensions/renderer/script_injection_manager.cc
+++ b/extensions/renderer/script_injection_manager.cc
@@ -103,6 +103,36 @@ class ScriptInjectionManager::RFOHelper : public content::RenderFrameObserver {
   base::WeakPtrFactory<RFOHelper> weak_factory_;
 };
 
+class ScriptInjectionManager::ScriptInjectionWatcher {
+ 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;
+      }
+    }
+  }
+
+  void InvalidateHost() {
+    injection_->OnHostRemoved();
+  }
+
+  const HostID& host_id() const { return host_id_; }
+
+ private:
+  ScriptInjection* injection_;
+  ScriptInjectionManager* manager_;
+  HostID host_id_;
+};
+
 ScriptInjectionManager::RFOHelper::RFOHelper(content::RenderFrame* render_frame,
                                              ScriptInjectionManager* manager)
     : content::RenderFrameObserver(render_frame),
@@ -285,6 +315,11 @@ void ScriptInjectionManager::OnExtensionUnloaded(
       ++iter;
     }
   }
+
+  for (auto watcher : injection_watchers_) {
+    if (watcher->host_id().id() == extension_id)
+      watcher->InvalidateHost();
+  }
 }
 
 void ScriptInjectionManager::OnInjectionFinished(
@@ -303,10 +338,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();
   }
 }
 
@@ -401,6 +443,7 @@ void ScriptInjectionManager::InjectScripts(
       break;
     std::unique_ptr<ScriptInjection> injection(std::move(*iter));
     iter = frame_injections.erase(iter);
+    ScriptInjectionWatcher watcher(injection.get(), this);

[Devlin, 2016/07/08 20:51:09]

Hmm... I think this isn't quite sufficient.  Consider the following scenario:
Extensions A and B both want to inject on a page.
Extension A injects, and blocks.
During blocking, Extension B is removed
  We only have a watcher for Extension A (the injecting one)
Extension B starts injecting, and blows up

Unless I'm missing some way that's mitigated?

Instead, I think we have to keep track of both pending and active injections
during these times, and notify both of extension host removal.  This will
probably render the Watcher class unnecessary.

[robwu, 2016/07/09 05:22:02]

You're right - I solved the issue by constructing a queue of
ScriptInjectionWatchers. I also moved the frame logic to
active_injection_frames_, so that we can have an early out in the
ScriptInjection::Inject (similarly to |injection_host_|).
I added the class to avoid spreeing push_back with a balanced find & erase all
over the place. With the watcher class, the watcher is set up when it's
constructed and automatically removed upon destruction.

     TryToInject(std::move(injection), run_location, &scripts_run_info);
   }
 
@@ -431,6 +474,7 @@ void ScriptInjectionManager::TryToInject(
       running_injections_.push_back(std::move(injection));
       break;
     case ScriptInjection::INJECTION_FINISHED:
+    case ScriptInjection::INJECTION_CANCELED:
       break;
   }
 }
@@ -460,6 +504,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 +519,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 +548,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)