Stop injection when injections are invalidated

extensions/renderer/script_injection_manager.cc

 // Copyright 2014 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "extensions/renderer/script_injection_manager.h"
 
 #include <memory>
 #include <utility>
 
 #include "base/auto_reset.h"
@@ skipping 85 matching lines @@
   void InvalidateAndResetFrame();
 
   // The owning ScriptInjectionManager.
   ScriptInjectionManager* manager_;
 
   bool should_run_idle_;
 
   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),
       manager_(manager),
       should_run_idle_(true),
       weak_factory_(this) {
 }
 
 ScriptInjectionManager::RFOHelper::~RFOHelper() {
 }
@@ skipping 162 matching lines @@
     const std::string& extension_id) {
   for (auto iter = pending_injections_.begin();
       iter != pending_injections_.end();) {
     if ((*iter)->host_id().id() == extension_id) {
       (*iter)->OnHostRemoved();
       iter = pending_injections_.erase(iter);
     } else {
       ++iter;
     }
   }
+
+  for (auto watcher : injection_watchers_) {
+    if (watcher->host_id().id() == extension_id)
+      watcher->InvalidateHost();
+  }
 }
 
 void ScriptInjectionManager::OnInjectionFinished(
     ScriptInjection* injection) {
   auto iter =
       std::find_if(running_injections_.begin(), running_injections_.end(),
                    [injection](const std::unique_ptr<ScriptInjection>& mode) {
                      return injection == mode.get();
                    });
   if (iter != running_injections_.end())
     running_injections_.erase(iter);
 }
 
 void ScriptInjectionManager::OnUserScriptsUpdated(
     const std::set<HostID>& changed_hosts,
     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();
   }
 }
 
 void ScriptInjectionManager::RemoveObserver(RFOHelper* helper) {
   for (auto iter = rfo_helpers_.begin(); iter != rfo_helpers_.end(); ++iter) {
     if (iter->get() == helper) {
       rfo_helpers_.erase(iter);
       break;
     }
   }
@@ skipping 74 matching lines @@
   active_injection_frames_.insert(frame);
 
   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))
       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);
   }
 
   // We are done running in the frame.
   active_injection_frames_.erase(frame);
 
   scripts_run_info.LogRun(activity_logging_enabled_);
 }
 
 void ScriptInjectionManager::TryToInject(
@@ skipping 10 matching lines @@
       scripts_run_info,
       base::Bind(&ScriptInjectionManager::OnInjectionFinished,
                  base::Unretained(this)))) {
     case ScriptInjection::INJECTION_WAITING:
       pending_injections_.push_back(std::move(injection));
       break;
     case ScriptInjection::INJECTION_BLOCKED:
       running_injections_.push_back(std::move(injection));
       break;
     case ScriptInjection::INJECTION_FINISHED:
+    case ScriptInjection::INJECTION_CANCELED:
       break;
   }
 }
 
 void ScriptInjectionManager::HandleExecuteCode(
     const ExtensionMsg_ExecuteCode_Params& params,
     content::RenderFrame* render_frame) {
   std::unique_ptr<const InjectionHost> injection_host;
   if (params.host_id.type() == HostID::EXTENSIONS) {
     injection_host = ExtensionInjectionHost::Create(params.host_id.id());
     if (!injection_host)
       return;
   } else if (params.host_id.type() == HostID::WEBUI) {
     injection_host.reset(
         new WebUIInjectionHost(params.host_id));
   }
 
   std::unique_ptr<ScriptInjection> injection(new ScriptInjection(
       std::unique_ptr<ScriptInjector>(
           new ProgrammaticScriptInjector(params, render_frame)),
       render_frame, std::move(injection_host),
       static_cast<UserScript::RunLocation>(params.run_at),
       activity_logging_enabled_));
 
   FrameStatusMap::const_iterator iter = frame_statuses_.find(render_frame);
   UserScript::RunLocation run_location =
       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);
 }
 
 void ScriptInjectionManager::HandleExecuteDeclarativeScript(
     content::RenderFrame* render_frame,
     int tab_id,
     const ExtensionId& extension_id,
     int script_id,
     const GURL& url) {
   std::unique_ptr<ScriptInjection> injection =
       user_script_set_manager_->GetInjectionForDeclarativeScript(
           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);
 
     scripts_run_info.LogRun(activity_logging_enabled_);
   }
 }
 
 void ScriptInjectionManager::HandlePermitScriptInjection(int64_t request_id) {
   auto iter = pending_injections_.begin();
   for (; iter != pending_injections_.end(); ++iter) {
     if ((*iter)->request_id() == request_id) {
       DCHECK((*iter)->host_id().type() == HostID::EXTENSIONS);
       break;
     }
   }
   if (iter == pending_injections_.end())
     return;
 
   // At this point, because the request is present in pending_injections_, we
   // know that this is the same page that issued the request (otherwise,
   // RFOHelper::InvalidateAndResetFrame would have caused it to be cleared out).
 
   std::unique_ptr<ScriptInjection> injection(std::move(*iter));
   pending_injections_.erase(iter);
 
   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)
     running_injections_.push_back(std::move(injection));
   scripts_run_info.LogRun(activity_logging_enabled_);
 }
 
 }  // namespace extensions