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

#include "components/trusted_vault/physical_device_recovery_factor.h"

#include "base/task/bind_post_task.h"
#include "components/trusted_vault/proto_string_bytes_conversion.h"
#include "components/trusted_vault/securebox.h"
#include "components/trusted_vault/trusted_vault_connection.h"

namespace trusted_vault {

namespace {

constexpr int kCurrentDeviceRegistrationVersion = 1;

TrustedVaultDownloadKeysStatusForUMA GetDownloadKeysStatusForUMAFromResponse(
    TrustedVaultDownloadKeysStatus response_status) {
  switch (response_status) {
    case TrustedVaultDownloadKeysStatus::kSuccess:
      return TrustedVaultDownloadKeysStatusForUMA::kSuccess;
    case TrustedVaultDownloadKeysStatus::kMemberNotFound:
      return TrustedVaultDownloadKeysStatusForUMA::kMemberNotFound;
    case TrustedVaultDownloadKeysStatus::kMembershipNotFound:
      return TrustedVaultDownloadKeysStatusForUMA::kMembershipNotFound;
    case TrustedVaultDownloadKeysStatus::kMembershipCorrupted:
      return TrustedVaultDownloadKeysStatusForUMA::kMembershipCorrupted;
    case TrustedVaultDownloadKeysStatus::kMembershipEmpty:
      return TrustedVaultDownloadKeysStatusForUMA::kMembershipEmpty;
    case TrustedVaultDownloadKeysStatus::kNoNewKeys:
      return TrustedVaultDownloadKeysStatusForUMA::kNoNewKeys;
    case TrustedVaultDownloadKeysStatus::kKeyProofsVerificationFailed:
      return TrustedVaultDownloadKeysStatusForUMA::kKeyProofsVerificationFailed;
    case TrustedVaultDownloadKeysStatus::kAccessTokenFetchingFailure:
      return TrustedVaultDownloadKeysStatusForUMA::kAccessTokenFetchingFailure;
    case TrustedVaultDownloadKeysStatus::kNetworkError:
      return TrustedVaultDownloadKeysStatusForUMA::kNetworkError;
    case TrustedVaultDownloadKeysStatus::kOtherError:
      return TrustedVaultDownloadKeysStatusForUMA::kOtherError;
  }

  NOTREACHED();
}

}  // namespace

PhysicalDeviceRecoveryFactor::PhysicalDeviceRecoveryFactor(
    SecurityDomainId security_domain_id,
    PhysicalDeviceStorage* storage,
    KeyStorage* key_storage,
    TrustedVaultThrottlingConnection* connection,
    CoreAccountInfo primary_account)
    : security_domain_id_(security_domain_id),
      storage_(storage),
      key_storage_(key_storage),
      connection_(connection),
      primary_account_(primary_account) {
  CHECK(storage_);
  CHECK(key_storage_);
  CHECK(connection_);
}
PhysicalDeviceRecoveryFactor::~PhysicalDeviceRecoveryFactor() = default;

LocalRecoveryFactorType PhysicalDeviceRecoveryFactor::GetRecoveryFactorType()
    const {
  return LocalRecoveryFactorType::kPhysicalDevice;
}

void PhysicalDeviceRecoveryFactor::AttemptRecovery(AttemptRecoveryCallback cb) {
  const LocalDeviceRegistrationInfo& registration_info =
      storage_->GetLocalDeviceRegistrationInfo(primary_account_.gaia);

  if (!registration_info.device_registered()) {
    FulfillRecoveryWithFailure(
        TrustedVaultDownloadKeysStatusForUMA::kDeviceNotRegistered,
        std::move(cb));
    return;
  }

  if (connection_->AreRequestsThrottled(primary_account_)) {
    FulfillRecoveryWithFailure(
        TrustedVaultDownloadKeysStatusForUMA::kThrottledClientSide,
        std::move(cb));
    return;
  }

  std::unique_ptr<SecureBoxKeyPair> key_pair =
      SecureBoxKeyPair::CreateByPrivateKeyImport(
          ProtoStringToBytes(registration_info.private_key_material()));
  if (!key_pair) {
    // Corrupted state: device is registered, but `key_pair` can't be imported.
    // TODO(crbug.com/40699425): restore from this state (throw away the key
    // and trigger device registration again).
    FulfillRecoveryWithFailure(
        TrustedVaultDownloadKeysStatusForUMA::kCorruptedLocalDeviceRegistration,
        std::move(cb));
    return;
  }

  std::vector<std::vector<uint8_t>> vault_keys =
      key_storage_->GetVaultKeys(primary_account_.gaia);
  int last_vault_key_version =
      key_storage_->GetLastKeyVersion(primary_account_.gaia);

  // Guaranteed by `device_registered` check above.
  CHECK(!vault_keys.empty());
  ongoing_request_ = connection_->DownloadNewKeys(
      primary_account_,
      TrustedVaultKeyAndVersion(vault_keys.back(), last_vault_key_version),
      std::move(key_pair),
      // `this` outlives `ongoing_request_`.
      base::BindOnce(&PhysicalDeviceRecoveryFactor::OnKeysDownloaded,
                     base::Unretained(this), std::move(cb)));
  CHECK(ongoing_request_);
}

bool PhysicalDeviceRecoveryFactor::IsRegistered() {
  return storage_->GetLocalDeviceRegistrationInfo(primary_account_.gaia)
      .device_registered();
}

void PhysicalDeviceRecoveryFactor::MarkAsNotRegistered() {
  storage_->MutateLocalDeviceRegistrationInfo(
      primary_account_.gaia,
      [](LocalDeviceRegistrationInfo& registration_info) {
        registration_info.set_device_registered(false);
        registration_info.clear_device_registered_version();
      });
}

TrustedVaultRecoveryFactorRegistrationStateForUMA
PhysicalDeviceRecoveryFactor::MaybeRegister(RegisterCallback cb) {
  const LocalDeviceRegistrationInfo& registration_info =
      storage_->GetLocalDeviceRegistrationInfo(primary_account_.gaia);

  if (registration_info.device_registered()) {
    static_assert(kCurrentDeviceRegistrationVersion == 1);
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kRegistrationNotAttempted,
        std::move(cb));
    return TrustedVaultRecoveryFactorRegistrationStateForUMA::
        kAlreadyRegisteredV1;
  }

  if (storage_->GetLastRegistrationReturnedLocalDataObsolete(
          primary_account_.gaia)) {
    // Client already knows that existing vault keys (or their absence) isn't
    // sufficient for device registration. Fresh keys should be obtained
    // first.
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kRegistrationNotAttempted,
        std::move(cb));
    return TrustedVaultRecoveryFactorRegistrationStateForUMA::
        kLocalKeysAreStale;
  }

  if (connection_->AreRequestsThrottled(primary_account_)) {
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kRegistrationNotAttempted,
        std::move(cb));
    return TrustedVaultRecoveryFactorRegistrationStateForUMA::
        kThrottledClientSide;
  }

  std::unique_ptr<SecureBoxKeyPair> key_pair;
  if (!registration_info.private_key_material().empty()) {
    key_pair = SecureBoxKeyPair::CreateByPrivateKeyImport(
        /*private_key_bytes=*/ProtoStringToBytes(
            registration_info.private_key_material()));
  }

  const bool had_generated_key_pair = key_pair != nullptr;

  if (!key_pair) {
    key_pair = SecureBoxKeyPair::GenerateRandom();
    // It's possible that device will be successfully registered, but the
    // client won't persist this state (for example response doesn't reach the
    // client or registration callback is cancelled). To avoid duplicated
    // registrations device key is stored before sending the registration
    // request, so the same key will be used for future registration attempts.
    storage_->MutateLocalDeviceRegistrationInfo(
        primary_account_.gaia, [&](LocalDeviceRegistrationInfo& info) {
          AssignBytesToProtoString(key_pair->private_key().ExportToBytes(),
                                   info.mutable_private_key_material());
        });
  }

  if (ongoing_registration_callback_) {
    // Cancel ongoing request before starting a new one.
    ongoing_registration_request_ = nullptr;
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kRegistrationCancelled,
        std::move(ongoing_registration_callback_));
  }

  ongoing_registration_callback_ = std::move(cb);
  std::vector<std::vector<uint8_t>> vault_keys =
      key_storage_->GetVaultKeys(primary_account_.gaia);
  int last_vault_key_version =
      key_storage_->GetLastKeyVersion(primary_account_.gaia);

  // `this` outlives `ongoing_registration_request_`, so it's safe to
  // use base::Unretained() here.
  if (key_storage_->HasNonConstantKey(primary_account_.gaia)) {
    ongoing_registration_request_ = connection_->RegisterAuthenticationFactor(
        primary_account_,
        GetTrustedVaultKeysWithVersions(vault_keys, last_vault_key_version),
        key_pair->public_key(), LocalPhysicalDevice(),
        base::BindOnce(&PhysicalDeviceRecoveryFactor::OnRegistered,
                       base::Unretained(this), true));
  } else {
    ongoing_registration_request_ = connection_->RegisterLocalDeviceWithoutKeys(
        primary_account_, key_pair->public_key(),
        base::BindOnce(&PhysicalDeviceRecoveryFactor::OnRegistered,
                       base::Unretained(this), false));
  }

  CHECK(ongoing_registration_request_);

  return had_generated_key_pair
             ? TrustedVaultRecoveryFactorRegistrationStateForUMA::
                   kAttemptingRegistrationWithExistingKeyPair
             : TrustedVaultRecoveryFactorRegistrationStateForUMA::
                   kAttemptingRegistrationWithNewKeyPair;
}

void PhysicalDeviceRecoveryFactor::OnKeysDownloaded(
    AttemptRecoveryCallback cb,
    TrustedVaultDownloadKeysStatus status,
    const std::vector<std::vector<uint8_t>>& new_vault_keys,
    int last_vault_key_version) {
  // This method should be called only as a result of
  // `ongoing_request_` completion/failure, verify this
  // condition and destroy `ongoing_request_` as it's not
  // needed anymore.
  CHECK(ongoing_request_);
  ongoing_request_ = nullptr;

  RecordTrustedVaultDownloadKeysStatus(
      LocalRecoveryFactorType::kPhysicalDevice, security_domain_id_,
      GetDownloadKeysStatusForUMAFromResponse(status));

  RecoveryStatus recovery_status = RecoveryStatus::kFailure;
  switch (status) {
    case TrustedVaultDownloadKeysStatus::kSuccess: {
      recovery_status = RecoveryStatus::kSuccess;
      break;
    }
    case TrustedVaultDownloadKeysStatus::kMemberNotFound:
    case TrustedVaultDownloadKeysStatus::kMembershipNotFound:
    case TrustedVaultDownloadKeysStatus::kMembershipCorrupted:
    case TrustedVaultDownloadKeysStatus::kMembershipEmpty:
    case TrustedVaultDownloadKeysStatus::kKeyProofsVerificationFailed: {
      // Unable to download new keys due to known protocol errors. The only way
      // to go out of these states is to receive new vault keys through external
      // means. It's safe to mark device as not registered regardless of the
      // cause (device registration will be triggered once new vault keys are
      // available).
      MarkAsNotRegistered();
      break;
    }
    case TrustedVaultDownloadKeysStatus::kNoNewKeys: {
      // The registration itself exists, but there's no additional keys to
      // download. This is bad because key download attempts are triggered for
      // the case where local keys have been marked as stale, which means the
      // user is likely in an unrecoverable state.
      connection_->RecordFailedRequestForThrottling(primary_account_);
      recovery_status = RecoveryStatus::kNoNewKeys;
      break;
    }
    case TrustedVaultDownloadKeysStatus::kAccessTokenFetchingFailure:
    case TrustedVaultDownloadKeysStatus::kNetworkError:
      // Request wasn't sent to the server, so there is no need for throttling.
      break;
    case TrustedVaultDownloadKeysStatus::kOtherError:
      connection_->RecordFailedRequestForThrottling(primary_account_);
      break;
  }

  std::move(cb).Run(recovery_status, new_vault_keys, last_vault_key_version);
}

void PhysicalDeviceRecoveryFactor::FulfillRecoveryWithFailure(
    TrustedVaultDownloadKeysStatusForUMA status_for_uma,
    AttemptRecoveryCallback cb) {
  RecordTrustedVaultDownloadKeysStatus(LocalRecoveryFactorType::kPhysicalDevice,
                                       security_domain_id_, status_for_uma);

  base::BindPostTaskToCurrentDefault(
      base::BindOnce(std::move(cb), RecoveryStatus::kFailure,
                     std::vector<std::vector<uint8_t>>(), 0))
      .Run();
}

void PhysicalDeviceRecoveryFactor::OnRegistered(
    bool had_local_keys,
    TrustedVaultRegistrationStatus status,
    int key_version) {
  // This method should be called only as a result of
  // `ongoing_registration_request_` completion/failure, verify this
  // condition and destroy `ongoing_registration_request_` as it's not
  // needed anymore.
  CHECK(ongoing_registration_request_);
  ongoing_registration_request_ = nullptr;
  CHECK(ongoing_registration_callback_);
  RegisterCallback cb = std::move(ongoing_registration_callback_);

  switch (status) {
    case TrustedVaultRegistrationStatus::kRegistrationNotAttempted:
    case TrustedVaultRegistrationStatus::kRegistrationCancelled:
      NOTREACHED();
    case TrustedVaultRegistrationStatus::kSuccess:
    case TrustedVaultRegistrationStatus::kAlreadyRegistered:
      // kAlreadyRegistered handled as success, because it only means that
      // client doesn't fully handled successful device registration before.
      storage_->MutateLocalDeviceRegistrationInfo(
          primary_account_.gaia, [](LocalDeviceRegistrationInfo& info) {
            info.set_device_registered(true);
            info.set_device_registered_version(
                kCurrentDeviceRegistrationVersion);
          });
      storage_->SetLastRegistrationReturnedLocalDataObsolete(
          primary_account_.gaia, false);
      break;
    case TrustedVaultRegistrationStatus::kLocalDataObsolete:
      storage_->SetLastRegistrationReturnedLocalDataObsolete(
          primary_account_.gaia, true);
      break;
    case TrustedVaultRegistrationStatus::kTransientAccessTokenFetchError:
    case TrustedVaultRegistrationStatus::kPersistentAccessTokenFetchError:
    case TrustedVaultRegistrationStatus::
        kPrimaryAccountChangeAccessTokenFetchError:
    case TrustedVaultRegistrationStatus::kNetworkError:
    case TrustedVaultRegistrationStatus::kOtherError:
      break;
  }

  std::move(cb).Run(status, key_version, had_local_keys);
}

void PhysicalDeviceRecoveryFactor::FulfillRegistrationWithFailure(
    TrustedVaultRegistrationStatus status,
    RegisterCallback cb) {
  base::BindPostTaskToCurrentDefault(base::BindOnce(std::move(cb), status,
                                                    /*key_version=*/0,
                                                    /*had_local_keys=*/true))
      .Run();
}

}  // namespace trusted_vault