Enhanced and refactored the AttestationFlow interface.

chromeos/attestation/attestation_flow.cc

Index: chromeos/attestation/attestation_flow.cc
diff --git a/chromeos/attestation/attestation_flow.cc b/chromeos/attestation/attestation_flow.cc
index 97842a37f292e748a91e9d2a918193c429cc66c2..2fba76152b95cea5e5d9c4d1579c1ae0bcb5c915 100644
--- a/chromeos/attestation/attestation_flow.cc
+++ b/chromeos/attestation/attestation_flow.cc
@@ -38,30 +38,82 @@ void DBusBoolRedirectCallback(const base::Closure& on_true,
     task.Run();
 }
 
-}  // namespace
+void DBusDataMethodCallback(
+    const AttestationFlow::CertificateCallback& callback,
+    DBusMethodCallStatus status,
+    bool result,
+    const std::string& data) {
+  if (status != DBUS_METHOD_CALL_SUCCESS) {
+    LOG(ERROR) << "Attestation: DBus data operation failed.";
+    if (!callback.is_null())
+      callback.Run(false, "");
+    return;
+  }
+  if (!callback.is_null())
+    callback.Run(result, data);
+}
+
+AttestationKeyType GetKeyTypeForProfile(
+    AttestationCertificateProfile profile) {
+  switch (profile) {
+    case PROFILE_ENTERPRISE_MACHINE_CERTIFICATE:
+      return KEY_DEVICE;
+    case PROFILE_ENTERPRISE_USER_CERTIFICATE:
+      return KEY_USER;
+  }
+  NOTREACHED();
+  return KEY_USER;
+}
 
-const char AttestationFlow::kEnterpriseMachineKey[] = "attest-ent-machine";
+std::string GetKeyNameForProfile(
+    AttestationCertificateProfile profile) {
+  switch (profile) {
+    case PROFILE_ENTERPRISE_MACHINE_CERTIFICATE:
+      return kEnterpriseMachineKey;
+    case PROFILE_ENTERPRISE_USER_CERTIFICATE:
+      return kEnterpriseUserKey;
+  }
+  NOTREACHED();
+  return "";
+}
+
+int GetCertificateOptionsForProfile(
+    AttestationCertificateProfile profile) {
+  switch (profile) {
+    case PROFILE_ENTERPRISE_MACHINE_CERTIFICATE:
+      return CERTIFICATE_INCLUDE_STABLE_ID | CERTIFICATE_INCLUDE_DEVICE_STATE;
+    case PROFILE_ENTERPRISE_USER_CERTIFICATE:
+      return CERTIFICATE_INCLUDE_DEVICE_STATE;
+  }
+  NOTREACHED();
+  return CERTIFICATE_OPTION_NONE;
+}
+
+}  // namespace
 
 AttestationFlow::AttestationFlow(cryptohome::AsyncMethodCaller* async_caller,
                                  CryptohomeClient* cryptohome_client,
                                  scoped_ptr<ServerProxy> server_proxy)
-    : weak_factory_(this),
-      async_caller_(async_caller),
+    : async_caller_(async_caller),
       cryptohome_client_(cryptohome_client),
-      server_proxy_(server_proxy.Pass()) {
+      server_proxy_(server_proxy.Pass()),
+      weak_factory_(this) {
 }
 
 AttestationFlow::~AttestationFlow() {
 }
 
-void AttestationFlow::GetCertificate(const std::string& name,
-                                     const CertificateCallback& callback) {
+void AttestationFlow::GetCertificate(
+    AttestationCertificateProfile certificate_profile,
+    bool force_new_key,
+    const CertificateCallback& callback) {
   // If this device has not enrolled with the Privacy CA, we need to do that
   // first.  Once enrolled we can proceed with the certificate request.
   base::Closure do_cert_request = base::Bind(
       &AttestationFlow::StartCertificateRequest,
       weak_factory_.GetWeakPtr(),
-      name,
+      certificate_profile,
+      force_new_key,
       callback);
   base::Closure on_enroll_failure = base::Bind(callback, false, "");
   base::Closure do_enroll = base::Bind(&AttestationFlow::StartEnroll,
@@ -143,22 +195,49 @@ void AttestationFlow::OnEnrollComplete(const base::Closure& on_failure,
 }
 
 void AttestationFlow::StartCertificateRequest(
-    const std::string& name,
+    AttestationCertificateProfile certificate_profile,
+    bool generate_new_key,
     const CertificateCallback& callback) {
-  // Get the attestation service to create a Privacy CA certificate request.
-  int options = CryptohomeClient::INCLUDE_DEVICE_STATE;
-  if (name == kEnterpriseMachineKey)
-    options |= CryptohomeClient::INCLUDE_STABLE_ID;
-  async_caller_->AsyncTpmAttestationCreateCertRequest(
-      options,
-      base::Bind(&AttestationFlow::SendCertificateRequestToPCA,
-                 weak_factory_.GetWeakPtr(),
-                 name,
-                 callback));
+  AttestationKeyType key_type = GetKeyTypeForProfile(certificate_profile);
+  std::string key_name = GetKeyNameForProfile(certificate_profile);
+  if (generate_new_key) {
+    // Get the attestation service to create a Privacy CA certificate request.
+    async_caller_->AsyncTpmAttestationCreateCertRequest(
+        GetCertificateOptionsForProfile(certificate_profile),
+        base::Bind(&AttestationFlow::SendCertificateRequestToPCA,
+                   weak_factory_.GetWeakPtr(),
+                   key_type,
+                   key_name,
+                   callback));
+  } else {
+    // If the key already exists, query the existing certificate.
+    base::Closure on_key_exists = base::Bind(
+        &AttestationFlow::GetExistingCertificate,
+        weak_factory_.GetWeakPtr(),
+        key_type,
+        key_name,
+        callback);
+    // If the key does not exist, call this method back with |generate_new_key|
+    // set to true.
+    base::Closure on_key_not_exists = base::Bind(
+        &AttestationFlow::StartCertificateRequest,
+        weak_factory_.GetWeakPtr(),
+        certificate_profile,
+        true,
+        callback);
+    cryptohome_client_->TpmAttestationDoesKeyExist(
+        key_type,
+        key_name,
+        base::Bind(&DBusBoolRedirectCallback,
+            on_key_exists,
+            on_key_not_exists,
+            base::Bind(callback, false, "")));
+  }
 }
 
 void AttestationFlow::SendCertificateRequestToPCA(
-    const std::string& name,
+    AttestationKeyType key_type,
+    const std::string& key_name,
     const CertificateCallback& callback,
     bool success,
     const std::string& data) {
@@ -174,12 +253,14 @@ void AttestationFlow::SendCertificateRequestToPCA(
       data,
       base::Bind(&AttestationFlow::SendCertificateResponseToDaemon,
                  weak_factory_.GetWeakPtr(),
-                 name,
+                 key_type,
+                 key_name,
                  callback));
 }
 
 void AttestationFlow::SendCertificateResponseToDaemon(
-    const std::string& name,
+    AttestationKeyType key_type,
+    const std::string& key_name,
     const CertificateCallback& callback,
     bool success,
     const std::string& data) {
@@ -191,14 +272,21 @@ void AttestationFlow::SendCertificateResponseToDaemon(
   }
 
   // Forward the response to the attestation service to complete the operation.
-  CryptohomeClient::AttestationKeyType key_type = CryptohomeClient::USER_KEY;
-  if (name == kEnterpriseMachineKey)
-    key_type = CryptohomeClient::DEVICE_KEY;
   async_caller_->AsyncTpmAttestationFinishCertRequest(data,
                                                       key_type,
-                                                      name,
+                                                      key_name,
                                                       base::Bind(callback));
 }
 
+void AttestationFlow::GetExistingCertificate(
+    AttestationKeyType key_type,
+    const std::string& key_name,
+    const CertificateCallback& callback) {
+  cryptohome_client_->TpmAttestationGetCertificate(
+      key_type,
+      key_name,
+      base::Bind(&DBusDataMethodCallback, callback));
+}
+
 }  // namespace attestation
 }  // namespace chromeos