[Extensions] Handle some funny cases in script injection

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 326452b8fad2de0a6aa183c2d2434fa0cd083b5d..fb1ae01ce8f328df110db6634f6f1729261a9665 100644
--- a/extensions/renderer/script_injection_manager.cc
+++ b/extensions/renderer/script_injection_manager.cc
@@ -90,19 +90,16 @@ class ScriptInjectionManager::RFOHelper : public content::RenderFrameObserver {
   // The owning ScriptInjectionManager.
   ScriptInjectionManager* manager_;
 
-  // The set of frames that we are about to notify for DOCUMENT_IDLE. We keep
-  // a set of those that are valid, so we don't notify that an invalid frame
-  // became idle.
-  std::set<content::RenderFrame*> pending_idle_frames_;
+  bool should_run_idle_;
 
   base::WeakPtrFactory<RFOHelper> weak_factory_;
 };
 
-ScriptInjectionManager::RFOHelper::RFOHelper(
-    content::RenderFrame* render_frame,
-    ScriptInjectionManager* manager)
+ScriptInjectionManager::RFOHelper::RFOHelper(content::RenderFrame* render_frame,
+                                             ScriptInjectionManager* manager)
     : content::RenderFrameObserver(render_frame),
       manager_(manager),
+      should_run_idle_(true),
       weak_factory_(this) {
 }
 
@@ -139,7 +136,6 @@ void ScriptInjectionManager::RFOHelper::DidCreateDocumentElement() {
 
 void ScriptInjectionManager::RFOHelper::DidFinishDocumentLoad() {
   manager_->StartInjectScripts(render_frame(), UserScript::DOCUMENT_END);
-  pending_idle_frames_.insert(render_frame());
   // We try to run idle in two places: here and DidFinishLoad.
   // DidFinishDocumentLoad() corresponds to completing the document's load,
   // whereas DidFinishLoad corresponds to completing the document and all
@@ -156,10 +152,6 @@ void ScriptInjectionManager::RFOHelper::DidFinishDocumentLoad() {
 
 void ScriptInjectionManager::RFOHelper::DidFinishLoad() {
   // Ensure that we don't block any UI progress by running scripts.
-  // We *don't* add the frame to |pending_idle_frames_| here because
-  // DidFinishDocumentLoad should strictly come before DidFinishLoad, so the
-  // first posted task to RunIdle() pops it out of the set. This ensures we
-  // don't try to run idle twice.
   content::RenderThread::Get()->GetTaskRunner()->PostTask(
       FROM_HERE,
       base::Bind(&ScriptInjectionManager::RFOHelper::RunIdle,
@@ -205,14 +197,18 @@ void ScriptInjectionManager::RFOHelper::OnPermitScriptInjection(
 void ScriptInjectionManager::RFOHelper::RunIdle() {
   // Only notify the manager if the frame hasn't either been removed or already
   // had idle run since the task to RunIdle() was posted.
-  if (pending_idle_frames_.count(render_frame()) > 0) {
+  if (should_run_idle_) {
+    should_run_idle_ = false;
     manager_->StartInjectScripts(render_frame(), UserScript::DOCUMENT_IDLE);
-    pending_idle_frames_.erase(render_frame());
   }
 }
 
 void ScriptInjectionManager::RFOHelper::InvalidateFrame() {

[not at google - send to devlin, 2015/06/29 20:09:07]

Could you rename this "ResetFrame"? Then you wouldn't even really need these
comments.

[Devlin, 2015/06/29 20:43:13]

Made it InvalidateAndResetFrame() (I like the invalidate, since it hints at
invalidating it in the ScriptInjectionManager).

-  pending_idle_frames_.erase(render_frame());
+  // Invalidate any pending idle injections, and reset the frame inject on idle.
+  weak_factory_.InvalidateWeakPtrs();
+  // We reset to inject on idle, because the frame can be reused (in the case of
+  // navigation).
+  should_run_idle_ = true;
   manager_->InvalidateForFrame(render_frame());
 }
 
@@ -220,6 +216,7 @@ ScriptInjectionManager::ScriptInjectionManager(
     const ExtensionSet* extensions,
     UserScriptSetManager* user_script_set_manager)
     : extensions_(extensions),
+      valid_running_injection_frame_(nullptr),
       user_script_set_manager_(user_script_set_manager),
       user_script_set_manager_observer_(this) {
   user_script_set_manager_observer_.Add(user_script_set_manager_);
@@ -281,6 +278,11 @@ 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.
+  if (frame == valid_running_injection_frame_)
+    valid_running_injection_frame_ = nullptr;
+
   for (ScopedVector<ScriptInjection>::iterator iter =
            pending_injections_.begin();
        iter != pending_injections_.end();) {
@@ -349,13 +351,26 @@ void ScriptInjectionManager::InjectScripts(
   user_script_set_manager_->GetAllInjections(
       &frame_injections, frame, tab_id, run_location);
 
-  ScriptsRunInfo scripts_run_info;
+  // Note that we are running in |frame|. We should only ever be running in a
+  // single frame at a time.
+  DCHECK(!valid_running_injection_frame_);
+  valid_running_injection_frame_ = frame;
+
+  ScriptsRunInfo scripts_run_info(frame, run_location);
   std::vector<ScriptInjection*> released_injections;
   frame_injections.release(&released_injections);
-  for (ScriptInjection* injection : released_injections)
+  for (ScriptInjection* injection : released_injections) {
+    // 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 (!valid_running_injection_frame_)
+      break;
     TryToInject(make_scoped_ptr(injection), run_location, &scripts_run_info);
+  }
+
+  // We are done running in the frame.
+  valid_running_injection_frame_ = nullptr;
 
-  scripts_run_info.LogRun(frame->GetWebFrame(), run_location);
+  scripts_run_info.LogRun();
 }
 
 void ScriptInjectionManager::TryToInject(
@@ -405,13 +420,12 @@ void ScriptInjectionManager::HandleExecuteCode(
       static_cast<UserScript::RunLocation>(params.run_at),
       ExtensionFrameHelper::Get(render_frame)->tab_id()));
 
-  ScriptsRunInfo scripts_run_info;
   FrameStatusMap::const_iterator iter = frame_statuses_.find(render_frame);
+  UserScript::RunLocation run_location =
+      iter == frame_statuses_.end() ? UserScript::UNDEFINED : iter->second;
 
-  TryToInject(
-      injection.Pass(),
-      iter == frame_statuses_.end() ? UserScript::UNDEFINED : iter->second,
-      &scripts_run_info);
+  ScriptsRunInfo scripts_run_info(render_frame, run_location);
+  TryToInject(injection.Pass(), run_location, &scripts_run_info);
 }
 
 void ScriptInjectionManager::HandleExecuteDeclarativeScript(
@@ -428,14 +442,13 @@ void ScriptInjectionManager::HandleExecuteDeclarativeScript(
           url,
           extension_id);
   if (injection.get()) {
-    ScriptsRunInfo scripts_run_info;
+    ScriptsRunInfo scripts_run_info(render_frame, UserScript::BROWSER_DRIVEN);
     // TODO(markdittmer): Use return value of TryToInject for error handling.
     TryToInject(injection.Pass(),
                 UserScript::BROWSER_DRIVEN,
                 &scripts_run_info);
 
-    scripts_run_info.LogRun(render_frame->GetWebFrame(),
-                            UserScript::BROWSER_DRIVEN);
+    scripts_run_info.LogRun();
   }
 }
 
@@ -459,13 +472,13 @@ void ScriptInjectionManager::HandlePermitScriptInjection(int64 request_id) {
   scoped_ptr<ScriptInjection> injection(*iter);
   pending_injections_.weak_erase(iter);
 
-  ScriptsRunInfo scripts_run_info;
+  ScriptsRunInfo scripts_run_info(injection->render_frame(),
+                                  UserScript::RUN_DEFERRED);
   ScriptInjection::InjectionResult res = injection->OnPermissionGranted(
       &scripts_run_info);
   if (res == ScriptInjection::INJECTION_BLOCKED)
     running_injections_.push_back(injection.Pass());
-  scripts_run_info.LogRun(injection->render_frame()->GetWebFrame(),
-                          UserScript::RUN_DEFERRED);
+  scripts_run_info.LogRun();
 }
 
 }  // namespace extensions