use AttestationClient to get key info in MachineCertificateUploader.

We are deprecating attestation methods by CryptohomeClient.

This CL also includes a simplied flow by consolidating
TpmAttestationDoesKeyExist and TpmAttestationGetCertificate, of which
results come from the same methods of attestation service anyway.

BUG=b:158955123
TEST=unit_tests.

Change-Id: I36933b4ee20e86fd6946644d957ab70459609482
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2515801
Commit-Queue: Leo Lai <cylai@google.com>
Reviewed-by: Maksim Ivanov <emaxx@chromium.org>
Cr-Commit-Position: refs/heads/master@{#823899}

chrome/browser/chromeos/attestation/machine_certificate_uploader_impl.cc

@@ -44,51 +44,6 @@ const int kExpiryThresholdInDays = 30;
 const int kRetryDelay = 5;  // Seconds.
 const int kRetryLimit = 100;
 
-// A dbus callback which handles a boolean result.
-//
-// Parameters
-//   on_true - Called when status=success and value=true.
-//   on_false - Called when status=success and value=false.
-//   status - The dbus operation status.
-//   result - The value returned by the dbus operation.
-void DBusBoolRedirectCallback(const base::RepeatingClosure& on_true,
-                              const base::RepeatingClosure& on_false,
-                              const base::RepeatingClosure& on_failure,
-                              const base::Location& from_here,
-                              base::Optional<bool> result) {
-  if (!result.has_value()) {
-    LOG(ERROR) << "Cryptohome DBus method failed: " << from_here.ToString();
-    if (!on_failure.is_null())
-      on_failure.Run();
-    return;
-  }
-  const base::RepeatingClosure& task = result.value() ? on_true : on_false;
-  if (!task.is_null())
-    task.Run();
-}
-
-// A dbus callback which handles a string result.
-//
-// Parameters
-//   on_success - Called when status=success and result=true.
-//   status - The dbus operation status.
-//   result - The result returned by the dbus operation.
-//   data - The data returned by the dbus operation.
-void DBusStringCallback(
-    const base::RepeatingCallback<void(const std::string&)> on_success,
-    const base::RepeatingClosure& on_failure,
-    const base::Location& from_here,
-    base::Optional<chromeos::CryptohomeClient::TpmAttestationDataResult>
-        result) {
-  if (!result.has_value() || !result->success) {
-    LOG(ERROR) << "Cryptohome DBus method failed: " << from_here.ToString();
-    if (!on_failure.is_null())
-      on_failure.Run();
-    return;
-  }
-  on_success.Run(result->data);
-}
-
 void DBusPrivacyCACallback(
     const base::RepeatingCallback<void(const std::string&)> on_success,
     const base::RepeatingCallback<
@@ -181,22 +136,13 @@ void MachineCertificateUploaderImpl::Start() {
     return;
   }
 
-  // Start a dbus call to check if an Enterprise Machine Key already exists.
-  base::RepeatingClosure on_does_not_exist =
-      base::BindRepeating(&MachineCertificateUploaderImpl::GetNewCertificate,
-                          weak_factory_.GetWeakPtr());
-  base::RepeatingClosure on_does_exist = base::BindRepeating(
-      &MachineCertificateUploaderImpl::GetExistingCertificate,
-      weak_factory_.GetWeakPtr());
-  cryptohome_client_->TpmAttestationDoesKeyExist(
-      KEY_DEVICE,
-      cryptohome::AccountIdentifier(),  // Not used.
-      kEnterpriseMachineKey,
-      base::BindOnce(
-          DBusBoolRedirectCallback, on_does_exist, on_does_not_exist,
-          base::BindRepeating(&MachineCertificateUploaderImpl::Reschedule,
-                              weak_factory_.GetWeakPtr()),
-          FROM_HERE));
+  ::attestation::GetKeyInfoRequest request;
+  request.set_username("");
+  request.set_key_label(kEnterpriseMachineKey);
+  AttestationClient::Get()->GetKeyInfo(
+      request,
+      base::BindOnce(&MachineCertificateUploaderImpl::OnGetExistingCertificate,
+                     weak_factory_.GetWeakPtr()));
 }
 
 void MachineCertificateUploaderImpl::GetNewCertificate() {
@@ -224,23 +170,26 @@ void MachineCertificateUploaderImpl::GetNewCertificate() {
           FROM_HERE));
 }
 
-void MachineCertificateUploaderImpl::GetExistingCertificate() {
-  cryptohome_client_->TpmAttestationGetCertificate(
-      KEY_DEVICE,
-      cryptohome::AccountIdentifier(),  // Not used.
-      kEnterpriseMachineKey,
-      base::BindOnce(
-          DBusStringCallback,
-          base::BindRepeating(
-              &MachineCertificateUploaderImpl::CheckCertificateExpiry,
-              weak_factory_.GetWeakPtr()),
-          base::BindRepeating(&MachineCertificateUploaderImpl::Reschedule,
-                              weak_factory_.GetWeakPtr()),
-          FROM_HERE));
+void MachineCertificateUploaderImpl::OnGetExistingCertificate(
+    const ::attestation::GetKeyInfoReply& reply) {
+  if (reply.status() != ::attestation::STATUS_SUCCESS &&
+      reply.status() != ::attestation::STATUS_INVALID_PARAMETER) {
+    LOG(ERROR) << "Error getting the existing certificate; status: "
+               << reply.status();
+    Reschedule();
+    return;
+  }
+  // Get a new certificate if not exists.
+  if (reply.status() == ::attestation::STATUS_INVALID_PARAMETER) {
+    GetNewCertificate();
+    return;
+  }
+  CheckCertificateExpiry(reply);
 }
 
 void MachineCertificateUploaderImpl::CheckCertificateExpiry(
-    const std::string& pem_certificate_chain) {
+    const ::attestation::GetKeyInfoReply& reply) {
+  const std::string& pem_certificate_chain = reply.certificate();
   int num_certificates = 0;
   net::PEMTokenizer pem_tokenizer(pem_certificate_chain, {"CERTIFICATE"});
   while (pem_tokenizer.GetNext()) {
@@ -270,13 +219,7 @@ void MachineCertificateUploaderImpl::CheckCertificateExpiry(
   if (num_certificates == 0) {
     LOG(WARNING) << "Failed to parse certificate chain, cannot check expiry.";
   }
-  ::attestation::GetKeyInfoRequest request;
-  request.set_username("");
-  request.set_key_label(kEnterpriseMachineKey);
-  AttestationClient::Get()->GetKeyInfo(
-      request,
-      base::BindOnce(&MachineCertificateUploaderImpl::CheckIfUploaded,
-                     weak_factory_.GetWeakPtr(), pem_certificate_chain));
+  CheckIfUploaded(reply);
 }
 
 void MachineCertificateUploaderImpl::UploadCertificate(
@@ -289,9 +232,12 @@ void MachineCertificateUploaderImpl::UploadCertificate(
 }
 
 void MachineCertificateUploaderImpl::CheckIfUploaded(
-    const std::string& pem_certificate_chain,
     const ::attestation::GetKeyInfoReply& reply) {
+  // The caller in the entire flow should have checked the reply status already.
   if (reply.status() != ::attestation::STATUS_SUCCESS) {
+    LOG(DFATAL) << "Checking key payload in a bad reply from attestation "
+                   "service; status: "
+                << reply.status();
     Reschedule();
     return;
   }
@@ -304,7 +250,7 @@ void MachineCertificateUploaderImpl::CheckIfUploaded(
     RunCallbacks(certificate_uploaded_.value());
     return;
   }
-  UploadCertificate(pem_certificate_chain);
+  UploadCertificate(reply.certificate());
 }
 
 void MachineCertificateUploaderImpl::OnUploadComplete(bool status) {
@@ -333,6 +279,10 @@ void MachineCertificateUploaderImpl::WaitForUploadComplete(
 
 void MachineCertificateUploaderImpl::MarkAsUploaded(
     const ::attestation::GetKeyInfoReply& reply) {
+  if (reply.status() != ::attestation::STATUS_SUCCESS) {
+    LOG(WARNING) << "Failed to get existing payload.";
+    return;
+  }
   AttestationKeyPayload payload_pb;
   if (!reply.payload().empty())
     payload_pb.ParseFromString(reply.payload());

chrome/browser/chromeos/attestation/machine_certificate_uploader_impl.h

@@ -67,19 +67,23 @@ class MachineCertificateUploaderImpl : public MachineCertificateUploader {
   // Gets a new certificate for the Enterprise Machine Key (EMK).
   void GetNewCertificate();
 
-  // Gets the existing EMK certificate and sends it to CheckCertificateExpiry.
-  void GetExistingCertificate();
+  // Called when `GetKeyInfo()` returned to check the existing certificate.
+  // There are 3 cases by the status replied from attestation service:
+  //     1. If the existing EMK is found, calls `CheckCertificateExpiry()`.
+  //     2. If the key does not exist, calls `GetNewCertificate()`.
+  //     3. Otherwise, there is an error and `Reschedule()` is called to retry.
+  void OnGetExistingCertificate(const ::attestation::GetKeyInfoReply& reply);
 
-  // Checks if any certificate in the given pem_certificate_chain is expired
-  // and, if so, gets a new one. If not renewing, calls CheckIfUploaded.
-  void CheckCertificateExpiry(const std::string& pem_certificate_chain);
+  // Checks if any certificate in the given `reply` is expired and, if so, gets
+  // a new one. If not renewing, calls `CheckIfUploaded()`.
+  void CheckCertificateExpiry(const ::attestation::GetKeyInfoReply& reply);
 
   // Uploads a machine certificate to the policy server.
   void UploadCertificate(const std::string& pem_certificate_chain);
 
-  // Checks if a certificate has already been uploaded and, if not, upload.
-  void CheckIfUploaded(const std::string& pem_certificate_chain,
-                       const ::attestation::GetKeyInfoReply& reply);
+  // Checks if a certificate in `reply` has already been uploaded and, if not,
+  // upload.
+  void CheckIfUploaded(const ::attestation::GetKeyInfoReply& reply);
 
   // Called when a certificate upload operation completes.  On success, |status|
   // will be true.

chrome/browser/chromeos/attestation/machine_certificate_uploader_impl_unittest.cc

@@ -17,6 +17,7 @@
 #include "chrome/browser/chromeos/attestation/machine_certificate_uploader_impl.h"
 #include "chrome/browser/chromeos/settings/scoped_cros_settings_test_helper.h"
 #include "chromeos/attestation/mock_attestation_flow.h"
+#include "chromeos/cryptohome/cryptohome_parameters.h"
 #include "chromeos/dbus/attestation/fake_attestation_client.h"
 #include "chromeos/dbus/cryptohome/fake_cryptohome_client.h"
 #include "chromeos/settings/cros_settings_names.h"
@@ -34,9 +35,11 @@ namespace attestation {
 
 namespace {
 
-const int64_t kCertValid = 90;
-const int64_t kCertExpiringSoon = 20;
-const int64_t kCertExpired = -20;
+constexpr int64_t kCertValid = 90;
+constexpr int64_t kCertExpiringSoon = 20;
+constexpr int64_t kCertExpired = -20;
+constexpr int kRetryLimit = 3;
+constexpr char kFakeCertificate[] = "fake_cert";
 
 void CertCallbackSuccess(AttestationFlow::CertificateCallback callback) {
   AttestationClient::Get()
@@ -77,16 +80,50 @@ class CallbackObserver {
   }
 };
 
+class MockableFakeAttestationFlow : public MockAttestationFlow {
+ public:
+  MockableFakeAttestationFlow() {
+    ON_CALL(*this, GetCertificate(_, _, _, _, _, _))
+        .WillByDefault(
+            Invoke(this, &MockableFakeAttestationFlow::GetCertificateInternal));
+  }
+  ~MockableFakeAttestationFlow() override = default;
+  void set_status(AttestationStatus status) { status_ = status; }
+
+ private:
+  void GetCertificateInternal(AttestationCertificateProfile certificate_profile,
+                              const AccountId& account_id,
+                              const std::string& request_origin,
+                              bool force_new_key,
+                              const std::string& key_name,
+                              CertificateCallback callback) {
+    std::string certificate;
+    if (status_ == ATTESTATION_SUCCESS) {
+      certificate = certificate_;
+      AttestationClient::Get()
+          ->GetTestInterface()
+          ->GetMutableKeyInfoReply(cryptohome::Identification(account_id).id(),
+                                   kEnterpriseMachineKey)
+          ->set_certificate(certificate_);
+    }
+    std::move(callback).Run(status_, certificate);
+  }
+  AttestationStatus status_ = ATTESTATION_SUCCESS;
+  const std::string certificate_ = kFakeCertificate;
+};
+
 }  // namespace
 
-class MachineCertificateUploaderTest : public ::testing::TestWithParam<bool> {
+class MachineCertificateUploaderTestBase : public ::testing::Test {
  public:
-  MachineCertificateUploaderTest() {
+  MachineCertificateUploaderTestBase() {
     AttestationClient::InitializeFake();
     settings_helper_.ReplaceDeviceSettingsProviderWithStub();
     policy_client_.SetDMToken("fake_dm_token");
   }
-  ~MachineCertificateUploaderTest() override { AttestationClient::Shutdown(); }
+  ~MachineCertificateUploaderTestBase() override {
+    AttestationClient::Shutdown();
+  }
 
  protected:
   enum MockOptions {
@@ -95,7 +132,9 @@ class MachineCertificateUploaderTest : public ::testing::TestWithParam<bool> {
     MOCK_NEW_KEY = (1 << 2)        // Configure expecting new key generation.
   };
 
-  bool GetShouldRefreshCert() { return GetParam(); }
+  // The derived fixture has different needs to control this function's
+  // behavior.
+  virtual bool GetShouldRefreshCert() const = 0;
 
   // Configures mock expectations according to |mock_options|.  If options
   // require that a certificate exists, |certificate| will be used.
@@ -103,8 +142,6 @@ class MachineCertificateUploaderTest : public ::testing::TestWithParam<bool> {
     bool key_exists = (mock_options & MOCK_KEY_EXISTS);
     // Setup expected key / cert queries.
     if (key_exists) {
-      cryptohome_client_.SetTpmAttestationDeviceCertificate(
-          kEnterpriseMachineKey, certificate);
       ::attestation::GetKeyInfoReply* key_info =
           AttestationClient::Get()->GetTestInterface()->GetMutableKeyInfoReply(
               /*username=*/"", kEnterpriseMachineKey);
@@ -112,8 +149,7 @@ class MachineCertificateUploaderTest : public ::testing::TestWithParam<bool> {
     }
 
     // Setup expected key payload queries.
-    bool key_uploaded =
-        GetShouldRefreshCert() || (mock_options & MOCK_KEY_UPLOADED);
+    bool key_uploaded = (mock_options & MOCK_KEY_UPLOADED);
     if (key_uploaded) {
       ::attestation::GetKeyInfoReply* key_info =
           AttestationClient::Get()->GetTestInterface()->GetMutableKeyInfoReply(
@@ -128,8 +164,9 @@ class MachineCertificateUploaderTest : public ::testing::TestWithParam<bool> {
     // status in the key payload matches the upload operation.
     bool new_key = GetShouldRefreshCert() || (mock_options & MOCK_NEW_KEY);
     if (new_key || !key_uploaded) {
-      EXPECT_CALL(policy_client_, UploadEnterpriseMachineCertificate(
-                                      new_key ? "fake_cert" : certificate, _))
+      EXPECT_CALL(policy_client_,
+                  UploadEnterpriseMachineCertificate(
+                      new_key ? kFakeCertificate : certificate, _))
           .WillOnce(WithArgs<1>(Invoke(StatusCallbackSuccess)));
     }
 
@@ -137,15 +174,14 @@ class MachineCertificateUploaderTest : public ::testing::TestWithParam<bool> {
     // another costly operation and if it gets triggered more than once during
     // a single pass this indicates a logical problem in the observer.
     if (new_key) {
-      EXPECT_CALL(attestation_flow_, GetCertificate(_, _, _, _, _, _))
-          .WillOnce(WithArgs<5>(Invoke(CertCallbackSuccess)));
+      EXPECT_CALL(attestation_flow_, GetCertificate(_, _, _, _, _, _));
     }
   }
 
-  void Run() {
+  void RunUploader() {
     MachineCertificateUploaderImpl uploader(
         &policy_client_, &cryptohome_client_, &attestation_flow_);
-    uploader.set_retry_limit(3);
+    uploader.set_retry_limit(kRetryLimit);
     uploader.set_retry_delay(0);
     if (GetShouldRefreshCert())
       uploader.RefreshAndUploadCertificate(base::DoNothing());
@@ -166,30 +202,37 @@ class MachineCertificateUploaderTest : public ::testing::TestWithParam<bool> {
   content::BrowserTaskEnvironment task_environment_;
   ScopedCrosSettingsTestHelper settings_helper_;
   FakeCryptohomeClient cryptohome_client_;
-  StrictMock<MockAttestationFlow> attestation_flow_;
+  StrictMock<MockableFakeAttestationFlow> attestation_flow_;
   StrictMock<policy::MockCloudPolicyClient> policy_client_;
 };
 
+class MachineCertificateUploaderTest
+    : public MachineCertificateUploaderTestBase,
+      public ::testing::WithParamInterface<bool> {
+ public:
+  bool GetShouldRefreshCert() const final { return GetParam(); }
+};
+
 TEST_P(MachineCertificateUploaderTest, UnregisteredPolicyClient) {
   policy_client_.SetDMToken("");
-  Run();
+  RunUploader();
 }
 
 TEST_P(MachineCertificateUploaderTest, GetCertificateUnspecifiedFailure) {
   EXPECT_CALL(attestation_flow_, GetCertificate(_, _, _, _, _, _))
       .WillRepeatedly(WithArgs<5>(Invoke(CertCallbackUnspecifiedFailure)));
-  Run();
+  RunUploader();
 }
 
 TEST_P(MachineCertificateUploaderTest, GetCertificateBadRequestFailure) {
   EXPECT_CALL(attestation_flow_, GetCertificate(_, _, _, _, _, _))
       .WillOnce(WithArgs<5>(Invoke(CertCallbackBadRequestFailure)));
-  Run();
+  RunUploader();
 }
 
 TEST_P(MachineCertificateUploaderTest, NewCertificate) {
   SetupMocks(MOCK_NEW_KEY, "");
-  Run();
+  RunUploader();
   EXPECT_EQ(CreatePayload(),
             AttestationClient::Get()
                 ->GetTestInterface()
@@ -266,7 +309,7 @@ TEST_P(MachineCertificateUploaderTest, KeyExistsNotUploaded) {
   ASSERT_TRUE(GetFakeCertificatePEM(base::TimeDelta::FromDays(kCertValid),
                                     &certificate));
   SetupMocks(MOCK_KEY_EXISTS, certificate);
-  Run();
+  RunUploader();
   EXPECT_EQ(CreatePayload(),
             AttestationClient::Get()
                 ->GetTestInterface()
@@ -279,7 +322,7 @@ TEST_P(MachineCertificateUploaderTest, KeyExistsAlreadyUploaded) {
   ASSERT_TRUE(GetFakeCertificatePEM(base::TimeDelta::FromDays(kCertValid),
                                     &certificate));
   SetupMocks(MOCK_KEY_EXISTS | MOCK_KEY_UPLOADED, certificate);
-  Run();
+  RunUploader();
 }
 
 TEST_P(MachineCertificateUploaderTest, KeyExistsCertExpiringSoon) {
@@ -287,7 +330,7 @@ TEST_P(MachineCertificateUploaderTest, KeyExistsCertExpiringSoon) {
   ASSERT_TRUE(GetFakeCertificatePEM(
       base::TimeDelta::FromDays(kCertExpiringSoon), &certificate));
   SetupMocks(MOCK_KEY_EXISTS | MOCK_KEY_UPLOADED | MOCK_NEW_KEY, certificate);
-  Run();
+  RunUploader();
 }
 
 TEST_P(MachineCertificateUploaderTest, KeyExistsCertExpired) {
@@ -295,41 +338,38 @@ TEST_P(MachineCertificateUploaderTest, KeyExistsCertExpired) {
   ASSERT_TRUE(GetFakeCertificatePEM(base::TimeDelta::FromDays(kCertExpired),
                                     &certificate));
   SetupMocks(MOCK_KEY_EXISTS | MOCK_KEY_UPLOADED | MOCK_NEW_KEY, certificate);
-  Run();
+  RunUploader();
 }
 
 TEST_P(MachineCertificateUploaderTest, IgnoreUnknownCertFormat) {
   SetupMocks(MOCK_KEY_EXISTS | MOCK_KEY_UPLOADED, "unsupported");
-  Run();
+  RunUploader();
 }
 
-TEST_P(MachineCertificateUploaderTest, DBusFailureRetry) {
+INSTANTIATE_TEST_SUITE_P(All,
+                         MachineCertificateUploaderTest,
+                         testing::Values(false, true));
+
+class MachineCertificateUploaderTestNoRefresh
+    : public MachineCertificateUploaderTestBase {
+ public:
+  bool GetShouldRefreshCert() const final { return false; }
+};
+
+TEST_F(MachineCertificateUploaderTestNoRefresh, DBusFailureRetry) {
   SetupMocks(MOCK_NEW_KEY, "");
-  // Simulate a DBus failure.
-  cryptohome_client_.SetServiceIsAvailable(false);
-
-  // Emulate delayed service initialization.
-  // Run() instantiates an Observer, which synchronously calls
-  // TpmAttestationDoesKeyExist() and fails. During this call, we make the
-  // service available in the next run, so on retry, it will successfully
-  // return the result.
-  base::ThreadTaskRunnerHandle::Get()->PostTask(
-      FROM_HERE, base::BindOnce(
-                     [](FakeCryptohomeClient* cryptohome_client) {
-                       cryptohome_client->SetServiceIsAvailable(true);
-                     },
-                     base::Unretained(&cryptohome_client_)));
-  Run();
+  AttestationClient::Get()->GetTestInterface()->set_key_info_dbus_error_count(
+      kRetryLimit - 1);
+  RunUploader();
   EXPECT_EQ(CreatePayload(),
             AttestationClient::Get()
                 ->GetTestInterface()
                 ->GetMutableKeyInfoReply(/*username=*/"", kEnterpriseMachineKey)
                 ->payload());
+  EXPECT_EQ(
+      AttestationClient::Get()->GetTestInterface()->key_info_dbus_error_count(),
+      0);
 }
 
-INSTANTIATE_TEST_SUITE_P(All,
-                         MachineCertificateUploaderTest,
-                         testing::Values(false, true));
-
 }  // namespace attestation
 }  // namespace chromeos

chromeos/dbus/attestation/attestation_client.h

@@ -135,6 +135,11 @@ class COMPONENT_EXPORT(CHROMEOS_DBUS_ATTESTATION) AttestationClient {
     virtual ::attestation::GetKeyInfoReply* GetMutableKeyInfoReply(
         const std::string& username,
         const std::string& label) = 0;
+    // Sets the returned status of `GetKeyInfo()` to be D-Bus error for
+    // `count` times to emulate D-Bus late availability.
+    virtual void set_key_info_dbus_error_count(int count) = 0;
+    // Gets the remaining count of `GetKeyInfo()` replying D-Bus error.
+    virtual int key_info_dbus_error_count() const = 0;
 
     // Verifies if `signed_data` is signed against `challenge`.
     virtual bool VerifySimpleChallengeResponse(

chromeos/dbus/attestation/fake_attestation_client.cc

@@ -76,12 +76,17 @@ void FakeAttestationClient::GetKeyInfo(
     const ::attestation::GetKeyInfoRequest& request,
     GetKeyInfoCallback callback) {
   ::attestation::GetKeyInfoReply reply;
+  if (key_info_dbus_error_count_ > 0) {
+    reply.set_status(::attestation::STATUS_DBUS_ERROR);
+    key_info_dbus_error_count_--;
+  } else {
     auto iter = key_info_database_.find(request);
     if (iter != key_info_database_.end()) {
       reply = iter->second;
     } else {
       reply.set_status(::attestation::STATUS_INVALID_PARAMETER);
     }
+  }
 
   PostProtoResponse(std::move(callback), reply);
 }
@@ -413,6 +418,14 @@ void FakeAttestationClient::set_enrollment_id_dbus_error_count(int count) {
   return &(key_info_database_[request]);
 }
 
+void FakeAttestationClient::set_key_info_dbus_error_count(int count) {
+  key_info_dbus_error_count_ = count;
+}
+
+int FakeAttestationClient::key_info_dbus_error_count() const {
+  return key_info_dbus_error_count_;
+}
+
 bool FakeAttestationClient::VerifySimpleChallengeResponse(
     const std::string& challenge,
     const ::attestation::SignedData& signed_data) {

chromeos/dbus/attestation/fake_attestation_client.h

@@ -121,6 +121,8 @@ class COMPONENT_EXPORT(CHROMEOS_DBUS_ATTESTATION) FakeAttestationClient
   ::attestation::GetKeyInfoReply* GetMutableKeyInfoReply(
       const std::string& username,
       const std::string& label) override;
+  void set_key_info_dbus_error_count(int count) override;
+  int key_info_dbus_error_count() const override;
   bool VerifySimpleChallengeResponse(
       const std::string& challenge,
       const ::attestation::SignedData& signed_data) override;
@@ -189,6 +191,7 @@ class COMPONENT_EXPORT(CHROMEOS_DBUS_ATTESTATION) FakeAttestationClient
            ::attestation::GetKeyInfoReply,
            GetKeyInfoRequestComparator>
       key_info_database_;
+  int key_info_dbus_error_count_ = 0;
 
   // The status returned by `SignSimpleChallenge()`.
   ::attestation::AttestationStatus sign_simple_challenge_status_ =