cros: Transfer auth cookies for SAML webview sign-in.

- Update ProfileAuthData to transfer cookie from sign-in webview;
- Passing back usingSAML flag for auth code authentication and use
  it to set has_auth_cookie flag when SAML is used;
- Update UserCloudPolicyManagerChromeOS to fetch initial policy from
  correct auth request context;
- Enable SAMLPolicyTest for WebView;

BUG=461505

Review URL: https://codereview.chromium.org/1021383002

Cr-Commit-Position: refs/heads/master@{#322223}

17 files changed, 213 insertions, 67 deletions

diff --git a/chrome/browser/chromeos/login/existing_user_controller.cc b/chrome/browser/chromeos/login/existing_user_controller.cc
index aa9615b..db71d44 100644
--- a/chrome/browser/chromeos/login/existing_user_controller.cc
+++ b/chrome/browser/chromeos/login/existing_user_controller.cc
@@ -589,9 +589,15 @@ void ExistingUserController::OnAuthSuccess(const UserContext& user_context) {
 
   StopPublicSessionAutoLoginTimer();
 
+  // Truth table of |has_auth_cookies|:
+  //                          Regular        SAML
+  //  /ServiceLogin              T            T
+  //  /ChromeOsEmbeddedSetup     F            T
+  //  Bootstrap experiment       F            N/A
   const bool has_auth_cookies =
       login_performer_->auth_mode() == LoginPerformer::AUTH_MODE_EXTENSION &&
-      user_context.GetAuthCode().empty() &&
+      (user_context.GetAuthCode().empty() ||
+       user_context.GetAuthFlow() == UserContext::AUTH_FLOW_GAIA_WITH_SAML) &&
       user_context.GetAuthFlow() != UserContext::AUTH_FLOW_EASY_BOOTSTRAP;
 
   // LoginPerformer instance will delete itself in case of successful auth.
diff --git a/chrome/browser/chromeos/login/helper.cc b/chrome/browser/chromeos/login/helper.cc
index efb9890..692c831 100644
--- a/chrome/browser/chromeos/login/helper.cc
+++ b/chrome/browser/chromeos/login/helper.cc
@@ -4,13 +4,22 @@
 
 #include "chrome/browser/chromeos/login/helper.h"
 
+#include "base/bind.h"
 #include "base/command_line.h"
 #include "base/strings/utf_string_conversions.h"
+#include "chrome/browser/chromeos/login/startup_utils.h"
+#include "chrome/browser/chromeos/login/ui/login_display_host_impl.h"
+#include "chrome/browser/chromeos/login/ui/webui_login_view.h"
+#include "chrome/browser/chromeos/profiles/profile_helper.h"
 #include "chrome/grit/generated_resources.h"
 #include "chromeos/chromeos_switches.h"
 #include "chromeos/network/network_handler.h"
 #include "chromeos/network/network_state.h"
 #include "chromeos/network/network_state_handler.h"
+#include "content/public/browser/storage_partition.h"
+#include "content/public/browser/web_contents.h"
+#include "extensions/browser/guest_view/guest_view_manager.h"
+#include "extensions/browser/guest_view/web_view/web_view_guest.h"
 #include "third_party/cros_system_api/dbus/service_constants.h"
 #include "ui/base/l10n/l10n_util.h"
 #include "ui/gfx/image/image_skia.h"
@@ -18,6 +27,57 @@
 
 namespace chromeos {
 
+namespace {
+
+// Gets the WebContents instance of current login display. If there is none,
+// returns nullptr.
+content::WebContents* GetLoginWebContents() {
+  LoginDisplayHost* host = LoginDisplayHostImpl::default_host();
+  if (!host || !host->GetWebUILoginView())
+    return nullptr;
+
+  return host->GetWebUILoginView()->GetWebContents();
+}
+
+// Callback used by GetPartition below to return the first guest contents with a
+// matching partition name.
+bool FindGuestByPartitionName(const std::string& partition_name,
+                              content::WebContents** out_guest_contents,
+                              content::WebContents* guest_contents) {
+  std::string domain;
+  std::string name;
+  bool in_memory;
+  extensions::WebViewGuest::GetGuestPartitionConfigForSite(
+      guest_contents->GetSiteInstance()->GetSiteURL(), &domain, &name,
+      &in_memory);
+  if (partition_name != name)
+    return false;
+
+  *out_guest_contents = guest_contents;
+  return true;
+}
+
+// Gets the storage partition of guest contents of a given embedder.
+// If a name is given, returns the partition associated with the name.
+// Otherwise, returns the default shared in-memory partition. Returns nullptr if
+// a matching partition could not be found.
+content::StoragePartition* GetPartition(content::WebContents* embedder,
+                                        const std::string& partition_name) {
+  extensions::GuestViewManager* manager =
+      extensions::GuestViewManager::FromBrowserContext(
+          embedder->GetBrowserContext());
+  content::WebContents* guest_contents = nullptr;
+  manager->ForEachGuest(embedder, base::Bind(&FindGuestByPartitionName,
+                                             partition_name, &guest_contents));
+
+  return guest_contents ? content::BrowserContext::GetStoragePartition(
+                              guest_contents->GetBrowserContext(),
+                              guest_contents->GetSiteInstance())
+                        : nullptr;
+}
+
+}  // namespace
+
 gfx::Rect CalculateScreenBounds(const gfx::Size& size) {
   gfx::Rect bounds =
       gfx::Screen::GetNativeScreen()->GetPrimaryDisplay().bounds();
@@ -82,6 +142,19 @@ bool NetworkStateHelper::IsConnecting() const {
       chromeos::NetworkTypePattern::Default()) != NULL;
 }
 
+content::StoragePartition* GetSigninPartition() {
+  // Note the partition name must match the sign-in webview used. For now,
+  // this is the default unnamed, shared, in-memory partition.
+  return GetPartition(GetLoginWebContents(), std::string());
+}
+
+net::URLRequestContextGetter* GetSigninContext() {
+  if (StartupUtils::IsWebviewSigninEnabled())
+    return GetSigninPartition()->GetURLRequestContext();
+
+  return ProfileHelper::GetSigninProfile()->GetRequestContext();
+}
+
 }  // namespace login
 
 }  // namespace chromeos
diff --git a/chrome/browser/chromeos/login/helper.h b/chrome/browser/chromeos/login/helper.h
index ef4e870..2809d66 100644
--- a/chrome/browser/chromeos/login/helper.h
+++ b/chrome/browser/chromeos/login/helper.h
@@ -21,6 +21,14 @@ class Rect;
 class Size;
 }  // namespace gfx
 
+namespace content {
+class StoragePartition;
+}
+
+namespace net {
+class URLRequestContextGetter;
+}
+
 namespace chromeos {
 
 // Returns bounds of the screen to use for login wizard.
@@ -66,6 +74,20 @@ class NetworkStateHelper {
   DISALLOW_COPY_AND_ASSIGN(NetworkStateHelper);
 };
 
+//
+// Webview based login helpers.
+//
+
+// Returns the storage partition for the sign-in webview. Note the function gets
+// the partition via the sign-in WebContents thus returns nullptr if the sign-in
+// webui is torn down.
+content::StoragePartition* GetSigninPartition();
+
+// Returns the request context that contains sign-in cookies. For old iframe
+// based flow, the context of the sign-in profile is returned. For webview based
+// flow, the context of the sign-in webview storage partition is returned.
+net::URLRequestContextGetter* GetSigninContext();
+
 }  // namespace login
 
 }  // namespace chromeos
diff --git a/chrome/browser/chromeos/login/profile_auth_data.cc b/chrome/browser/chromeos/login/profile_auth_data.cc
index d75c9d0..7d0e9ef 100644
--- a/chrome/browser/chromeos/login/profile_auth_data.cc
+++ b/chrome/browser/chromeos/login/profile_auth_data.cc
@@ -40,8 +40,8 @@ const char kSAMLEndCookie[] = "google-accounts-saml-end";
 class ProfileAuthDataTransferer {
  public:
   ProfileAuthDataTransferer(
-      content::BrowserContext* from_context,
-      content::BrowserContext* to_context,
+      net::URLRequestContextGetter* from_context,
+      net::URLRequestContextGetter* to_context,
       bool transfer_auth_cookies_and_channel_ids_on_first_login,
       bool transfer_saml_auth_cookies_on_subsequent_login,
       const base::Closure& completion_callback);
@@ -119,13 +119,13 @@ class ProfileAuthDataTransferer {
 };
 
 ProfileAuthDataTransferer::ProfileAuthDataTransferer(
-    content::BrowserContext* from_context,
-    content::BrowserContext* to_context,
+    net::URLRequestContextGetter* from_context,
+    net::URLRequestContextGetter* to_context,
     bool transfer_auth_cookies_and_channel_ids_on_first_login,
     bool transfer_saml_auth_cookies_on_subsequent_login,
     const base::Closure& completion_callback)
-    : from_context_(from_context->GetRequestContext()),
-      to_context_(to_context->GetRequestContext()),
+    : from_context_(from_context),
+      to_context_(to_context),
       transfer_auth_cookies_and_channel_ids_on_first_login_(
           transfer_auth_cookies_and_channel_ids_on_first_login),
       transfer_saml_auth_cookies_on_subsequent_login_(
@@ -314,8 +314,8 @@ void ProfileAuthDataTransferer::Finish() {
 }  // namespace
 
 void ProfileAuthData::Transfer(
-    content::BrowserContext* from_context,
-    content::BrowserContext* to_context,
+    net::URLRequestContextGetter* from_context,
+    net::URLRequestContextGetter* to_context,
     bool transfer_auth_cookies_and_channel_ids_on_first_login,
     bool transfer_saml_auth_cookies_on_subsequent_login,
     const base::Closure& completion_callback) {
diff --git a/chrome/browser/chromeos/login/profile_auth_data.h b/chrome/browser/chromeos/login/profile_auth_data.h
index 59a35a2..5984feb 100644
--- a/chrome/browser/chromeos/login/profile_auth_data.h
+++ b/chrome/browser/chromeos/login/profile_auth_data.h
@@ -8,8 +8,8 @@
 #include "base/callback_forward.h"
 #include "base/macros.h"
 
-namespace content {
-class BrowserContext;
+namespace net {
+class URLRequestContextGetter;
 }
 
 namespace chromeos {
@@ -31,8 +31,8 @@ class ProfileAuthData {
   //   |transfer_saml_auth_cookies_on_subsequent_login| is true and
   //   |to_context|'s cookie jar is not empty.
   static void Transfer(
-      content::BrowserContext* from_context,
-      content::BrowserContext* to_context,
+      net::URLRequestContextGetter* from_context,
+      net::URLRequestContextGetter* to_context,
       bool transfer_auth_cookies_and_channel_ids_on_first_login,
       bool transfer_saml_auth_cookies_on_subsequent_login,
       const base::Closure& completion_callback);
diff --git a/chrome/browser/chromeos/login/profile_auth_data_unittest.cc b/chrome/browser/chromeos/login/profile_auth_data_unittest.cc
index bf6bc33..d2385fa 100644
--- a/chrome/browser/chromeos/login/profile_auth_data_unittest.cc
+++ b/chrome/browser/chromeos/login/profile_auth_data_unittest.cc
@@ -128,8 +128,8 @@ void ProfileAuthDataTest::Transfer(
     bool transfer_saml_auth_cookies_on_subsequent_login) {
   base::RunLoop run_loop;
   ProfileAuthData::Transfer(
-      &login_browser_context_,
-      &user_browser_context_,
+      login_browser_context_.GetRequestContext(),
+      user_browser_context_.GetRequestContext(),
       transfer_auth_cookies_and_channel_ids_on_first_login,
       transfer_saml_auth_cookies_on_subsequent_login,
       run_loop.QuitClosure());
diff --git a/chrome/browser/chromeos/login/saml/saml_browsertest.cc b/chrome/browser/chromeos/login/saml/saml_browsertest.cc
index 047a7aa..9bfde3e 100644
--- a/chrome/browser/chromeos/login/saml/saml_browsertest.cc
+++ b/chrome/browser/chromeos/login/saml/saml_browsertest.cc
@@ -293,23 +293,28 @@ class SamlTest : public InProcessBrowserTest,
 
   bool SetUpUserDataDirectory() override {
     if (UseWebView()) {
-      // Enable webview signin.
+      // Fake Dev channel to enable webview signin.
       scoped_channel_.reset(new extensions::ScopedCurrentChannel(
           chrome::VersionInfo::CHANNEL_DEV));
 
-      base::DictionaryValue local_state_dict;
-      local_state_dict.SetBoolean(prefs::kWebviewSigninEnabled, true);
-      // OobeCompleted to skip controller-pairing-screen which still uses
-      // iframe and ends up in a JS error in oobe page init.
-      // See http://crbug.com/467147
-      local_state_dict.SetBoolean(prefs::kOobeComplete, true);
-
       base::FilePath user_data_dir;
       CHECK(PathService::Get(chrome::DIR_USER_DATA, &user_data_dir));
       base::FilePath local_state_path =
           user_data_dir.Append(chrome::kLocalStateFilename);
-      CHECK(JSONFileValueSerializer(local_state_path)
-                .Serialize(local_state_dict));
+
+      // Set webview enabled flag only when local state file does not exist.
+      // Otherwise, we break PRE tests that leave state in it.
+      if (!base::PathExists(local_state_path)) {
+        base::DictionaryValue local_state_dict;
+        local_state_dict.SetBoolean(prefs::kWebviewSigninEnabled, true);
+        // OobeCompleted to skip controller-pairing-screen which still uses
+        // iframe and ends up in a JS error in oobe page init.
+        // See http://crbug.com/467147
+        local_state_dict.SetBoolean(prefs::kOobeComplete, true);
+
+        CHECK(JSONFileValueSerializer(local_state_path)
+                  .Serialize(local_state_dict));
+      }
     }
 
     return InProcessBrowserTest::SetUpUserDataDirectory();
@@ -1250,6 +1255,6 @@ IN_PROC_BROWSER_TEST_P(SAMLPolicyTest, TransferCookiesUnaffiliated) {
 // TODO(xiyuan): Update once cookies are properly handled.
 INSTANTIATE_TEST_CASE_P(SamlSuite,
                         SAMLPolicyTest,
-                        testing::Values(false));
+                        testing::Bool());
 
 }  // namespace chromeos
diff --git a/chrome/browser/chromeos/login/session/user_session_manager.cc b/chrome/browser/chromeos/login/session/user_session_manager.cc
index 6764105..b595ad8 100644
--- a/chrome/browser/chromeos/login/session/user_session_manager.cc
+++ b/chrome/browser/chromeos/login/session/user_session_manager.cc
@@ -34,6 +34,7 @@
 #include "chrome/browser/chromeos/login/chrome_restart_request.h"
 #include "chrome/browser/chromeos/login/demo_mode/demo_app_launcher.h"
 #include "chrome/browser/chromeos/login/easy_unlock/easy_unlock_key_manager.h"
+#include "chrome/browser/chromeos/login/helper.h"
 #include "chrome/browser/chromeos/login/lock/screen_locker.h"
 #include "chrome/browser/chromeos/login/profile_auth_data.h"
 #include "chrome/browser/chromeos/login/saml/saml_offline_signin_limiter.h"
@@ -72,7 +73,6 @@
 #include "chrome/common/chrome_switches.h"
 #include "chrome/common/logging_chrome.h"
 #include "chrome/common/pref_names.h"
-#include "chrome/common/url_constants.h"
 #include "chromeos/cert_loader.h"
 #include "chromeos/chromeos_switches.h"
 #include "chromeos/cryptohome/cryptohome_util.h"
@@ -974,8 +974,8 @@ void UserSessionManager::UserProfileInitialized(Profile* profile,
     const bool transfer_auth_cookies_and_channel_ids_on_first_login =
         has_auth_cookies_;
     ProfileAuthData::Transfer(
-        authenticator_->authentication_context(),
-        profile,
+        GetAuthRequestContext(),
+        profile->GetRequestContext(),
         transfer_auth_cookies_and_channel_ids_on_first_login,
         transfer_saml_auth_cookies_on_subsequent_login,
         base::Bind(&UserSessionManager::CompleteProfileCreateAfterAuthTransfer,
@@ -1204,27 +1204,10 @@ void UserSessionManager::RestoreAuthSessionImpl(
   OAuth2LoginManager* login_manager =
       OAuth2LoginManagerFactory::GetInstance()->GetForProfile(profile);
   login_manager->AddObserver(this);
-  net::URLRequestContextGetter* auth_request_context = NULL;
-
-  if (StartupUtils::IsWebviewSigninEnabled()) {
-    // Webview uses different partition storage than iframe. We need to get
-    // cookies from the right storage for url request to get auth token into
-    // session.
-    GURL oobe_url(chrome::kChromeUIOobeURL);
-    GURL guest_url(std::string(content::kGuestScheme) +
-                   url::kStandardSchemeSeparator + oobe_url.GetContent());
-    content::StoragePartition* partition =
-        content::BrowserContext::GetStoragePartitionForSite(
-            ProfileHelper::GetSigninProfile(), guest_url);
-    auth_request_context = partition->GetURLRequestContext();
-  } else if (authenticator_.get() && authenticator_->authentication_context()) {
-    auth_request_context =
-        authenticator_->authentication_context()->GetRequestContext();
-  }
 
-  login_manager->RestoreSession(auth_request_context, session_restore_strategy_,
-                                user_context_.GetRefreshToken(),
-                                user_context_.GetAuthCode());
+  login_manager->RestoreSession(
+      GetAuthRequestContext(), session_restore_strategy_,
+      user_context_.GetRefreshToken(), user_context_.GetAuthCode());
 }
 
 void UserSessionManager::InitRlzImpl(Profile* profile, bool disabled) {
@@ -1411,6 +1394,22 @@ void UserSessionManager::UpdateEasyUnlockKeys(const UserContext& user_context) {
                  user_context.GetUserID()));
 }
 
+net::URLRequestContextGetter*
+UserSessionManager::GetAuthRequestContext() const {
+  net::URLRequestContextGetter* auth_request_context = NULL;
+
+  if (StartupUtils::IsWebviewSigninEnabled()) {
+    // Webview uses different partition storage than iframe. We need to get
+    // cookies from the right storage for url request to get auth token into
+    // session.
+    auth_request_context = login::GetSigninPartition()->GetURLRequestContext();
+  } else if (authenticator_.get() && authenticator_->authentication_context()) {
+    auth_request_context =
+        authenticator_->authentication_context()->GetRequestContext();
+  }
+  return auth_request_context;
+}
+
 void UserSessionManager::AttemptRestart(Profile* profile) {
   if (CheckEasyUnlockKeyOps(base::Bind(&UserSessionManager::AttemptRestart,
                                        AsWeakPtr(), profile))) {
diff --git a/chrome/browser/chromeos/login/session/user_session_manager.h b/chrome/browser/chromeos/login/session/user_session_manager.h
index 3e17160..9874aef 100644
--- a/chrome/browser/chromeos/login/session/user_session_manager.h
+++ b/chrome/browser/chromeos/login/session/user_session_manager.h
@@ -28,6 +28,10 @@ class PrefRegistrySimple;
 class PrefService;
 class Profile;
 
+namespace net {
+class URLRequestContextGetter;
+}
+
 namespace user_manager {
 class User;
 }  // namespace user_manager
@@ -220,6 +224,9 @@ class UserSessionManager
   // Update Easy unlock cryptohome keys for given user context.
   void UpdateEasyUnlockKeys(const UserContext& user_context);
 
+  // Returns the auth request context associated with auth data.
+  net::URLRequestContextGetter* GetAuthRequestContext() const;
+
   // Removes a profile from the per-user input methods states map.
   void RemoveProfileForTesting(Profile* profile);
 
diff --git a/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos.cc b/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos.cc
index b0a2959..3929680 100644
--- a/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos.cc
+++ b/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos.cc
@@ -12,10 +12,10 @@
 #include "base/sequenced_task_runner.h"
 #include "base/values.h"
 #include "chrome/browser/browser_process.h"
+#include "chrome/browser/chromeos/login/helper.h"
 #include "chrome/browser/chromeos/policy/policy_oauth2_token_fetcher.h"
 #include "chrome/browser/chromeos/policy/user_cloud_policy_manager_factory_chromeos.h"
 #include "chrome/browser/chromeos/policy/wildcard_login_checker.h"
-#include "chrome/browser/chromeos/profiles/profile_helper.h"
 #include "chrome/browser/lifetime/application_lifetime.h"
 #include "chrome/common/chrome_content_client.h"
 #include "components/policy/core/common/cloud/cloud_external_data_manager.h"
@@ -212,7 +212,7 @@ void UserCloudPolicyManagerChromeOS::OnInitializationCompleted(
   // access token is already available.
   if (!client()->is_registered()) {
     if (wait_for_policy_fetch_) {
-      FetchPolicyOAuthTokenUsingSigninProfile();
+      FetchPolicyOAuthTokenUsingSigninContext();
     } else if (!access_token_.empty()) {
       OnAccessTokenAvailable(access_token_);
     }
@@ -286,13 +286,11 @@ void UserCloudPolicyManagerChromeOS::GetChromePolicy(PolicyMap* policy_map) {
   SetEnterpriseUsersDefaults(policy_map);
 }
 
-void UserCloudPolicyManagerChromeOS::FetchPolicyOAuthTokenUsingSigninProfile() {
-  scoped_refptr<net::URLRequestContextGetter> signin_context;
-  Profile* signin_profile = chromeos::ProfileHelper::GetSigninProfile();
-  if (signin_profile)
-    signin_context = signin_profile->GetRequestContext();
+void UserCloudPolicyManagerChromeOS::FetchPolicyOAuthTokenUsingSigninContext() {
+  scoped_refptr<net::URLRequestContextGetter> signin_context =
+      chromeos::login::GetSigninContext();
   if (!signin_context.get()) {
-    LOG(ERROR) << "No signin Profile for policy oauth token fetch!";
+    LOG(ERROR) << "No signin context for policy oauth token fetch!";
     OnOAuth2PolicyTokenFetched(
         std::string(), GoogleServiceAuthError(GoogleServiceAuthError::NONE));
     return;
diff --git a/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos.h b/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos.h
index ca07361..fdc7766 100644
--- a/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos.h
+++ b/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos.h
@@ -110,8 +110,8 @@ class UserCloudPolicyManagerChromeOS : public CloudPolicyManager,
 
  private:
   // Fetches a policy token using the authentication context of the signin
-  // Profile, and calls back to OnOAuth2PolicyTokenFetched when done.
-  void FetchPolicyOAuthTokenUsingSigninProfile();
+  // context, and calls back to OnOAuth2PolicyTokenFetched when done.
+  void FetchPolicyOAuthTokenUsingSigninContext();
 
   // Called once the policy access token is available, and starts the
   // registration with the policy server if the token was successfully fetched.
diff --git a/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos_unittest.cc b/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos_unittest.cc
index d630d65..7d7f399 100644
--- a/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos_unittest.cc
+++ b/chrome/browser/chromeos/policy/user_cloud_policy_manager_chromeos_unittest.cc
@@ -342,15 +342,10 @@ class UserCloudPolicyManagerChromeOSTest : public testing::Test {
   TestingProfile* profile_;
   TestingProfile* signin_profile_;
 
-  static const char kSigninProfile[];
-
  private:
   DISALLOW_COPY_AND_ASSIGN(UserCloudPolicyManagerChromeOSTest);
 };
 
-const char UserCloudPolicyManagerChromeOSTest::kSigninProfile[] =
-    "signin_profile";
-
 TEST_F(UserCloudPolicyManagerChromeOSTest, BlockingFirstFetch) {
   // Tests the initialization of a manager whose Profile is waiting for the
   // initial fetch, when the policy cache is empty.
diff --git a/chrome/browser/resources/chromeos/login/screen_gaia_signin.js b/chrome/browser/resources/chromeos/login/screen_gaia_signin.js
index f4139cb..367e6f1 100644
--- a/chrome/browser/resources/chromeos/login/screen_gaia_signin.js
+++ b/chrome/browser/resources/chromeos/login/screen_gaia_signin.js
@@ -600,7 +600,8 @@ login.createScreen('GaiaSigninScreen', 'gaia-signin', function() {
                       [credentials.gaiaId,
                        credentials.email,
                        credentials.password,
-                       credentials.authCode]);
+                       credentials.authCode,
+                       credentials.usingSAML]);
         }
       } else {
         chrome.send('completeLogin',
diff --git a/chrome/browser/ui/webui/chromeos/login/base_screen_handler.h b/chrome/browser/ui/webui/chromeos/login/base_screen_handler.h
index b1b6e8b..ce57647 100644
--- a/chrome/browser/ui/webui/chromeos/login/base_screen_handler.h
+++ b/chrome/browser/ui/webui/chromeos/login/base_screen_handler.h
@@ -172,6 +172,22 @@ class BaseScreenHandler : public content::WebUIMessageHandler,
         name, base::Bind(&::login::CallbackWrapper4<A1, A2, A3, A4>, callback));
   }
 
+  template <typename T,
+            typename A1,
+            typename A2,
+            typename A3,
+            typename A4,
+            typename A5>
+  void AddCallback(
+      const std::string& name,
+      void (T::*method)(A1 arg1, A2 arg2, A3 arg3, A4 arg4, A5 arg5)) {
+    base::Callback<void(A1, A2, A3, A4, A5)> callback =
+        base::Bind(method, base::Unretained(static_cast<T*>(this)));
+    web_ui()->RegisterMessageCallback(
+        name,
+        base::Bind(&::login::CallbackWrapper5<A1, A2, A3, A4, A5>, callback));
+  }
+
   template <typename Method>
   void AddPrefixedCallback(const std::string& unprefixed_name,
                            const Method& method) {
diff --git a/chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.cc b/chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.cc
index d98dce5..47af3d4 100644
--- a/chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.cc
+++ b/chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.cc
@@ -402,7 +402,8 @@ void GaiaScreenHandler::HandleCompleteAuthentication(
     const std::string& gaia_id,
     const std::string& email,
     const std::string& password,
-    const std::string& auth_code) {
+    const std::string& auth_code,
+    bool using_saml) {
   if (!Delegate())
     return;
 
@@ -416,6 +417,9 @@ void GaiaScreenHandler::HandleCompleteAuthentication(
   user_context.SetGaiaID(gaia_id);
   user_context.SetKey(Key(password));
   user_context.SetAuthCode(auth_code);
+  user_context.SetAuthFlow(using_saml
+                               ? UserContext::AUTH_FLOW_GAIA_WITH_SAML
+                               : UserContext::AUTH_FLOW_GAIA_WITHOUT_SAML);
   Delegate()->CompleteLogin(user_context);
 }
 
diff --git a/chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.h b/chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.h
index 8af7c07..f06086a 100644
--- a/chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.h
+++ b/chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.h
@@ -103,7 +103,8 @@ class GaiaScreenHandler : public BaseScreenHandler {
   void HandleCompleteAuthentication(const std::string& gaia_id,
                                     const std::string& email,
                                     const std::string& password,
-                                    const std::string& auth_code);
+                                    const std::string& auth_code,
+                                    bool using_saml);
   void HandleCompleteAuthenticationAuthCodeOnly(const std::string& auth_code);
   void HandleCompleteLogin(const std::string& gaia_id,
                            const std::string& typed_email,
diff --git a/components/login/base_screen_handler_utils.h b/components/login/base_screen_handler_utils.h
index 63e5512..c2b8971 100644
--- a/components/login/base_screen_handler_utils.h
+++ b/components/login/base_screen_handler_utils.h
@@ -122,6 +122,25 @@ void CallbackWrapper4(base::Callback<void(A1, A2, A3, A4)> callback,
   callback.Run(arg1, arg2, arg3, arg4);
 }
 
+template <typename A1, typename A2, typename A3, typename A4, typename A5>
+void CallbackWrapper5(base::Callback<void(A1, A2, A3, A4, A5)> callback,
+                      const base::ListValue* args) {
+  DCHECK(args);
+  DCHECK_EQ(5u, args->GetSize());
+  typename UnwrapConstRef<A1>::Type arg1;
+  typename UnwrapConstRef<A2>::Type arg2;
+  typename UnwrapConstRef<A3>::Type arg3;
+  typename UnwrapConstRef<A4>::Type arg4;
+  typename UnwrapConstRef<A5>::Type arg5;
+  if (!GetArg(args, 0, &arg1) || !GetArg(args, 1, &arg2) ||
+      !GetArg(args, 2, &arg3) || !GetArg(args, 3, &arg4) ||
+      !GetArg(args, 4, &arg5)) {
+    NOTREACHED();
+    return;
+  }
+  callback.Run(arg1, arg2, arg3, arg4, arg5);
+}
+
 }  // namespace login
 
 #endif  // COMPONENTS_LOGIN_BASE_SCREEN_HANDLER_UTILS_H_