Fix focus bug with shell dialogs

chrome/browser/views/shell_dialogs_win.cc

 // Copyright (c) 2006-2008 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 "chrome/browser/shell_dialogs.h"
 
 #include <windows.h>
 #include <commdlg.h>
 #include <shlobj.h>
 
 #include <algorithm>
 #include <set>
 
 #include "app/gfx/font.h"
 #include "app/l10n_util.h"
 #include "app/win_util.h"
 #include "base/file_util.h"
 #include "base/registry.h"
 #include "base/scoped_comptr_win.h"
 #include "base/string_util.h"
 #include "base/thread.h"
 #include "chrome/browser/browser_process.h"
 #include "grit/generated_resources.h"
+#include "views/focus/focus_manager.h"
+#include "views/focus/view_storage.h"
 
 // Helpers to show certain types of Windows shell dialogs in a way that doesn't
 // block the UI of the entire app.
 
 class ShellDialogThread : public base::Thread {
  public:
   ShellDialogThread() : base::Thread("Chrome_ShellDialogThread") { }
 
  protected:
   void Init() {
@@ skipping 75 matching lines @@
   // boxes active, this list will consist of the two browser windows' HWNDs.
   // The derived class must call EndRun once the dialog is done showing to
   // remove the owning HWND from this list.
   // This object is static since it is maintained for all instances of this
   // object - i.e. you can't have a font picker and a file picker open for the
   // same owner, even though they might be represented by different instances
   // of this object.
   // This set only contains non-null HWNDs. NULL hwnds are not added to this
   // list.
   typedef std::set<HWND> Owners;
+
+  // Storage id used to store the last focused view so we can restore focus
+  // appropriately.
+  int view_storage_id_;
+
   static Owners owners_;
   static int instance_count_;
 
   DISALLOW_COPY_AND_ASSIGN(BaseShellDialogImpl);
 };
 
 // static
 BaseShellDialogImpl::Owners BaseShellDialogImpl::owners_;
 int BaseShellDialogImpl::instance_count_ = 0;
 
 BaseShellDialogImpl::BaseShellDialogImpl()
-    : ui_loop_(MessageLoop::current()) {
+    : ui_loop_(MessageLoop::current()),
+      view_storage_id_(views::ViewStorage::GetSharedInstance()->
+          CreateStorageID()) {
   ++instance_count_;
 }
 
 BaseShellDialogImpl::~BaseShellDialogImpl() {
   // All runs should be complete by the time this is called!
   if (--instance_count_ == 0)
     DCHECK(owners_.empty());
 }
 
 BaseShellDialogImpl::RunState BaseShellDialogImpl::BeginRun(HWND owner) {
   // Cannot run a modal shell dialog if one is already running for this owner.
   DCHECK(!IsRunningDialogForOwner(owner));
   // The owner must be a top level window, otherwise we could end up with two
   // entries in our map for the same top level window.
   DCHECK(!owner || owner == GetAncestor(owner, GA_ROOT));
   RunState run_state;
   run_state.dialog_thread = CreateDialogThread();
   run_state.owner = owner;
   if (owner) {
     owners_.insert(owner);
+    // Disabling the owner causes the browser to be disabled when the dialog is
+    // closed and the browser gets the activation.  This messes up the normal
+    // focus restoration process.  To work-around this, we'll restore the focus
+    // ourselves after we have reenabled the owner.
+    views::FocusManager* focus_manager =
+        views::FocusManager::GetFocusManagerForNativeView(owner);
+    if (focus_manager) {
+      views::View* focused_view = focus_manager->GetFocusedView();
+      if (focused_view)
+        views::ViewStorage::GetSharedInstance()->StoreView(view_storage_id_,
+                                                           focused_view);
+    } else {
+      NOTREACHED();
+    }
     DisableOwner(owner);
   }
   return run_state;
 }
 
 void BaseShellDialogImpl::EndRun(RunState run_state) {
   if (run_state.owner) {
     DCHECK(IsRunningDialogForOwner(run_state.owner));
     EnableOwner(run_state.owner);
     DCHECK(owners_.find(run_state.owner) != owners_.end());
     owners_.erase(run_state.owner);
+    // Now that the owner is enabled, restore the focus if applicable.
+    views::View* view_to_focus =
+        views::ViewStorage::GetSharedInstance()->RetrieveView(view_storage_id_);
+    if (view_to_focus) {
+      views::ViewStorage::GetSharedInstance()->RemoveView(view_storage_id_);
+      views::FocusManager* focus_manager =
+          views::FocusManager::GetFocusManagerForNativeView(run_state.owner);
+      if (focus_manager) {
+        // We need to clear focus as when the focus is restored when the dialog
+        // is closed, only setting the native focus fails.  Meaning the focused
+        // manager still believes the right view has focus, and would ignore
+        // requesting focus to what it thinks is already focused.
+        focus_manager->ClearFocus();
+        view_to_focus->RequestFocus();
+      } else {
+        NOTREACHED();
+      }
+    }
   }
   DCHECK(run_state.dialog_thread);
   delete run_state.dialog_thread;
 }
 
 bool BaseShellDialogImpl::IsRunningDialogForOwner(HWND owner) const {
   return (owner && owners_.find(owner) != owners_.end());
 }
 
 void BaseShellDialogImpl::DisableOwner(HWND owner) {
@@ skipping 569 matching lines @@
 void SelectFontDialogImpl::FontNotSelected(void* params, RunState run_state) {
   if (listener_)
     listener_->FontSelectionCanceled(params);
   EndRun(run_state);
 }
 
 // static
 SelectFontDialog* SelectFontDialog::Create(Listener* listener) {
   return new SelectFontDialogImpl(listener);
 }