use AttestationClient to get/set key payloads.

We are deprecating attestation methods by CryptohomeClient.

BUG=b:158955123
TEST=unit_tests.

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

chrome/browser/chromeos/attestation/machine_certificate_uploader_impl.cc

@@ -19,6 +19,8 @@
 #include "chromeos/attestation/attestation_flow.h"
 #include "chromeos/cryptohome/async_method_caller.h"
 #include "chromeos/cryptohome/cryptohome_parameters.h"
+#include "chromeos/dbus/attestation/attestation_client.h"
+#include "chromeos/dbus/attestation/interface.pb.h"
 #include "chromeos/dbus/cryptohome/cryptohome_client.h"
 #include "chromeos/dbus/dbus_method_call_status.h"
 #include "chromeos/dbus/dbus_thread_manager.h"
@@ -266,12 +268,13 @@ void MachineCertificateUploaderImpl::CheckCertificateExpiry(
   if (num_certificates == 0) {
     LOG(WARNING) << "Failed to parse certificate chain, cannot check expiry.";
   }
-  // Get the payload and check if the certificate has already been uploaded.
-  GetKeyPayload(
-      base::BindRepeating(&MachineCertificateUploaderImpl::CheckIfUploaded,
-                          weak_factory_.GetWeakPtr(), pem_certificate_chain),
-      base::BindRepeating(&MachineCertificateUploaderImpl::Reschedule,
-                          weak_factory_.GetWeakPtr()));
+  ::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));
 }
 
 void MachineCertificateUploaderImpl::UploadCertificate(
@@ -284,9 +287,14 @@ void MachineCertificateUploaderImpl::UploadCertificate(
 
 void MachineCertificateUploaderImpl::CheckIfUploaded(
     const std::string& pem_certificate_chain,
-    const std::string& key_payload) {
+    const ::attestation::GetKeyInfoReply& reply) {
+  if (reply.status() != ::attestation::STATUS_SUCCESS) {
+    Reschedule();
+    return;
+  }
+
   AttestationKeyPayload payload_pb;
-  if (!key_payload.empty() && payload_pb.ParseFromString(key_payload) &&
+  if (!reply.payload().empty() && payload_pb.ParseFromString(reply.payload()) &&
       payload_pb.is_certificate_uploaded()) {
     // Already uploaded... nothing more to do.
     return;
@@ -294,45 +302,35 @@ void MachineCertificateUploaderImpl::CheckIfUploaded(
   UploadCertificate(pem_certificate_chain);
 }
 
-void MachineCertificateUploaderImpl::GetKeyPayload(
-    base::RepeatingCallback<void(const std::string&)> callback,
-    base::RepeatingCallback<void()> on_failure) {
-  cryptohome_client_->TpmAttestationGetKeyPayload(
-      KEY_DEVICE,
-      cryptohome::AccountIdentifier(),  // Not used.
-      kEnterpriseMachineKey,
-      base::BindOnce(DBusStringCallback, callback, on_failure, FROM_HERE));
-}
-
 void MachineCertificateUploaderImpl::OnUploadComplete(bool status) {
   if (status) {
     VLOG(1) << "Enterprise Machine Certificate uploaded to DMServer.";
-    GetKeyPayload(
-        base::BindRepeating(&MachineCertificateUploaderImpl::MarkAsUploaded,
-                            weak_factory_.GetWeakPtr()),
-        base::DoNothing());
+    ::attestation::GetKeyInfoRequest request;
+    request.set_username("");
+    request.set_key_label(kEnterpriseMachineKey);
+    AttestationClient::Get()->GetKeyInfo(
+        request, base::BindOnce(&MachineCertificateUploaderImpl::MarkAsUploaded,
+                                weak_factory_.GetWeakPtr()));
   }
   std::move(callback_).Run(status);
 }
 
 void MachineCertificateUploaderImpl::MarkAsUploaded(
-    const std::string& key_payload) {
+    const ::attestation::GetKeyInfoReply& reply) {
   AttestationKeyPayload payload_pb;
-  if (!key_payload.empty())
-    payload_pb.ParseFromString(key_payload);
+  if (!reply.payload().empty())
+    payload_pb.ParseFromString(reply.payload());
   payload_pb.set_is_certificate_uploaded(true);
   std::string new_payload;
   if (!payload_pb.SerializeToString(&new_payload)) {
     LOG(WARNING) << "Failed to serialize key payload.";
     return;
   }
-  cryptohome_client_->TpmAttestationSetKeyPayload(
-      KEY_DEVICE,
-      cryptohome::AccountIdentifier(),  // Not used.
-      kEnterpriseMachineKey, new_payload,
-      base::BindOnce(DBusBoolRedirectCallback, base::RepeatingClosure(),
-                     base::RepeatingClosure(), base::RepeatingClosure(),
-                     FROM_HERE));
+  ::attestation::SetKeyPayloadRequest request;
+  request.set_username("");
+  request.set_key_label(kEnterpriseMachineKey);
+  request.set_payload(new_payload);
+  AttestationClient::Get()->SetKeyPayload(request, base::DoNothing());
 }
 
 void MachineCertificateUploaderImpl::HandleGetCertificateFailure(

chrome/browser/chromeos/attestation/machine_certificate_uploader_impl.h

@@ -10,6 +10,7 @@
 #include "base/callback.h"
 #include "base/memory/weak_ptr.h"
 #include "chrome/browser/chromeos/attestation/machine_certificate_uploader.h"
+#include "chromeos/dbus/attestation/interface.pb.h"
 #include "chromeos/dbus/constants/attestation_constants.h"
 
 namespace policy {
@@ -71,19 +72,14 @@ class MachineCertificateUploaderImpl : public MachineCertificateUploader {
 
   // Checks if a certificate has already been uploaded and, if not, upload.
   void CheckIfUploaded(const std::string& pem_certificate_chain,
-                       const std::string& key_payload);
-
-  // Gets the payload associated with the EMK and sends it to |callback|,
-  // or call |on_failure| with no arguments if the payload cannot be obtained.
-  void GetKeyPayload(base::RepeatingCallback<void(const std::string&)> callback,
-                     base::RepeatingCallback<void()> on_failure);
+                       const ::attestation::GetKeyInfoReply& reply);
 
   // Called when a certificate upload operation completes.  On success, |status|
   // will be true.
   void OnUploadComplete(bool status);
 
   // Marks a key as uploaded in the payload proto.
-  void MarkAsUploaded(const std::string& key_payload);
+  void MarkAsUploaded(const ::attestation::GetKeyInfoReply& reply);
 
   // Handles failure of getting a certificate.
   void HandleGetCertificateFailure(AttestationStatus status);

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/dbus/attestation/fake_attestation_client.h"
 #include "chromeos/dbus/cryptohome/fake_cryptohome_client.h"
 #include "chromeos/settings/cros_settings_names.h"
 #include "components/policy/core/common/cloud/mock_cloud_policy_client.h"
@@ -38,6 +39,11 @@ const int64_t kCertExpiringSoon = 20;
 const int64_t kCertExpired = -20;
 
 void CertCallbackSuccess(AttestationFlow::CertificateCallback callback) {
+  AttestationClient::Get()
+      ->GetTestInterface()
+      ->GetMutableKeyInfoReply(/*username=*/"", kEnterpriseMachineKey)
+      ->set_certificate("fake_cert");
+
   base::ThreadTaskRunnerHandle::Get()->PostTask(
       FROM_HERE,
       base::BindOnce(std::move(callback), ATTESTATION_SUCCESS, "fake_cert"));
@@ -67,9 +73,11 @@ void StatusCallbackSuccess(policy::CloudPolicyClient::StatusCallback callback) {
 class MachineCertificateUploaderTest : public ::testing::TestWithParam<bool> {
  public:
   MachineCertificateUploaderTest() {
+    AttestationClient::InitializeFake();
     settings_helper_.ReplaceDeviceSettingsProviderWithStub();
     policy_client_.SetDMToken("fake_dm_token");
   }
+  ~MachineCertificateUploaderTest() override { AttestationClient::Shutdown(); }
 
  protected:
   enum MockOptions {
@@ -88,12 +96,20 @@ class MachineCertificateUploaderTest : public ::testing::TestWithParam<bool> {
     if (key_exists) {
       cryptohome_client_.SetTpmAttestationDeviceCertificate(
           kEnterpriseMachineKey, certificate);
+      ::attestation::GetKeyInfoReply* key_info =
+          AttestationClient::Get()->GetTestInterface()->GetMutableKeyInfoReply(
+              /*username=*/"", kEnterpriseMachineKey);
+      key_info->set_certificate(certificate);
     }
 
     // Setup expected key payload queries.
     bool key_uploaded = refresh || (mock_options & MOCK_KEY_UPLOADED);
-    cryptohome_client_.SetTpmAttestationDeviceKeyPayload(
-        kEnterpriseMachineKey, key_uploaded ? CreatePayload() : std::string());
+    if (key_uploaded) {
+      ::attestation::GetKeyInfoReply* key_info =
+          AttestationClient::Get()->GetTestInterface()->GetMutableKeyInfoReply(
+              /*username=*/"", kEnterpriseMachineKey);
+      key_info->set_payload(key_uploaded ? CreatePayload() : std::string());
+    }
 
     // Setup expected key uploads.  Use WillOnce() so StrictMock will trigger an
     // error if our expectations are not met exactly.  We want to verify that
@@ -165,8 +181,10 @@ TEST_P(MachineCertificateUploaderTest, NewCertificate) {
   SetupMocks(MOCK_NEW_KEY, "");
   Run();
   EXPECT_EQ(CreatePayload(),
-            cryptohome_client_.GetTpmAttestationDeviceKeyPayload(
-                kEnterpriseMachineKey));
+            AttestationClient::Get()
+                ->GetTestInterface()
+                ->GetMutableKeyInfoReply(/*username=*/"", kEnterpriseMachineKey)
+                ->payload());
 }
 
 TEST_P(MachineCertificateUploaderTest, KeyExistsNotUploaded) {
@@ -176,8 +194,10 @@ TEST_P(MachineCertificateUploaderTest, KeyExistsNotUploaded) {
   SetupMocks(MOCK_KEY_EXISTS, certificate);
   Run();
   EXPECT_EQ(CreatePayload(),
-            cryptohome_client_.GetTpmAttestationDeviceKeyPayload(
-                kEnterpriseMachineKey));
+            AttestationClient::Get()
+                ->GetTestInterface()
+                ->GetMutableKeyInfoReply(/*username=*/"", kEnterpriseMachineKey)
+                ->payload());
 }
 
 TEST_P(MachineCertificateUploaderTest, KeyExistsAlreadyUploaded) {
@@ -227,8 +247,10 @@ TEST_P(MachineCertificateUploaderTest, DBusFailureRetry) {
                      base::Unretained(&cryptohome_client_)));
   Run();
   EXPECT_EQ(CreatePayload(),
-            cryptohome_client_.GetTpmAttestationDeviceKeyPayload(
-                kEnterpriseMachineKey));
+            AttestationClient::Get()
+                ->GetTestInterface()
+                ->GetMutableKeyInfoReply(/*username=*/"", kEnterpriseMachineKey)
+                ->payload());
 }
 
 INSTANTIATE_TEST_SUITE_P(All,

chromeos/dbus/attestation/attestation_client.h

@@ -126,6 +126,11 @@ class COMPONENT_EXPORT(CHROMEOS_DBUS_ATTESTATION) AttestationClient {
     // Sets the returned status of `GetEnrollmentId()` to be D-Bus error for
     // `count` times to emulate D-Bus late availability.
     virtual void set_enrollment_id_dbus_error_count(int count) = 0;
+
+    // Gets the reply to the key info query as a fake database.
+    virtual ::attestation::GetKeyInfoReply* GetMutableKeyInfoReply(
+        const std::string& username,
+        const std::string& label) = 0;
   };
 
   // Not copyable or movable.

chromeos/dbus/attestation/fake_attestation_client.cc

@@ -57,7 +57,15 @@ FakeAttestationClient::mutable_certificate_request_reply() {
 void FakeAttestationClient::GetKeyInfo(
     const ::attestation::GetKeyInfoRequest& request,
     GetKeyInfoCallback callback) {
-  NOTIMPLEMENTED();
+  ::attestation::GetKeyInfoReply reply;
+  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);
 }
 
 void FakeAttestationClient::GetEndorsementInfo(
@@ -228,7 +236,17 @@ void FakeAttestationClient::SignSimpleChallenge(
 void FakeAttestationClient::SetKeyPayload(
     const ::attestation::SetKeyPayloadRequest& request,
     SetKeyPayloadCallback callback) {
-  NOTIMPLEMENTED();
+  ::attestation::GetKeyInfoRequest get_key_info_request;
+  get_key_info_request.set_username(request.username());
+  get_key_info_request.set_key_label(request.key_label());
+  auto iter = key_info_database_.find(get_key_info_request);
+  ::attestation::SetKeyPayloadReply reply;
+  if (iter == key_info_database_.end()) {
+    reply.set_status(::attestation::STATUS_INVALID_PARAMETER);
+  } else {
+    iter->second.set_payload(request.payload());
+  }
+  PostProtoResponse(std::move(callback), reply);
 }
 
 void FakeAttestationClient::DeleteKeys(
@@ -330,6 +348,16 @@ void FakeAttestationClient::set_enrollment_id_dbus_error_count(int count) {
   enrollment_id_dbus_error_count_ = count;
 }
 
+::attestation::GetKeyInfoReply* FakeAttestationClient::GetMutableKeyInfoReply(
+    const std::string& username,
+    const std::string& label) {
+  ::attestation::GetKeyInfoRequest request;
+  request.set_username(username);
+  request.set_key_label(label);
+  // If there doesn't exist the entry yet, just create a new one.
+  return &(key_info_database_[request]);
+}
+
 AttestationClient::TestInterface* FakeAttestationClient::GetTestInterface() {
   return this;
 }

chromeos/dbus/attestation/fake_attestation_client.h

@@ -8,6 +8,7 @@
 #include "chromeos/dbus/attestation/attestation_client.h"
 
 #include <deque>
+#include <map>
 #include <string>
 #include <vector>
 
@@ -113,6 +114,9 @@ class COMPONENT_EXPORT(CHROMEOS_DBUS_ATTESTATION) FakeAttestationClient
   void set_enrollment_id_ignore_cache(const std::string& id) override;
   void set_cached_enrollment_id(const std::string& id) override;
   void set_enrollment_id_dbus_error_count(int count) override;
+  ::attestation::GetKeyInfoReply* GetMutableKeyInfoReply(
+      const std::string& username,
+      const std::string& label) override;
 
   AttestationClient::TestInterface* GetTestInterface() override;
 
@@ -145,6 +149,22 @@ class COMPONENT_EXPORT(CHROMEOS_DBUS_ATTESTATION) FakeAttestationClient
   std::string enrollment_id_;
   std::string enrollment_id_ignore_cache_;
   int enrollment_id_dbus_error_count_ = 0;
+
+  class GetKeyInfoRequestComparator {
+   public:
+    bool operator()(const ::attestation::GetKeyInfoRequest& r1,
+                    const ::attestation::GetKeyInfoRequest& r2) const {
+      return r1.username() == r2.username() ? r1.key_label() < r2.key_label()
+                                            : r1.username() < r2.username();
+    }
+  };
+  // The fake key info database. std::map is chosen because we don't have to
+  // implement the hash function for the `GetKeyInfoRequest`, which could be
+  // expensive and contributes unreasonable overhead at smaller scale, anyway.
+  std::map<::attestation::GetKeyInfoRequest,
+           ::attestation::GetKeyInfoReply,
+           GetKeyInfoRequestComparator>
+      key_info_database_;
 };
 
 }  // namespace chromeos