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 <deque>
 #include <memory>
 #include <utility>
 
 #include "base/auto_reset.h"
 #include "base/bind.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/weak_ptr.h"
 #include "base/values.h"
 #include "content/public/renderer/render_frame.h"
 #include "content/public/renderer/render_frame_observer.h"
@@ skipping 79 matching lines @@
   void InvalidateAndResetFrame();
 
   // The owning ScriptInjectionManager.
   ScriptInjectionManager* manager_;
 
   bool should_run_idle_;
 
   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),
       manager_(manager),
       should_run_idle_(true),
       weak_factory_(this) {
 }
 
 ScriptInjectionManager::RFOHelper::~RFOHelper() {
 }
@@ skipping 144 matching lines @@
     : user_script_set_manager_(user_script_set_manager),
       user_script_set_manager_observer_(this) {
   user_script_set_manager_observer_.Add(user_script_set_manager_);
 }
 
 ScriptInjectionManager::~ScriptInjectionManager() {
   for (const auto& injection : pending_injections_)
     injection->invalidate_render_frame();
   for (const auto& injection : running_injections_)
     injection->invalidate_render_frame();
+  DCHECK(injection_watchers_.empty());
 }
 
 void ScriptInjectionManager::OnRenderFrameCreated(
     content::RenderFrame* render_frame) {
   rfo_helpers_.push_back(base::WrapUnique(new RFOHelper(render_frame, this)));
 }
 
 void ScriptInjectionManager::OnExtensionUnloaded(
     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;
     }
   }
 }
 
 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)
       iter = pending_injections_.erase(iter);
     else
       ++iter;
   }
 
+  for (auto watcher : injection_watchers_) {
+    watcher->InvalidateIfFrameEquals(frame);
+  }
+
   frame_statuses_.erase(frame);
 }
 
 void ScriptInjectionManager::StartInjectScripts(
     content::RenderFrame* frame,
     UserScript::RunLocation run_location) {
   FrameStatusMap::iterator iter = frame_statuses_.find(frame);
   // We also don't execute if we detect that the run location is somehow out of
   // order. This can happen if:
   // - The first run location reported for the frame isn't DOCUMENT_START, or
@@ skipping 38 matching lines @@
     } else {
       ++iter;
     }
   }
 
   // Add any injections for user scripts.
   int tab_id = ExtensionFrameHelper::Get(frame)->tab_id();
   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);
+
+    // 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();
   }
 
-  // We are done running in the frame.
-  active_injection_frames_.erase(frame);
-
   scripts_run_info.LogRun(activity_logging_enabled_);
 }
 
 void ScriptInjectionManager::TryToInject(
     std::unique_ptr<ScriptInjection> injection,
     UserScript::RunLocation run_location,
     ScriptsRunInfo* scripts_run_info) {
   // Try to inject the script. If the injection is waiting (i.e., for
   // permission), add it to the list of pending injections. If the injection
   // has blocked, add it to the list of running injections.
   // The Unretained below is safe because this object owns all the
   // ScriptInjections, so is guaranteed to outlive them.
   switch (injection->TryToInject(
       run_location,
       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