oobe: Wait for oauth_code during enrollment flow

This CL extracts CookieWaiter class from the GaiaScreenHandler and reuse
it in the EnrollmentScreenHandler.

Bug: b/167404108
Change-Id: Idaeca9cf93157b42819eac33a653eebe27def867
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2429025Reviewed-by: Roman Aleksandrov <raleksandrov@google.com>
Commit-Queue: Roman Sorokin [CET] <rsorokin@chromium.org>
Cr-Commit-Position: refs/heads/master@{#812160}

chrome/browser/ui/BUILD.gn

@@ -2215,6 +2215,8 @@ static_library("ui") {
       "webui/chromeos/login/base_screen_handler.h",
       "webui/chromeos/login/base_webui_handler.cc",
       "webui/chromeos/login/base_webui_handler.h",
+      "webui/chromeos/login/cookie_waiter.cc",
+      "webui/chromeos/login/cookie_waiter.h",
       "webui/chromeos/login/core_oobe_handler.cc",
       "webui/chromeos/login/core_oobe_handler.h",
       "webui/chromeos/login/debug/debug_overlay_handler.cc",

chrome/browser/ui/webui/chromeos/login/cookie_waiter.cc

+// Copyright 2020 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/ui/webui/chromeos/login/cookie_waiter.h"
+
+#include "google_apis/gaia/gaia_urls.h"
+
+namespace chromeos {
+
+namespace {
+
+constexpr base::TimeDelta kCookieDelay = base::TimeDelta::FromSeconds(20);
+
+}
+
+CookieWaiter::CookieWaiter(network::mojom::CookieManager* cookie_manager,
+                           const std::string& cookie_name,
+                           base::RepeatingClosure on_cookie_change,
+                           base::OnceClosure on_timeout)
+    : on_cookie_change_(std::move(on_cookie_change)) {
+  cookie_manager->AddCookieChangeListener(
+      GaiaUrls::GetInstance()->gaia_url(), cookie_name,
+      cookie_listener_.BindNewPipeAndPassRemote());
+  waiting_timer_.Start(FROM_HERE, kCookieDelay, std::move(on_timeout));
+}
+
+CookieWaiter::~CookieWaiter() = default;
+
+void CookieWaiter::OnCookieChange(const net::CookieChangeInfo&) {
+  on_cookie_change_.Run();
+}
+
+}  // namespace chromeos

chrome/browser/ui/webui/chromeos/login/cookie_waiter.h

+// Copyright 2020 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.
+
+#ifndef CHROME_BROWSER_UI_WEBUI_CHROMEOS_LOGIN_COOKIE_WAITER_H_
+#define CHROME_BROWSER_UI_WEBUI_CHROMEOS_LOGIN_COOKIE_WAITER_H_
+
+#include <string>
+
+#include "base/callback.h"
+#include "base/timer/timer.h"
+#include "mojo/public/cpp/bindings/receiver.h"
+#include "services/network/public/mojom/cookie_manager.mojom.h"
+
+namespace chromeos {
+
+class CookieWaiter : public network::mojom::CookieChangeListener {
+ public:
+  CookieWaiter(network::mojom::CookieManager* cookie_manager,
+               const std::string& cookie_name,
+               base::RepeatingClosure on_cookie_change,
+               base::OnceClosure on_timeout);
+  ~CookieWaiter() override;
+
+  CookieWaiter(const CookieWaiter&) = delete;
+  CookieWaiter& operator=(const CookieWaiter&) = delete;
+
+  // network::mojom::CookieChangeListener:
+  void OnCookieChange(const net::CookieChangeInfo&) override;
+
+ private:
+  base::RepeatingClosure on_cookie_change_;
+
+  mojo::Receiver<network::mojom::CookieChangeListener> cookie_listener_{this};
+  base::OneShotTimer waiting_timer_;
+};
+
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_UI_WEBUI_CHROMEOS_LOGIN_COOKIE_WAITER_H_

chrome/browser/ui/webui/chromeos/login/enrollment_screen_handler.cc

@@ -31,6 +31,7 @@
 #include "chrome/browser/chromeos/policy/policy_oauth2_token_fetcher.h"
 #include "chrome/browser/policy/enrollment_status.h"
 #include "chrome/browser/profiles/profile.h"
+#include "chrome/browser/ui/webui/chromeos/login/cookie_waiter.h"
 #include "chrome/browser/ui/webui/chromeos/login/error_screen_handler.h"
 #include "chrome/grit/generated_resources.h"
 #include "chromeos/network/network_state.h"
@@ -64,6 +65,8 @@ const char kEnrollmentModeUIRecovery[] = "recovery";
 constexpr char kActiveDirectoryJoinHistogram[] =
     "Enterprise.ActiveDirectoryJoin";
 
+constexpr char kOAUTHCodeCookie[] = "oauth_code";
+
 // Converts |mode| to a mode identifier for the UI.
 std::string EnrollmentModeToUIMode(policy::EnrollmentConfig::Mode mode) {
   switch (mode) {
@@ -773,7 +776,37 @@ void EnrollmentScreenHandler::HandleCompleteLogin(const std::string& user) {
   DCHECK_EQ(signin_partition_manager->GetCurrentStoragePartitionName(),
             signin_partition_name_);
 
-  partition->GetCookieManagerForBrowserProcess()->GetCookieList(
+  network::mojom::CookieManager* cookie_manager =
+      partition->GetCookieManagerForBrowserProcess();
+  if (!oauth_code_waiter_) {
+    // Set listener before requesting the cookies to avoid race conditions.
+    oauth_code_waiter_ = std::make_unique<CookieWaiter>(
+        cookie_manager, kOAUTHCodeCookie,
+        base::BindRepeating(&EnrollmentScreenHandler::
+                                ContinueAuthenticationWhenCookiesAvailable,
+                            weak_ptr_factory_.GetWeakPtr(), user),
+        base::BindOnce(&EnrollmentScreenHandler::OnCookieWaitTimeout,
+                       weak_ptr_factory_.GetWeakPtr()));
+  }
+
+  ContinueAuthenticationWhenCookiesAvailable(user);
+}
+
+void EnrollmentScreenHandler::ContinueAuthenticationWhenCookiesAvailable(
+    const std::string& user) {
+  login::SigninPartitionManager* signin_partition_manager =
+      login::SigninPartitionManager::Factory::GetForBrowserContext(
+          Profile::FromWebUI(web_ui()));
+  content::StoragePartition* partition =
+      signin_partition_manager->GetCurrentStoragePartition();
+
+  // Validity check that partition did not change during enrollment flow.
+  DCHECK_EQ(signin_partition_manager->GetCurrentStoragePartitionName(),
+            signin_partition_name_);
+
+  network::mojom::CookieManager* cookie_manager =
+      partition->GetCookieManagerForBrowserProcess();
+  cookie_manager->GetCookieList(
       GaiaUrls::GetInstance()->gaia_url(),
       net::CookieOptions::MakeAllInclusive(),
       base::BindOnce(&EnrollmentScreenHandler::OnGetCookiesForCompleteLogin,
@@ -786,18 +819,27 @@ void EnrollmentScreenHandler::OnGetCookiesForCompleteLogin(
     const net::CookieAccessResultList& excluded_cookies) {
   std::string auth_code;
   for (const auto& cookie_with_access_result : cookies) {
-    if (cookie_with_access_result.cookie.Name() == "oauth_code") {
+    if (cookie_with_access_result.cookie.Name() == kOAUTHCodeCookie) {
       auth_code = cookie_with_access_result.cookie.Value();
       break;
     }
   }
 
-  DCHECK(!auth_code.empty());
+  if (auth_code.empty()) {
+    // Will try again from oauth_code_waiter callback.
+    return;
+  }
 
+  oauth_code_waiter_.reset();
   DCHECK(controller_);
   controller_->OnLoginDone(gaia::SanitizeEmail(user), auth_code);
 }
 
+void EnrollmentScreenHandler::OnCookieWaitTimeout() {
+  oauth_code_waiter_.reset();
+  ShowError(IDS_LOGIN_FATAL_ERROR_NO_AUTH_TOKEN, true);
+}
+
 void EnrollmentScreenHandler::HandleAdCompleteLogin(
     const std::string& machine_name,
     const std::string& distinguished_name,

chrome/browser/ui/webui/chromeos/login/enrollment_screen_handler.h

@@ -21,6 +21,7 @@
 
 namespace chromeos {
 
+class CookieWaiter;
 class ErrorScreensHistogramHelper;
 class HelpAppLauncher;
 
@@ -98,6 +99,9 @@ class EnrollmentScreenHandler
   // Implements NetworkStateInformer::NetworkStateInformerObserver
   void UpdateState(NetworkError::ErrorReason reason) override;
 
+  void ContinueAuthenticationWhenCookiesAvailable(const std::string& user);
+  void OnCookieWaitTimeout();
+
  private:
   // Handlers for WebUI messages.
   void HandleToggleFakeEnrollment();
@@ -189,6 +193,8 @@ class EnrollmentScreenHandler
   // Help application used for help dialogs.
   scoped_refptr<HelpAppLauncher> help_app_;
 
+  std::unique_ptr<CookieWaiter> oauth_code_waiter_;
+
   base::WeakPtrFactory<EnrollmentScreenHandler> weak_ptr_factory_{this};
 
   DISALLOW_COPY_AND_ASSIGN(EnrollmentScreenHandler);

chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.cc

@@ -53,6 +53,7 @@
 #include "chrome/browser/net/system_network_context_manager.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/ui/ash/login_screen_client.h"
+#include "chrome/browser/ui/webui/chromeos/login/cookie_waiter.h"
 #include "chrome/browser/ui/webui/chromeos/login/enrollment_screen_handler.h"
 #include "chrome/browser/ui/webui/chromeos/login/signin_screen_handler.h"
 #include "chrome/browser/ui/webui/chromeos/login/user_creation_screen_handler.h"
@@ -122,8 +123,6 @@ enum class ChromeOSSamlApiUsed {
   kMaxValue = kSamlApiNotUsed,
 };
 
-constexpr base::TimeDelta kCookieDelay = base::TimeDelta::FromSeconds(20);
-
 void RecordAPILogin(bool is_third_party_idp, bool is_api_used) {
   ChromeOSSamlApiUsed login_type;
   if (!is_third_party_idp) {
@@ -772,10 +771,6 @@ void GaiaScreenHandler::OnPortalDetectionCompleted(
   LoadAuthExtension(true /* force */, false /* offline */);
 }
 
-void GaiaScreenHandler::OnCookieChange(const net::CookieChangeInfo& change) {
-  ContinueAuthenticationWhenCookiesAvailable();
-}
-
 void GaiaScreenHandler::HandleIdentifierEntered(const std::string& user_email) {
   if (LoginDisplayHost::default_host() &&
       !LoginDisplayHost::default_host()->IsUserAllowlisted(
@@ -970,14 +965,13 @@ void GaiaScreenHandler::ContinueAuthenticationWhenCookiesAvailable() {
 
   network::mojom::CookieManager* cookie_manager =
       partition->GetCookieManagerForBrowserProcess();
-  if (!oauth_code_listener_.is_bound()) {
+  if (!oauth_code_waiter_) {
     // Set listener before requesting the cookies to avoid race conditions.
-    cookie_manager->AddCookieChangeListener(
-        GaiaUrls::GetInstance()->gaia_url(), kOAUTHCodeCookie,
-        oauth_code_listener_.BindNewPipeAndPassRemote());
-    cookie_waiting_timer_ = std::make_unique<base::OneShotTimer>();
-    cookie_waiting_timer_->Start(
-        FROM_HERE, kCookieDelay,
+    oauth_code_waiter_ = std::make_unique<CookieWaiter>(
+        cookie_manager, kOAUTHCodeCookie,
+        base::BindRepeating(
+            &GaiaScreenHandler::ContinueAuthenticationWhenCookiesAvailable,
+            weak_factory_.GetWeakPtr()),
         base::BindOnce(&GaiaScreenHandler::OnCookieWaitTimeout,
                        weak_factory_.GetWeakPtr()));
   }
@@ -1003,15 +997,14 @@ void GaiaScreenHandler::OnGetCookiesForCompleteAuthentication(
   }
 
   if (auth_code.empty()) {
-    // Will try again from onCookieChange.
+    // Will try again from oauth_code_waiter callback.
     return;
   }
 
   DCHECK(pending_user_context_);
   UserContext user_context = *pending_user_context_;
   pending_user_context_.reset();
-  oauth_code_listener_.reset();
-  cookie_waiting_timer_.reset();
+  oauth_code_waiter_.reset();
 
   user_context.SetAuthCode(auth_code);
   if (!gaps_cookie.empty())
@@ -1023,8 +1016,7 @@ void GaiaScreenHandler::OnGetCookiesForCompleteAuthentication(
 void GaiaScreenHandler::OnCookieWaitTimeout() {
   DCHECK(pending_user_context_);
   pending_user_context_.reset();
-  oauth_code_listener_.reset();
-  cookie_waiting_timer_.reset();
+  oauth_code_waiter_.reset();
   LoadAuthExtension(true /* force */, false /* offline */);
   core_oobe_view_->ShowSignInError(
       0, l10n_util::GetStringUTF8(IDS_LOGIN_FATAL_ERROR_NO_AUTH_TOKEN),

chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.h

@@ -23,13 +23,11 @@
 #include "net/base/net_errors.h"
 #include "net/cookies/canonical_cookie.h"
 #include "net/cookies/cookie_access_result.h"
-#include "services/network/public/mojom/cookie_manager.mojom.h"
 
 class AccountId;
 
 namespace base {
 class DictionaryValue;
-class OneShotTimer;
 }  // namespace base
 
 namespace network {
@@ -38,6 +36,7 @@ class NSSTempCertsCacheChromeOS;
 
 namespace chromeos {
 
+class CookieWaiter;
 class Key;
 class SamlPasswordAttributes;
 class SigninScreenHandler;
@@ -97,7 +96,6 @@ class GaiaView {
 class GaiaScreenHandler : public BaseScreenHandler,
                           public GaiaView,
                           public NetworkPortalDetector::Observer,
-                          public network::mojom::CookieChangeListener,
                           public SecurityTokenPinDialogHost {
  public:
   using TView = GaiaView;
@@ -214,9 +212,6 @@ class GaiaScreenHandler : public BaseScreenHandler,
       const NetworkState* network,
       const NetworkPortalDetector::CaptivePortalState& state) override;
 
-  // network::mojom::CookieChangeListener:
-  void OnCookieChange(const net::CookieChangeInfo& change) override;
-
   // WebUI message handlers.
   void HandleWebviewLoadAborted(int error_code);
   void HandleCompleteAuthentication(
@@ -472,10 +467,8 @@ class GaiaScreenHandler : public BaseScreenHandler,
   std::unique_ptr<SamlChallengeKeyHandler> saml_challenge_key_handler_for_test_;
 
   // Connection to the CookieManager that signals when the GAIA cookies change.
-  mojo::Receiver<network::mojom::CookieChangeListener> oauth_code_listener_{
-      this};
+  std::unique_ptr<CookieWaiter> oauth_code_waiter_;
   std::unique_ptr<UserContext> pending_user_context_;
-  std::unique_ptr<base::OneShotTimer> cookie_waiting_timer_;
 
   base::WeakPtrFactory<GaiaScreenHandler> weak_factory_{this};