[Extensions] Handle some funny cases in script injection

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 "base/auto_reset.h"
 #include "base/bind.h"
 #include "base/memory/weak_ptr.h"
 #include "base/values.h"
@@ skipping 67 matching lines @@
                                           int script_id,
                                           const GURL& url);
   virtual void OnPermitScriptInjection(int64 request_id);
 
   // Tells the ScriptInjectionManager to run tasks associated with
   // document_idle.
   void RunIdle();
 
   // Indicate that the frame is no longer valid because it is starting
   // a new load or closing.
-  void InvalidateFrame();
+  void InvalidateAndResetFrame();
 
   // 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) {
 }
 
 ScriptInjectionManager::RFOHelper::~RFOHelper() {
 }
 
 bool ScriptInjectionManager::RFOHelper::OnMessageReceived(
     const IPC::Message& message) {
   bool handled = true;
   IPC_BEGIN_MESSAGE_MAP(ScriptInjectionManager::RFOHelper, message)
     IPC_MESSAGE_HANDLER(ExtensionMsg_ExecuteCode, OnExecuteCode)
     IPC_MESSAGE_HANDLER(ExtensionMsg_PermitScriptInjection,
                         OnPermitScriptInjection)
     IPC_MESSAGE_HANDLER(ExtensionMsg_ExecuteDeclarativeScript,
                         OnExecuteDeclarativeScript)
     IPC_MESSAGE_UNHANDLED(handled = false)
   IPC_END_MESSAGE_MAP()
   return handled;
 }
 
 void ScriptInjectionManager::RFOHelper::DidCreateNewDocument() {
   // A new document is going to be shown, so invalidate the old document state.
   // Check that the frame's state is known before invalidating the frame,
   // because it is possible that a script injection was scheduled before the
   // page was loaded, e.g. by navigating to a javascript: URL before the page
   // has loaded.
   if (manager_->frame_statuses_.count(render_frame()) != 0)
-    InvalidateFrame();
+    InvalidateAndResetFrame();
 }
 
 void ScriptInjectionManager::RFOHelper::DidCreateDocumentElement() {
   manager_->StartInjectScripts(render_frame(), UserScript::DOCUMENT_START);
 }
 
 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
   // subresources' load. We don't want to hold up script injection for a
   // particularly slow subresource, so we set a delayed task from here - but if
   // we finish everything before that point (i.e., DidFinishLoad() is
   // triggered), then there's no reason to keep waiting.
   content::RenderThread::Get()->GetTaskRunner()->PostDelayedTask(
       FROM_HERE,
       base::Bind(&ScriptInjectionManager::RFOHelper::RunIdle,
                  weak_factory_.GetWeakPtr()),
       base::TimeDelta::FromMilliseconds(kScriptIdleTimeoutInMs));
 }
 
 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,
                  weak_factory_.GetWeakPtr()));
 }
 
 void ScriptInjectionManager::RFOHelper::FrameDetached() {
   // The frame is closing - invalidate.
-  InvalidateFrame();
+  InvalidateAndResetFrame();
 }
 
 void ScriptInjectionManager::RFOHelper::OnDestruct() {
   manager_->RemoveObserver(this);
 }
 
 void ScriptInjectionManager::RFOHelper::OnExecuteCode(
     const ExtensionMsg_ExecuteCode_Params& params) {
   manager_->HandleExecuteCode(params, render_frame());
 }
@@ skipping 16 matching lines @@
 }
 
 void ScriptInjectionManager::RFOHelper::OnPermitScriptInjection(
     int64 request_id) {
   manager_->HandlePermitScriptInjection(request_id);
 }
 
 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() {
-  pending_idle_frames_.erase(render_frame());
+void ScriptInjectionManager::RFOHelper::InvalidateAndResetFrame() {
+  // 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());
 }
 
 ScriptInjectionManager::ScriptInjectionManager(
     const ExtensionSet* extensions,
     UserScriptSetManager* user_script_set_manager)
     : extensions_(extensions),
       user_script_set_manager_(user_script_set_manager),
       user_script_set_manager_observer_(this) {
   user_script_set_manager_observer_.Add(user_script_set_manager_);
@@ skipping 48 matching lines @@
        iter != rfo_helpers_.end();
        ++iter) {
     if (*iter == 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 (ScopedVector<ScriptInjection>::iterator iter =
            pending_injections_.begin();
        iter != pending_injections_.end();) {
     if ((*iter)->render_frame() == frame)
       iter = pending_injections_.erase(iter);
     else
       ++iter;
   }
 
   frame_statuses_.erase(frame);
@@ skipping 48 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);
 
-  ScriptsRunInfo scripts_run_info;
+  // Note that we are running in |frame|.
+  active_injection_frames_.insert(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 (!active_injection_frames_.count(frame))
+      break;
     TryToInject(make_scoped_ptr(injection), run_location, &scripts_run_info);
+  }
 
-  scripts_run_info.LogRun(frame->GetWebFrame(), run_location);
+  // We are done running in the frame.
+  active_injection_frames_.erase(frame);
+
+  scripts_run_info.LogRun();
 }
 
 void ScriptInjectionManager::TryToInject(
     scoped_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
@@ skipping 29 matching lines @@
   }
 
   scoped_ptr<ScriptInjection> injection(new ScriptInjection(
       scoped_ptr<ScriptInjector>(
           new ProgrammaticScriptInjector(params, render_frame)),
       render_frame,
       injection_host.Pass(),
       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(
     content::RenderFrame* render_frame,
     int tab_id,
     const ExtensionId& extension_id,
     int script_id,
     const GURL& url) {
   scoped_ptr<ScriptInjection> injection =
       user_script_set_manager_->GetInjectionForDeclarativeScript(
           script_id,
           render_frame,
           tab_id,
           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();
   }
 }
 
 void ScriptInjectionManager::HandlePermitScriptInjection(int64 request_id) {
   ScopedVector<ScriptInjection>::iterator 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's DidStartProvisionalLoad callback would have caused it to be
   // cleared out).
 
   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