// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/ash/login/lock/online_reauth/lock_screen_reauth_manager.h"

#include <memory>

#include "ash/constants/ash_features.h"
#include "ash/constants/ash_login_pref_names.h"
#include "ash/login/mock_login_screen_client.h"
#include "ash/public/cpp/reauth_reason.h"
#include "ash/test/ash_test_helper.h"
#include "base/memory/raw_ptr.h"
#include "base/test/metrics/histogram_tester.h"
#include "base/test/scoped_feature_list.h"
#include "base/test/test_future.h"
#include "base/time/time.h"
#include "base/values.h"
#include "chrome/browser/ash/login/saml/mock_lock_handler.h"
#include "chrome/browser/ash/login/users/fake_chrome_user_manager.h"
#include "chrome/browser/profiles/profile_manager.h"
#include "chrome/test/base/testing_browser_process.h"
#include "chrome/test/base/testing_profile.h"
#include "chrome/test/base/testing_profile_manager.h"
#include "chromeos/ash/components/dbus/userdataauth/fake_userdataauth_client.h"
#include "chromeos/ash/components/dbus/userdataauth/userdataauth_client.h"
#include "chromeos/ash/components/login/auth/public/cryptohome_key_constants.h"
#include "chromeos/ash/components/login/auth/public/user_context.h"
#include "chromeos/ash/components/osauth/public/auth_parts.h"
#include "components/account_id/account_id.h"
#include "components/session_manager/core/session_manager.h"
#include "components/user_manager/known_user.h"
#include "components/user_manager/scoped_user_manager.h"
#include "components/user_manager/user_names.h"
#include "content/public/test/browser_task_environment.h"
#include "google_apis/gaia/gaia_id.h"

namespace ash {
namespace {
constexpr GaiaId::Literal kSAMLUserId1("12345");
constexpr char kSAMLUserEmail1[] = "alice@corp.example.com";

constexpr GaiaId::Literal kSAMLUserId2("67891");
constexpr char kSAMLUserEmail2[] = "bob@corp.example.com";

constexpr base::TimeDelta kSamlOnlineShortDelay = base::Seconds(10);

constexpr char kLockScreenReauthHistogram[] =
    "ChromeOS.LockScreenReauth.LockScreenReauthReason";
constexpr char kFakePassword[] = "p4ssw07d";
constexpr char kFakePIN[] = "7008";
}  // namespace

class LockScreenReauthManagerTest : public testing::Test {
 protected:
  LockScreenReauthManagerTest();
  ~LockScreenReauthManagerTest() override;

  // testing::Test:
  void SetUp() override;
  void TearDown() override;

  void CreateLockScreenReauthManager();
  void DestroyLockScreenReauthManager();

  void SetReauthRequiredBySamlTokenMismatch();
  bool IsReauthRequiredBySamlTokenMismatch();
  bool IsReauthRequiredBySamlTimeLimitPolicy();
  void MaybeForceReauthOnLockScreen(ReauthReason reason);

  void SetCryptohomePassword(AccountId user, std::string type);
  void SetCryptohomePin(AccountId user);
  void ClearAuthFactors(AccountId user);

  void LockScreen();

  const AccountId saml_login_account_id1_ =
      AccountId::FromUserEmailGaiaId(kSAMLUserEmail1, kSAMLUserId1);

  const AccountId saml_login_account_id2_ =
      AccountId::FromUserEmailGaiaId(kSAMLUserEmail2, kSAMLUserId2);

  content::BrowserTaskEnvironment test_environment_{
      base::test::TaskEnvironment::MainThreadType::UI,
      base::test::TaskEnvironment::TimeSource::MOCK_TIME};
  user_manager::TypedScopedUserManager<ash::FakeChromeUserManager>
      fake_user_manager_{std::make_unique<ash::FakeChromeUserManager>()};
  TestingProfileManager profile_manager_{TestingBrowserProcess::GetGlobal()};
  raw_ptr<TestingProfile> primary_profile_ = nullptr;
  raw_ptr<TestingProfile> secondary_profile_ = nullptr;

  MockLockHandler lock_handler_;
  std::unique_ptr<LockScreenReauthManager> manager_;
  base::test::ScopedFeatureList feature_list_;
  std::unique_ptr<user_manager::KnownUser> known_user_;
  const base::HistogramTester histogram_tester_;
  base::test::TestFuture<void> auth_configuration_exit_future;
  // `AshTestHelper` makes sure that `SessionManager` is set up in these tests.
  ash::AshTestHelper ash_test_helper_;
};

LockScreenReauthManagerTest::LockScreenReauthManagerTest() : manager_(nullptr) {
  UserDataAuthClient::InitializeFake();
  known_user_ = std::make_unique<user_manager::KnownUser>(
      TestingBrowserProcess::GetGlobal()->local_state());
  feature_list_.InitAndEnableFeature({features::kManagedLocalPinAndPassword});
}

LockScreenReauthManagerTest::~LockScreenReauthManagerTest() {
  DestroyLockScreenReauthManager();
  UserDataAuthClient::Shutdown();
}

void LockScreenReauthManagerTest::SetUp() {
  ASSERT_TRUE(profile_manager_.SetUp());
  primary_profile_ = profile_manager_.CreateTestingProfile("test1");
  secondary_profile_ = profile_manager_.CreateTestingProfile("test2");

  fake_user_manager_->AddUserWithAffiliationAndTypeAndProfile(
      saml_login_account_id1_, /* is_affiliated = */ false,
      user_manager::UserType::kRegular, primary_profile_);
  fake_user_manager_->AddUserWithAffiliationAndTypeAndProfile(
      saml_login_account_id2_, /* is_affiliated = */ false,
      user_manager::UserType::kRegular, secondary_profile_);
  fake_user_manager_->AddUser(saml_login_account_id2_);
  fake_user_manager_->LoginUser(saml_login_account_id1_);
  // ActiveUser in FakeChromeUserManager needs to be set explicitly.
  fake_user_manager_->SwitchActiveUser(saml_login_account_id1_);
  ASSERT_TRUE(fake_user_manager_->GetActiveUser());
  FakeUserDataAuthClient::TestApi::Get()->CreatePostponedDirectories();
  auto account_id =
      cryptohome::CreateAccountIdentifierFromAccountId(saml_login_account_id1_);
  FakeUserDataAuthClient::TestApi::Get()->AddExistingUser(
      std::move(account_id));
  SetCryptohomePassword(saml_login_account_id1_, kCryptohomeGaiaKeyLabel);
  ash_test_helper_.SetUp();
}

void LockScreenReauthManagerTest::TearDown() {
  ash_test_helper_.TearDown();
  proximity_auth::ScreenlockBridge::Get()->SetLockHandler(nullptr);
}

void LockScreenReauthManagerTest::CreateLockScreenReauthManager() {
  DestroyLockScreenReauthManager();
  manager_ = std::make_unique<LockScreenReauthManager>(
      TestingBrowserProcess::GetGlobal()->local_state(), primary_profile_);
  manager_->SetClockForTesting(test_environment_.GetMockClock());
  manager_->SetGetAuthfactorsConfigurationCallbackForTesting(
      auth_configuration_exit_future.GetRepeatingCallback());
}

void LockScreenReauthManagerTest::DestroyLockScreenReauthManager() {
  if (manager_) {
    manager_->Shutdown();
    manager_ = nullptr;
  }
}

void LockScreenReauthManagerTest::LockScreen() {
  proximity_auth::ScreenlockBridge::Get()->SetLockHandler(&lock_handler_);
  session_manager::SessionManager::Get()->SetSessionState(
      session_manager::SessionState::LOCKED);
}

void LockScreenReauthManagerTest::SetReauthRequiredBySamlTokenMismatch() {
  manager_->is_reauth_required_by_saml_token_mismatch_ = true;
}

bool LockScreenReauthManagerTest::IsReauthRequiredBySamlTokenMismatch() {
  return manager_->is_reauth_required_by_saml_token_mismatch_;
}

bool LockScreenReauthManagerTest::IsReauthRequiredBySamlTimeLimitPolicy() {
  return manager_->is_reauth_required_by_saml_time_limit_policy_;
}

void LockScreenReauthManagerTest::SetCryptohomePassword(AccountId user,
                                                        std::string label) {
  Key key(kFakePassword);
  user_data_auth::AuthFactor auth_factor;
  user_data_auth::AuthInput auth_input;

  auth_factor.set_label(label);
  auth_factor.set_type(user_data_auth::AUTH_FACTOR_TYPE_PASSWORD);

  auth_input.mutable_password_input()->set_secret(key.GetSecret());

  // Add the password key to the user.
  FakeUserDataAuthClient::TestApi::Get()->AddAuthFactor(
      cryptohome::CreateAccountIdentifierFromAccountId(user), auth_factor,
      auth_input);
}

void LockScreenReauthManagerTest::SetCryptohomePin(AccountId user) {
  Key key(kFakePIN);

  user_data_auth::AuthFactor auth_factor;
  user_data_auth::AuthInput auth_input;

  auth_factor.set_label(ash::kCryptohomePinLabel);
  auth_factor.set_type(user_data_auth::AUTH_FACTOR_TYPE_PIN);

  auth_input.mutable_pin_input()->set_secret(key.GetSecret());

  // Add the pin to the user.
  FakeUserDataAuthClient::TestApi::Get()->AddAuthFactor(
      cryptohome::CreateAccountIdentifierFromAccountId(user), auth_factor,
      auth_input);
}

void LockScreenReauthManagerTest::ClearAuthFactors(AccountId user) {
  FakeUserDataAuthClient::TestApi::Get()->ClearAuthFactors(
      cryptohome::CreateAccountIdentifierFromAccountId(user));
}

void LockScreenReauthManagerTest::MaybeForceReauthOnLockScreen(
    ReauthReason reason) {
  manager_->MaybeForceReauthOnLockScreen(reason);
  if (features::IsManagedLocalPinAndPasswordEnabled()) {
    ASSERT_TRUE(auth_configuration_exit_future.Wait())
        << "Failed to wait for the auth confiugration exit callback";
    auth_configuration_exit_future.Clear();
  }
}

TEST_F(LockScreenReauthManagerTest, ReauthenticateRequiredByTimelimitPolicy) {
  CreateLockScreenReauthManager();
  fake_user_manager_->SaveForceOnlineSignin(saml_login_account_id1_, true);
  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);
  EXPECT_TRUE(IsReauthRequiredBySamlTimeLimitPolicy());
}

TEST_F(LockScreenReauthManagerTest, ReauthenticateResetByToken) {
  primary_profile_->GetPrefs()->SetBoolean(
      ash::prefs::kLockScreenReauthenticationEnabled, true);
  CreateLockScreenReauthManager();
  fake_user_manager_->SaveForceOnlineSignin(saml_login_account_id1_, true);
  MaybeForceReauthOnLockScreen(
      ReauthReason::kSamlPasswordSyncTokenValidationFailed);
  EXPECT_TRUE(IsReauthRequiredBySamlTokenMismatch());
}

TEST_F(LockScreenReauthManagerTest, ReauthenticateSetOnLock) {
  primary_profile_->GetPrefs()->SetBoolean(
      ash::prefs::kLockScreenReauthenticationEnabled, true);
  CreateLockScreenReauthManager();
  EXPECT_CALL(lock_handler_,
              SetAuthType(saml_login_account_id1_,
                          proximity_auth::mojom::AuthType::ONLINE_SIGN_IN,
                          std::u16string()))
      .Times(1);
  LockScreen();
  fake_user_manager_->SaveForceOnlineSignin(saml_login_account_id1_, true);
  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);
  EXPECT_TRUE(IsReauthRequiredBySamlTimeLimitPolicy());
}

// User tries to unlock the screen using valid SAML credentials but not for the
// user who locked the screen. As a result screen remains locked.
TEST_F(LockScreenReauthManagerTest, AuthenticateWithIncorrectUser) {
  primary_profile_->GetPrefs()->SetBoolean(
      ash::prefs::kLockScreenReauthenticationEnabled, true);
  CreateLockScreenReauthManager();
  EXPECT_CALL(lock_handler_,
              SetAuthType(saml_login_account_id1_,
                          proximity_auth::mojom::AuthType::ONLINE_SIGN_IN,
                          std::u16string()))
      .Times(1);
  LockScreen();
  EXPECT_CALL(lock_handler_, Unlock(saml_login_account_id1_)).Times(0);
  fake_user_manager_->SaveForceOnlineSignin(saml_login_account_id1_, true);
  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);
  EXPECT_TRUE(IsReauthRequiredBySamlTimeLimitPolicy());
  UserContext user_context(user_manager::UserType::kRegular,
                           saml_login_account_id2_);
  EXPECT_DEATH(manager_->OnAuthSuccess(user_context), "");

  EXPECT_TRUE(IsReauthRequiredBySamlTimeLimitPolicy());
  EXPECT_TRUE(proximity_auth::ScreenlockBridge::Get()->IsLocked());
  histogram_tester_.ExpectBucketCount(
      kLockScreenReauthHistogram, ReauthReason::kGaiaLockScreenReauthPolicy, 0);
  histogram_tester_.ExpectBucketCount(
      kLockScreenReauthHistogram, ReauthReason::kSamlLockScreenReauthPolicy, 0);
  histogram_tester_.ExpectBucketCount(
      kLockScreenReauthHistogram,
      ReauthReason::kSamlPasswordSyncTokenValidationFailed, 0);
}

TEST_F(LockScreenReauthManagerTest, AuthenticateWithCorrectUser) {
  base::Time now = test_environment_.GetMockClock()->Now();
  known_user_->SetLastOnlineSignin(saml_login_account_id1_, now);
  known_user_->SetOfflineSigninLimit(saml_login_account_id1_,
                                     kSamlOnlineShortDelay);
  base::Time expected_signin_time = now + kSamlOnlineShortDelay;

  CreateLockScreenReauthManager();
  EXPECT_CALL(lock_handler_,
              SetAuthType(saml_login_account_id1_,
                          proximity_auth::mojom::AuthType::ONLINE_SIGN_IN,
                          std::u16string()))
      .Times(1);
  EXPECT_CALL(lock_handler_, Unlock(saml_login_account_id1_)).Times(1);
  LockScreen();
  fake_user_manager_->SaveForceOnlineSignin(saml_login_account_id1_, true);
  test_environment_.FastForwardBy(kSamlOnlineShortDelay);
  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);
  EXPECT_TRUE(IsReauthRequiredBySamlTimeLimitPolicy());
  UserContext user_context(user_manager::UserType::kRegular,
                           saml_login_account_id1_);
  manager_->OnAuthSuccess(user_context);
  EXPECT_FALSE(IsReauthRequiredBySamlTimeLimitPolicy());
  now = known_user_->GetLastOnlineSignin(saml_login_account_id1_);
  EXPECT_EQ(now, expected_signin_time);
  histogram_tester_.ExpectBucketCount(
      kLockScreenReauthHistogram, ReauthReason::kGaiaLockScreenReauthPolicy, 0);
  histogram_tester_.ExpectBucketCount(
      kLockScreenReauthHistogram, ReauthReason::kSamlLockScreenReauthPolicy, 1);
  histogram_tester_.ExpectBucketCount(
      kLockScreenReauthHistogram,
      ReauthReason::kSamlPasswordSyncTokenValidationFailed, 0);
}

TEST_F(LockScreenReauthManagerTest, FlowTriggeredByPolicyAndInvalidToken) {
  base::Time now = test_environment_.GetMockClock()->Now();
  known_user_->SetLastOnlineSignin(saml_login_account_id1_, now);
  known_user_->SetOfflineSigninLimit(saml_login_account_id1_,
                                     kSamlOnlineShortDelay);
  base::Time expected_signin_time = now + kSamlOnlineShortDelay;

  CreateLockScreenReauthManager();
  EXPECT_CALL(lock_handler_,
              SetAuthType(saml_login_account_id1_,
                          proximity_auth::mojom::AuthType::ONLINE_SIGN_IN,
                          std::u16string()))
      .Times(1);
  EXPECT_CALL(lock_handler_, Unlock(saml_login_account_id1_)).Times(1);
  LockScreen();
  fake_user_manager_->SaveForceOnlineSignin(saml_login_account_id1_, true);
  SetReauthRequiredBySamlTokenMismatch();
  test_environment_.FastForwardBy(kSamlOnlineShortDelay);
  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);
  EXPECT_TRUE(IsReauthRequiredBySamlTimeLimitPolicy());
  UserContext user_context(user_manager::UserType::kRegular,
                           saml_login_account_id1_);
  manager_->OnAuthSuccess(user_context);
  EXPECT_FALSE(IsReauthRequiredBySamlTimeLimitPolicy());
  now = known_user_->GetLastOnlineSignin(saml_login_account_id1_);
  EXPECT_EQ(now, expected_signin_time);
  histogram_tester_.ExpectBucketCount(
      kLockScreenReauthHistogram, ReauthReason::kGaiaLockScreenReauthPolicy, 0);
  histogram_tester_.ExpectBucketCount(
      kLockScreenReauthHistogram, ReauthReason::kSamlLockScreenReauthPolicy, 1);
  histogram_tester_.ExpectBucketCount(
      kLockScreenReauthHistogram,
      ReauthReason::kSamlPasswordSyncTokenValidationFailed, 1);
}

TEST_F(LockScreenReauthManagerTest, PolicySetToFalse) {
  primary_profile_->GetPrefs()->SetBoolean(
      ash::prefs::kLockScreenReauthenticationEnabled, false);
  CreateLockScreenReauthManager();
  EXPECT_FALSE(manager_->ShouldPasswordSyncTriggerReauth());
}

TEST_F(LockScreenReauthManagerTest, PolicyNotSet) {
  CreateLockScreenReauthManager();
  EXPECT_FALSE(manager_->ShouldPasswordSyncTriggerReauth());
}

TEST_F(LockScreenReauthManagerTest, ReauthWithLocalPasswordEnabled) {
  primary_profile_->GetPrefs()->SetBoolean(
      ash::prefs::kLockScreenReauthenticationEnabled, true);
  // Remove the online password as an auth factor to test local password only.
  ClearAuthFactors(saml_login_account_id1_);
  SetCryptohomePassword(saml_login_account_id1_,
                        kCryptohomeLocalPasswordKeyLabel);
  CreateLockScreenReauthManager();
  LockScreen();

  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);

  EXPECT_FALSE(IsReauthRequiredBySamlTimeLimitPolicy());
}

TEST_F(LockScreenReauthManagerTest, ReauthWithGaiaPasswordEnabled) {
  primary_profile_->GetPrefs()->SetBoolean(
      ash::prefs::kLockScreenReauthenticationEnabled, true);
  // Online Password is already added as a factor.
  CreateLockScreenReauthManager();
  LockScreen();

  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);

  EXPECT_TRUE(IsReauthRequiredBySamlTimeLimitPolicy());
}

TEST_F(LockScreenReauthManagerTest, ReauthWithPinAndGaiaPasswordEnabled) {
  primary_profile_->GetPrefs()->SetBoolean(
      ash::prefs::kLockScreenReauthenticationEnabled, true);
  // Online Password is already added as a factor.
  SetCryptohomePin(saml_login_account_id1_);
  CreateLockScreenReauthManager();
  LockScreen();

  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);

  EXPECT_TRUE(IsReauthRequiredBySamlTimeLimitPolicy());
}

TEST_F(LockScreenReauthManagerTest, ReauthWithPinEnabled) {
  primary_profile_->GetPrefs()->SetBoolean(
      ash::prefs::kLockScreenReauthenticationEnabled, true);
  // Remove the online password as an auth factor to test pin only.
  ClearAuthFactors(saml_login_account_id1_);
  SetCryptohomePin(saml_login_account_id1_);
  CreateLockScreenReauthManager();
  LockScreen();

  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);

  EXPECT_FALSE(IsReauthRequiredBySamlTimeLimitPolicy());
}

class AutoStartLockScreenReauthManagerTest
    : public LockScreenReauthManagerTest,
      public testing::WithParamInterface<bool> {
 protected:
  void SetUp() override;
  void TearDown() override;

  std::unique_ptr<ash::MockLoginScreenClient> login_screen_client_;
};

void AutoStartLockScreenReauthManagerTest::SetUp() {
  LockScreenReauthManagerTest::SetUp();
  login_screen_client_ = std::make_unique<ash::MockLoginScreenClient>();
}

void AutoStartLockScreenReauthManagerTest::TearDown() {
  LockScreenReauthManagerTest::TearDown();
}

TEST_P(AutoStartLockScreenReauthManagerTest,
       ForceOnlineReauthOnSessionStateChanged) {
  const bool is_auto_start_enabled = GetParam();
  primary_profile_->GetPrefs()->SetBoolean(
      ash::prefs::kLockScreenAutoStartOnlineReauth, is_auto_start_enabled);
  CreateLockScreenReauthManager();
  MaybeForceReauthOnLockScreen(ReauthReason::kSamlLockScreenReauthPolicy);
  EXPECT_CALL(lock_handler_,
              SetAuthType(saml_login_account_id1_,
                          proximity_auth::mojom::AuthType::ONLINE_SIGN_IN,
                          std::u16string()))
      .Times(1);
  EXPECT_CALL(*login_screen_client_, ShowGaiaSignin(saml_login_account_id1_))
      .Times(is_auto_start_enabled);
  LockScreen();
}

INSTANTIATE_TEST_SUITE_P(All,
                         AutoStartLockScreenReauthManagerTest,
                         /*is_auto_start_enabled=*/testing::Bool());

}  // namespace ash