Made session restore block about:flags driven browser restart.

chrome/browser/chromeos/login/login_utils.cc

Index: chrome/browser/chromeos/login/login_utils.cc
diff --git a/chrome/browser/chromeos/login/login_utils.cc b/chrome/browser/chromeos/login/login_utils.cc
index da1a4b5f59eb2905dd15c28d6e8e3fe2eb68026f..144b614c3e5aa8002a99f332f8081f1e7fa93b93 100644
--- a/chrome/browser/chromeos/login/login_utils.cc
+++ b/chrome/browser/chromeos/login/login_utils.cc
@@ -105,6 +105,7 @@ class LoginUtilsImpl
         has_web_auth_cookies_(false),
         delegate_(NULL),
         should_restore_auth_session_(false),
+        exit_after_session_restore_(false),
         session_restore_strategy_(
             OAuth2LoginManager::RESTORE_FROM_SAVED_OAUTH2_REFRESH_TOKEN) {
     net::NetworkChangeNotifier::AddConnectionTypeObserver(this);
@@ -132,10 +133,11 @@ class LoginUtilsImpl
   virtual void RestoreAuthenticationSession(Profile* profile) OVERRIDE;
   virtual void InitRlzDelayed(Profile* user_profile) OVERRIDE;
 
-  // OAuth2LoginManager::Delegate overrides.
+  // OAuth2LoginManager::Observer overrides.
   virtual void OnSessionRestoreStateChanged(
       Profile* user_profile,
       OAuth2LoginManager::SessionRestoreState state) OVERRIDE;
+  virtual void OnNewRefreshTokenAvaiable(Profile* user_profile) OVERRIDE;
   virtual void OnSessionAuthenticated(Profile* user_profile) OVERRIDE;
 
   // net::NetworkChangeNotifier::ConnectionTypeObserver overrides.
@@ -183,6 +185,10 @@ class LoginUtilsImpl
   // Starts signing related services. Initiates TokenService token retrieval.
   void StartSignedInServices(Profile* profile);
 
+  // Attempts exiting browser process and esures this does not happen
+  // while we are still fetching new OAuth refresh tokens.
+  void AttemptExit(Profile* profile);
+
   UserContext user_context_;
   bool using_oauth_;
 
@@ -199,6 +205,9 @@ class LoginUtilsImpl
   // online state change.
   bool should_restore_auth_session_;
 
+  // True if we should restart chrome right after session restore.
+  bool exit_after_session_restore_;
+
   // Sesion restore strategy.
   OAuth2LoginManager::SessionRestoreStrategy session_restore_strategy_;
   // OAuth2 refresh token for session restore.
@@ -261,7 +270,7 @@ void LoginUtilsImpl::DoBrowserLaunch(Profile* profile,
     VLOG(1) << "Restarting to apply per-session flags...";
     DBusThreadManager::Get()->GetSessionManagerClient()->SetFlagsForUser(
         UserManager::Get()->GetActiveUser()->email(), flags);
-    chrome::AttemptExit();
+    AttemptExit(profile);
     return;
   }
 
@@ -476,6 +485,7 @@ void LoginUtilsImpl::RestoreAuthSession(Profile* user_profile,
           chromeos::switches::kOobeSkipPostLogin))
     return;
 
+  exit_after_session_restore_ = false;
   // Remove legacy OAuth1 token if we have one. If it's valid, we should already
   // have OAuth2 refresh token in TokenService that could be used to retrieve
   // all other tokens and user_context.
@@ -707,28 +717,49 @@ void LoginUtilsImpl::RestoreAuthenticationSession(Profile* user_profile) {
 void LoginUtilsImpl::OnSessionRestoreStateChanged(
     Profile* user_profile,
     OAuth2LoginManager::SessionRestoreState state) {
+  User::OAuthTokenStatus user_status = User::OAUTH_TOKEN_STATUS_UNKNOWN;
   OAuth2LoginManager* login_manager =
       OAuth2LoginManagerFactory::GetInstance()->GetForProfile(user_profile);
   switch (state) {
-    case OAuth2LoginManager::SESSION_RESTORE_NOT_STARTED:
-      break;
-    case OAuth2LoginManager::SESSION_RESTORE_PREPARING:
-      break;
-    case OAuth2LoginManager::SESSION_RESTORE_IN_PROGRESS:
-      break;
     case OAuth2LoginManager::SESSION_RESTORE_DONE:
-      UserManager::Get()->SaveUserOAuthStatus(
-          UserManager::Get()->GetLoggedInUser()->email(),
-          User::OAUTH2_TOKEN_STATUS_VALID);
-      login_manager->RemoveObserver(this);
+      user_status = User::OAUTH2_TOKEN_STATUS_VALID;
       break;
     case OAuth2LoginManager::SESSION_RESTORE_FAILED:
-      UserManager::Get()->SaveUserOAuthStatus(
-          UserManager::Get()->GetLoggedInUser()->email(),
-          User::OAUTH2_TOKEN_STATUS_INVALID);
-      login_manager->RemoveObserver(this);
+      user_status = User::OAUTH2_TOKEN_STATUS_INVALID;
       break;
+    case OAuth2LoginManager::SESSION_RESTORE_NOT_STARTED:
+      return;
+    case OAuth2LoginManager::SESSION_RESTORE_PREPARING:
+      return;
+    case OAuth2LoginManager::SESSION_RESTORE_IN_PROGRESS:
+      return;
   }
+
+  // We are in one of "done" states here.
+  UserManager::Get()->SaveUserOAuthStatus(
+      UserManager::Get()->GetLoggedInUser()->email(),
+      user_status);
+  login_manager->RemoveObserver(this);
+}
+
+void LoginUtilsImpl::OnNewRefreshTokenAvaiable(Profile* user_profile) {
+  // Check if we were waiting to restart chrome.
+  if (!exit_after_session_restore_)
+    return;
+
+  OAuth2LoginManager* login_manager =
+      OAuth2LoginManagerFactory::GetInstance()->GetForProfile(user_profile);
+  login_manager->RemoveObserver(this);
+
+  // Mark user auth token status as valid.
+  UserManager::Get()->SaveUserOAuthStatus(
+      UserManager::Get()->GetLoggedInUser()->email(),
+      User::OAUTH2_TOKEN_STATUS_VALID);
+
+  LOG(WARNING) << "Exiting after new refresh token fetched";
+  // We need to exit cleanly in this case to make sure OAuth2 RT is actually
+  // saved.
+  chrome::ExitCleanly();
 }
 
 void LoginUtilsImpl::OnSessionAuthenticated(Profile* user_profile) {
@@ -755,6 +786,29 @@ void LoginUtilsImpl::OnConnectionTypeChanged(
   }
 }
 
+void LoginUtilsImpl::AttemptExit(Profile* profile) {
+  if (session_restore_strategy_ !=
+      OAuth2LoginManager::RESTORE_FROM_COOKIE_JAR) {
+    chrome::AttemptExit();
+    return;
+  }
+
+  // We can't really quit if the session restore process that mints new
+  // refresh token is still in progress.
+  OAuth2LoginManager* login_manager =
+      OAuth2LoginManagerFactory::GetInstance()->GetForProfile(profile);
+  if (login_manager->state() !=
+          OAuth2LoginManager::SESSION_RESTORE_PREPARING &&
+      login_manager->state() !=
+          OAuth2LoginManager::SESSION_RESTORE_IN_PROGRESS) {
+    chrome::AttemptExit();
+    return;
+  }
+
+  LOG(WARNING) << "Attempting browser restart during session restore.";
+  exit_after_session_restore_ = true;
+}
+
 // static
 void LoginUtils::RegisterPrefs(PrefRegistrySimple* registry) {
   registry->RegisterBooleanPref(prefs::kFactoryResetRequested, false);