KeyPermissionsService: Drop key locations args and make functions async

This CL changes KeyPermissionsService interface such that no key
locations arguments are expected and all functions are async.

KeyPermissionsService now identifies key locations itself, which is why
it also gets a dependency on PlatformKeysService. Making the functions
asynchronous will also be helpful in a follow-up where setting/getting
corporate flag will require a call to chaps.

Bug: 1130479
Change-Id: Idfb3d36283d22b90109c74a664d4aca60b0f6168
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2421704
Commit-Queue: Omar Morsi <omorsi@google.com>
Reviewed-by: Pavol Marko <pmarko@chromium.org>
Reviewed-by: Edman Anjos <edman@chromium.org>
Cr-Commit-Position: refs/heads/master@{#811197}

chrome/browser/chromeos/arc/enterprise/cert_store/arc_cert_store_bridge_browsertest.cc

@@ -215,8 +215,8 @@ class ArcCertStoreBridgeTest : public MixinBasedInProcessBrowserTest {
       base::RunLoop run_loop;
       key_permissions_service->GetPermissionsForExtension(
           kFakeExtensionId,
-          base::Bind(&ArcCertStoreBridgeTest::GotPermissionsForExtension,
-                     base::Unretained(this), run_loop.QuitClosure()));
+          base::BindOnce(&ArcCertStoreBridgeTest::GotPermissionsForExtension,
+                         base::Unretained(this), run_loop.QuitClosure()));
       run_loop.Run();
     }
   }
@@ -261,11 +261,24 @@ class ArcCertStoreBridgeTest : public MixinBasedInProcessBrowserTest {
       const base::Closure& done_callback,
       std::unique_ptr<chromeos::platform_keys::KeyPermissionsService::
                           PermissionsForExtension> permissions_for_ext) {
+    auto* permissions_for_ext_unowned = permissions_for_ext.get();
     std::string client_cert1_spki(
         client_cert1_->derPublicKey.data,
         client_cert1_->derPublicKey.data + client_cert1_->derPublicKey.len);
-    permissions_for_ext->RegisterKeyForCorporateUsage(
-        client_cert1_spki, {chromeos::platform_keys::TokenId::kUser});
+    permissions_for_ext_unowned->RegisterKeyForCorporateUsage(
+        client_cert1_spki,
+        base::BindOnce(
+            &ArcCertStoreBridgeTest::OnKeyRegisteredForCorporateUsage,
+            base::Unretained(this), std::move(permissions_for_ext),
+            done_callback));
+  }
+
+  void OnKeyRegisteredForCorporateUsage(
+      std::unique_ptr<chromeos::platform_keys::KeyPermissionsService::
+                          PermissionsForExtension> permissions_for_ext,
+      const base::Closure& done_callback,
+      chromeos::platform_keys::Status status) {
+    ASSERT_EQ(status, chromeos::platform_keys::Status::kSuccess);
     done_callback.Run();
   }
 

chrome/browser/chromeos/cert_provisioning/cert_provisioning_worker.cc

@@ -114,6 +114,12 @@ int GetStateOrderedIndex(CertProvisioningWorkerState state) {
   return res;
 }
 
+void OnSetCoporateKeyDone(platform_keys::Status status) {
+  if (status != platform_keys::Status::kSuccess) {
+    LOG(ERROR) << "Cannot mark key corporate: "
+               << platform_keys::StatusToString(status);
+  }
+}
 // Marks the key |public_key_spki_der| as corporate. |profile| can be nullptr if
 // |scope| is CertScope::kDevice.
 void MarkKeyAsCorporate(CertScope scope,
@@ -129,7 +135,8 @@ void MarkKeyAsCorporate(CertScope scope,
   DCHECK(profile);
 
   platform_keys::KeyPermissionsServiceFactory::GetForBrowserContext(profile)
-      ->SetCorporateKey(public_key_spki_der, GetPlatformKeysTokenId(scope));
+      ->SetCorporateKey(public_key_spki_der,
+                        base::BindOnce(&OnSetCoporateKeyDone));
 }
 
 }  // namespace

chrome/browser/chromeos/cert_provisioning/cert_provisioning_worker_unittest.cc

@@ -479,7 +479,7 @@ TEST_F(CertProvisioningWorkerTest, Success) {
     EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());
 
     EXPECT_CALL(*key_permissions_service_,
-                SetCorporateKey(GetPublicKey(), platform_keys::TokenId::kUser));
+                SetCorporateKey(GetPublicKey(), /*callback=*/_));
 
     EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
         platform_keys::TokenId::kUser, GetPublicKey(),
@@ -555,7 +555,7 @@ TEST_F(CertProvisioningWorkerTest, NoVaSuccess) {
         /*callback=*/_));
 
     EXPECT_CALL(*key_permissions_service_,
-                SetCorporateKey(GetPublicKey(), platform_keys::TokenId::kUser));
+                SetCorporateKey(GetPublicKey(), /*callback=*/_));
 
     EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
         platform_keys::TokenId::kUser, GetPublicKey(),
@@ -751,7 +751,7 @@ TEST_F(CertProvisioningWorkerTest, TryLaterWait) {
     EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);
 
     EXPECT_CALL(*key_permissions_service_,
-                SetCorporateKey(GetPublicKey(), platform_keys::TokenId::kUser));
+                SetCorporateKey(GetPublicKey(), /*callback=*/_));
 
     EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
         platform_keys::TokenId::kUser, GetPublicKey(),
@@ -1053,7 +1053,7 @@ TEST_F(CertProvisioningWorkerTest, RemoveRegisteredKey) {
     EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);
 
     EXPECT_CALL(*key_permissions_service_,
-                SetCorporateKey(GetPublicKey(), platform_keys::TokenId::kUser));
+                SetCorporateKey(GetPublicKey(), /*callback=*/_));
 
     EXPECT_SET_ATTRIBUTE_FOR_KEY_FAIL(SetAttributeForKey(
         platform_keys::TokenId::kUser, GetPublicKey(),
@@ -1210,8 +1210,8 @@ TEST_F(CertProvisioningWorkerTest, SerializationSuccess) {
 
     EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);
 
-    EXPECT_CALL(*key_permissions_service_,
-                SetCorporateKey(GetPublicKey(), platform_keys::TokenId::kUser));
+    EXPECT_CALL(*key_permissions_service_, SetCorporateKey(GetPublicKey(),
+                                                           /*callback=*/_));
 
     EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
         platform_keys::TokenId::kUser, GetPublicKey(),

chrome/browser/chromeos/platform_keys/extension_platform_keys_service.cc

@@ -147,17 +147,31 @@ class ExtensionPlatformKeysService::GenerateKeyTask : public Task {
   // Gets the permissions for the extension with id |extension_id|.
   void GetExtensionPermissions() {
     service_->key_permissions_service_->GetPermissionsForExtension(
-        extension_id_,
-        base::Bind(&GenerateKeyTask::GotPermissions, base::Unretained(this)));
+        extension_id_, base::BindOnce(&GenerateKeyTask::GotPermissions,
+                                      base::Unretained(this)));
   }
 
-  void UpdatePermissionsAndCallBack() {
-    std::vector<platform_keys::TokenId> key_locations = {token_id_};
-    extension_permissions_->RegisterKeyForCorporateUsage(public_key_spki_der_,
-                                                         key_locations);
-    callback_.Run(public_key_spki_der_, platform_keys::Status::kSuccess);
+  void OnKeyRegisteredForCorporateUsage(platform_keys::Status status) {
+    if (status == platform_keys::Status::kSuccess) {
+      callback_.Run(public_key_spki_der_, status);
+      DoStep();
+      return;
+    }
+
+    // TODO(crbug.com/1131436): Delete public key if corporate registration
+    // failed.
+    LOG(ERROR) << "Corporate key registration failed: "
+               << platform_keys::StatusToString(status);
+    next_step_ = Step::DONE;
+    callback_.Run(std::string() /* no public key */, status);
     DoStep();
-    return;
+  }
+
+  void UpdatePermissionsAndCallBack() {
+    extension_permissions_->RegisterKeyForCorporateUsage(
+        public_key_spki_der_,
+        base::BindOnce(&GenerateKeyTask::OnKeyRegisteredForCorporateUsage,
+                       base::Unretained(this)));
   }
 
   void GotPermissions(
@@ -230,8 +244,9 @@ class ExtensionPlatformKeysService::SignTask : public Task {
  public:
   enum class Step {
     GET_EXTENSION_PERMISSIONS,
-    GET_KEY_LOCATIONS,
-    SIGN_OR_ABORT,
+    CHECK_SIGN_PERMISSIONS,
+    UPDATE_SIGN_PERMISSIONS,
+    SIGN,
     DONE,
   };
 
@@ -274,33 +289,21 @@ class ExtensionPlatformKeysService::SignTask : public Task {
   void DoStep() {
     switch (next_step_) {
       case Step::GET_EXTENSION_PERMISSIONS:
-        next_step_ = Step::GET_KEY_LOCATIONS;
+        next_step_ = Step::CHECK_SIGN_PERMISSIONS;
         GetExtensionPermissions();
         return;
-      case Step::GET_KEY_LOCATIONS:
-        next_step_ = Step::SIGN_OR_ABORT;
-        GetKeyLocations();
+      case Step::CHECK_SIGN_PERMISSIONS:
+        next_step_ = Step::UPDATE_SIGN_PERMISSIONS;
+        CheckSignPermissions();
+        return;
+      case Step::UPDATE_SIGN_PERMISSIONS:
+        next_step_ = Step::SIGN;
+        UpdateSignPermissions();
         return;
-      case Step::SIGN_OR_ABORT: {
+      case Step::SIGN:
         next_step_ = Step::DONE;
-        const extensions::Extension* extension =
-            extensions::ExtensionRegistry::Get(service_->browser_context_)
-                ->GetExtensionById(extension_id_,
-                                   extensions::ExtensionRegistry::ENABLED);
-        bool sign_granted = extension_permissions_->CanUseKeyForSigning(
-                                public_key_spki_der_, key_locations_) ||
-                            (service_->IsUsingSigninProfile() &&
-                             IsExtensionAllowlisted(extension));
-
-        if (sign_granted) {
-          Sign();
-        } else {
-          callback_.Run(std::string() /* no signature */,
-                        platform_keys::Status::kErrorKeyNotAllowedForSigning);
-          DoStep();
-        }
+        Sign();
         return;
-      }
       case Step::DONE:
         service_->TaskFinished(this);
         // |this| might be invalid now.
@@ -311,7 +314,7 @@ class ExtensionPlatformKeysService::SignTask : public Task {
   void GetExtensionPermissions() {
     service_->key_permissions_service_->GetPermissionsForExtension(
         extension_id_,
-        base::Bind(&SignTask::GotPermissions, base::Unretained(this)));
+        base::BindOnce(&SignTask::GotPermissions, base::Unretained(this)));
   }
 
   void GotPermissions(
@@ -322,32 +325,57 @@ class ExtensionPlatformKeysService::SignTask : public Task {
     DoStep();
   }
 
-  void GetKeyLocations() {
-    service_->platform_keys_service_->GetKeyLocations(
+  void CheckSignPermissions() {
+    const extensions::Extension* extension =
+        extensions::ExtensionRegistry::Get(service_->browser_context_)
+            ->GetExtensionById(extension_id_,
+                               extensions::ExtensionRegistry::ENABLED);
+    if (service_->IsUsingSigninProfile() && IsExtensionAllowlisted(extension)) {
+      DoStep();
+      return;
+    }
+
+    extension_permissions_->CanUseKeyForSigning(
+        public_key_spki_der_,
+        base::BindOnce(&SignTask::OnCanUseKeyForSigningKnown,
+                       base::Unretained(this)));
+  }
+
+  void OnCanUseKeyForSigningKnown(bool allowed) {
+    if (!allowed) {
+      callback_.Run(std::string() /* no signature */,
+                    platform_keys::Status::kErrorKeyNotAllowedForSigning);
+      next_step_ = Step::DONE;
+      DoStep();
+    }
+
+    DoStep();
+  }
+
+  // Updates the signing permissions for |public_key_spki_der_|.
+  void UpdateSignPermissions() {
+    extension_permissions_->SetKeyUsedForSigning(
         public_key_spki_der_,
-        base::BindRepeating(&SignTask::GotKeyLocation, base::Unretained(this)));
+        base::BindOnce(&SignTask::OnSetKeyUsedForSigningDone,
+                       base::Unretained(this)));
   }
 
-  void GotKeyLocation(const std::vector<platform_keys::TokenId>& token_ids,
-                      platform_keys::Status status) {
+  void OnSetKeyUsedForSigningDone(platform_keys::Status status) {
     if (status != platform_keys::Status::kSuccess) {
+      LOG(ERROR) << "Marking a key used for signing failed: "
+                 << platform_keys::StatusToString(status);
       next_step_ = Step::DONE;
       callback_.Run(std::string() /* no signature */, status);
       DoStep();
       return;
     }
 
-    key_locations_ = token_ids;
     DoStep();
   }
 
-  // Updates the permissions for |public_key_spki_der_|, starts the actual
-  // signing operation and afterwards passes the signature (or error) to
-  // |callback_|.
+  // Starts the actual signing operation and afterwards passes the signature (or
+  // error) to |callback_|.
   void Sign() {
-    extension_permissions_->SetKeyUsedForSigning(public_key_spki_der_,
-                                                 key_locations_);
-
     switch (key_type_) {
       case platform_keys::KeyType::kRsassaPkcs1V15: {
         if (raw_pkcs1_) {
@@ -392,7 +420,6 @@ class ExtensionPlatformKeysService::SignTask : public Task {
   const SignCallback callback_;
   std::unique_ptr<platform_keys::KeyPermissionsService::PermissionsForExtension>
       extension_permissions_;
-  std::vector<platform_keys::TokenId> key_locations_;
   ExtensionPlatformKeysService* const service_;
   base::WeakPtrFactory<SignTask> weak_factory_{this};
 
@@ -404,11 +431,11 @@ class ExtensionPlatformKeysService::SelectTask : public Task {
   enum class Step {
     GET_EXTENSION_PERMISSIONS,
     GET_MATCHING_CERTS,
-    GET_KEY_LOCATIONS,
+    CHECK_KEY_PERMISSIONS,
     INTERSECT_WITH_INPUT_CERTS,
     SELECT_CERTS,
     UPDATE_PERMISSION,
-    FILTER_BY_PERMISSIONS,
+    PASS_RESULTING_CERTS,
     DONE,
   };
 
@@ -450,20 +477,20 @@ class ExtensionPlatformKeysService::SelectTask : public Task {
         GetExtensionPermissions();
         return;
       case Step::GET_MATCHING_CERTS:
-        next_step_ = Step::GET_KEY_LOCATIONS;
+        next_step_ = Step::CHECK_KEY_PERMISSIONS;
         GetMatchingCerts();
         return;
-      case Step::GET_KEY_LOCATIONS:
-        // Don't advance to the next step yet - GetKeyLocations is repeated for
-        // all matching certs. The next step will be selected in
-        // GetKeyLocations.
-        GetKeyLocations(Step::INTERSECT_WITH_INPUT_CERTS /* next_step */);
+      case Step::CHECK_KEY_PERMISSIONS:
+        // Don't advance to the next step yet - CheckKeyPermissions is repeated
+        // for all matching certs. The next step will be selected in
+        // CheckKeyPermissions.
+        CheckKeyPermissions(Step::INTERSECT_WITH_INPUT_CERTS /* next_step */);
         return;
       case Step::INTERSECT_WITH_INPUT_CERTS:
         if (interactive_)
           next_step_ = Step::SELECT_CERTS;
         else  // Skip SelectCerts and UpdatePermission if not interactive.
-          next_step_ = Step::FILTER_BY_PERMISSIONS;
+          next_step_ = Step::PASS_RESULTING_CERTS;
         IntersectWithInputCerts();
         return;
       case Step::SELECT_CERTS:
@@ -471,12 +498,12 @@ class ExtensionPlatformKeysService::SelectTask : public Task {
         SelectCerts();
         return;
       case Step::UPDATE_PERMISSION:
-        next_step_ = Step::FILTER_BY_PERMISSIONS;
+        next_step_ = Step::PASS_RESULTING_CERTS;
         UpdatePermission();
         return;
-      case Step::FILTER_BY_PERMISSIONS:
+      case Step::PASS_RESULTING_CERTS:
         next_step_ = Step::DONE;
-        FilterSelectionByPermission();
+        PassResultingCerts();
         return;
       case Step::DONE:
         service_->TaskFinished(this);
@@ -488,7 +515,7 @@ class ExtensionPlatformKeysService::SelectTask : public Task {
   void GetExtensionPermissions() {
     service_->key_permissions_service_->GetPermissionsForExtension(
         extension_id_,
-        base::Bind(&SelectTask::GotPermissions, base::Unretained(this)));
+        base::BindOnce(&SelectTask::GotPermissions, base::Unretained(this)));
   }
 
   void GotPermissions(
@@ -539,57 +566,59 @@ class ExtensionPlatformKeysService::SelectTask : public Task {
           continue;
       }
 
-      matches_pending_key_locations_.push_back(std::move(certificate));
+      matches_pending_permissions_check_.push_back(std::move(certificate));
     }
     DoStep();
   }
 
   // This is called once for each certificate in
-  // |matches_pending_key_locations_|.  Each invocation processes the first
+  // |matches_pending_permissions_check_|. Each invocation processes the first
   // element and removes it from the deque. Each processed certificate is added
-  // to |matches_| and |key_locations_for_matches_| if it is selectable
-  // according to KeyPermissionsService. When all certificates have been
-  // processed, advances the SignTask state machine to |next_step|.
-  void GetKeyLocations(Step next_step) {
-    if (matches_pending_key_locations_.empty()) {
+  // to |matches_| if it is selectable according to KeyPermissionsService. When
+  // all certificates have been processed, advances the SignTask state machine
+  // to |next_step|.
+  void CheckKeyPermissions(Step next_step) {
+    if (matches_pending_permissions_check_.empty()) {
       next_step_ = next_step;
       DoStep();
       return;
     }
 
     scoped_refptr<net::X509Certificate> certificate =
-        std::move(matches_pending_key_locations_.front());
-    matches_pending_key_locations_.pop_front();
+        std::move(matches_pending_permissions_check_.front());
+    matches_pending_permissions_check_.pop_front();
     const std::string public_key_spki_der(
         platform_keys::GetSubjectPublicKeyInfo(certificate));
 
-    service_->platform_keys_service_->GetKeyLocations(
+    extension_permissions_->CanUseKeyForSigning(
         public_key_spki_der,
-        base::BindRepeating(&SelectTask::GotKeyLocations,
-                            base::Unretained(this), certificate));
+        base::BindOnce(&SelectTask::OnKeySigningPermissionKnown,
+                       base::Unretained(this), public_key_spki_der,
+                       certificate));
   }
 
-  void GotKeyLocations(const scoped_refptr<net::X509Certificate>& certificate,
-                       const std::vector<platform_keys::TokenId>& token_ids,
-                       platform_keys::Status status) {
-    if (status != platform_keys::Status::kSuccess) {
-      next_step_ = Step::DONE;
-      callback_.Run(nullptr /* no certificates */, status);
+  void OnKeySigningPermissionKnown(
+      const std::string& public_key_spki_der,
+      const scoped_refptr<net::X509Certificate>& certificate,
+      bool allowed) {
+    if (allowed) {
+      matches_.push_back(certificate);
+      DoStep();
+    } else if (interactive_) {
+      service_->key_permissions_service_->CanUserGrantPermissionForKey(
+          public_key_spki_der,
+          base::BindOnce(&SelectTask::OnAbilityToGrantPermissionKnown,
+                         base::Unretained(this), std::move(certificate)));
+    } else {
       DoStep();
-      return;
     }
+  }
 
-    const std::string public_key_spki_der(
-        platform_keys::GetSubjectPublicKeyInfo(certificate));
-
-    // Use this key if the user can use it for signing or can grant permission
-    // for it.
-    if (service_->key_permissions_service_->CanUserGrantPermissionFor(
-            public_key_spki_der, token_ids) ||
-        extension_permissions_->CanUseKeyForSigning(public_key_spki_der,
-                                                    token_ids)) {
+  void OnAbilityToGrantPermissionKnown(
+      const scoped_refptr<net::X509Certificate>& certificate,
+      bool allowed) {
+    if (allowed) {
       matches_.push_back(certificate);
-      key_locations_for_matches_[public_key_spki_der] = token_ids;
     }
     DoStep();
   }
@@ -644,18 +673,22 @@ class ExtensionPlatformKeysService::SelectTask : public Task {
     }
     const std::string public_key_spki_der(
         platform_keys::GetSubjectPublicKeyInfo(selected_cert_));
-    auto key_locations_iter =
-        key_locations_for_matches_.find(public_key_spki_der);
-    CHECK(key_locations_iter != key_locations_for_matches_.end());
     extension_permissions_->SetUserGrantedPermission(
-        public_key_spki_der, key_locations_iter->second);
+        public_key_spki_der, base::BindOnce(&SelectTask::OnPermissionsUpdated,
+                                            base::Unretained(this)));
+  }
+
+  void OnPermissionsUpdated(platform_keys::Status status) {
+    if (status != platform_keys::Status::kSuccess) {
+      LOG(WARNING) << "Error while updating permissions: "
+                   << platform_keys::StatusToString(status);
+    }
+
     DoStep();
   }
 
-  // Filters from all matches (if not interactive) or from the selection (if
-  // interactive), the certificates that the extension has unlimited sign
-  // permission for. Passes the filtered certs to |callback_|.
-  void FilterSelectionByPermission() {
+  // Passes the filtered certs to |callback_|.
+  void PassResultingCerts() {
     std::unique_ptr<net::CertificateList> selection(new net::CertificateList);
     if (interactive_) {
       if (selected_cert_)
@@ -664,39 +697,15 @@ class ExtensionPlatformKeysService::SelectTask : public Task {
       selection->assign(matches_.begin(), matches_.end());
     }
 
-    std::unique_ptr<net::CertificateList> filtered_certs(
-        new net::CertificateList);
-    for (scoped_refptr<net::X509Certificate> selected_cert : *selection) {
-      const std::string public_key_spki_der(
-          platform_keys::GetSubjectPublicKeyInfo(selected_cert));
-
-      auto key_locations_iter =
-          key_locations_for_matches_.find(public_key_spki_der);
-      CHECK(key_locations_iter != key_locations_for_matches_.end());
-      if (!extension_permissions_->CanUseKeyForSigning(
-              public_key_spki_der, key_locations_iter->second))
-        continue;
-
-      filtered_certs->push_back(selected_cert);
-    }
-    // Note: In the interactive case this should have filtered exactly the
-    // one selected cert. Checking the permissions again is not striclty
-    // necessary but this ensures that the permissions were updated correctly.
-    CHECK(!selected_cert_ || (filtered_certs->size() == 1 &&
-                              filtered_certs->front() == selected_cert_));
-    callback_.Run(std::move(filtered_certs), platform_keys::Status::kSuccess);
+    callback_.Run(std::move(selection), platform_keys::Status::kSuccess);
     DoStep();
   }
 
   Step next_step_ = Step::GET_EXTENSION_PERMISSIONS;
 
   std::deque<scoped_refptr<net::X509Certificate>>
-      matches_pending_key_locations_;
+      matches_pending_permissions_check_;
   net::CertificateList matches_;
-  // Mapping of DER-encoded Subject Public Key Info to the TokenIds determined
-  // for the corresponding private key.
-  base::flat_map<std::string, std::vector<platform_keys::TokenId>>
-      key_locations_for_matches_;
   scoped_refptr<net::X509Certificate> selected_cert_;
   platform_keys::ClientCertificateRequest request_;
   std::unique_ptr<net::CertificateList> input_client_certificates_;

chrome/browser/chromeos/platform_keys/key_permissions/key_permissions_service.h

@@ -16,6 +16,23 @@
 namespace chromeos {
 namespace platform_keys {
 
+using CanUseKeyForSigningCallback = base::OnceCallback<void(bool allowed)>;
+
+using RegisterKeyForCorporateUsageCallback =
+    base::OnceCallback<void(Status status)>;
+
+using SetUserGrantedPermissionCallback =
+    base::OnceCallback<void(Status status)>;
+
+using SetKeyUsedForSigningCallback = base::OnceCallback<void(Status status)>;
+
+using CanUserGrantPermissionForKeyCallback =
+    base::OnceCallback<void(bool allowed)>;
+
+using IsCorporateKeyCallback = base::OnceCallback<void(bool corporate)>;
+
+using SetCorporateKeyCallback = base::OnceCallback<void(Status status)>;
+
 // This service will be responsible for answering queries regarding platform key
 // permissions with respect to a specific profile.
 //
@@ -67,72 +84,63 @@ class KeyPermissionsService : public KeyedService {
     PermissionsForExtension();
     virtual ~PermissionsForExtension();
 
-    // Returns true if the private key matching |public_key_spki_der| can be
-    // used for signing by the extension with id |extension_id_|.
-    // |key_locations| must describe locations available to the user the private
-    // key is stored on.
-    virtual bool CanUseKeyForSigning(
+    // Determines if the private key matching |public_key_spki_der| can be used
+    // for signing by the extension with id |extension_id_|. |callback| will be
+    // invoked with the result.
+    virtual void CanUseKeyForSigning(const std::string& public_key_spki_der,
+                                     CanUseKeyForSigningCallback callback) = 0;
+
+    // Must be called when the extension with id |extension_id| used the private
+    // key matching |public_key_spki_der| for signing. |callback| will be
+    // invoked with the resulting status. Updates the permissions accordingly.
+    // E.g. if this extension generated the key and no other permission was
+    // granted then the permission to sign with this key is removed.
+    virtual void SetKeyUsedForSigning(
         const std::string& public_key_spki_der,
-        const std::vector<platform_keys::TokenId>& key_locations) = 0;
+        SetKeyUsedForSigningCallback callback) = 0;
 
     // Registers the private key matching |public_key_spki_der| as being
     // generated by the extension with id |extension_id| and marks it for
-    // corporate usage. |key_locations| must describe locations available to the
-    // user the private key is stored on.
+    // corporate usage. |callback| will be invoked with the resulting status.
     virtual void RegisterKeyForCorporateUsage(
         const std::string& public_key_spki_der,
-        const std::vector<platform_keys::TokenId>& key_locations) = 0;
+        RegisterKeyForCorporateUsageCallback callback) = 0;
 
     // Sets the user granted permission that the extension with id
     // |extension_id| can use the private key matching |public_key_spki_der| for
-    // signing. |key_locations| must describe locations available to the user
-    // the private key is stored on.
+    // signing. |callback| will be invoked with the resulting status.
     virtual void SetUserGrantedPermission(
         const std::string& public_key_spki_der,
-        const std::vector<platform_keys::TokenId>& key_locations) = 0;
-
-    // Must be called when the extension with id |extension_id| used the private
-    // key matching |public_key_spki_der| for signing. |key_locations| must
-    // describe locations available to the user the private key is stored on.
-    // Updates the permissions accordingly.  E.g. if this extension generated
-    // the key and no other permission was granted then the permission to sign
-    // with this key is removed.
-    virtual void SetKeyUsedForSigning(
-        const std::string& public_key_spki_der,
-        const std::vector<platform_keys::TokenId>& key_locations) = 0;
+        SetUserGrantedPermissionCallback callback) = 0;
   };
 
   KeyPermissionsService();
   ~KeyPermissionsService() override;
 
-  using PermissionsCallback =
-      base::Callback<void(std::unique_ptr<PermissionsForExtension>)>;
-
+  using GetPermissionsForExtensionCallback = base::OnceCallback<void(
+      std::unique_ptr<PermissionsForExtension> permissions_for_extension)>;
   // Passes an object managing the key permissions of the extension with id
   // |extension_id| to |callback|. This can happen synchronously or
   // asynchronously.
   virtual void GetPermissionsForExtension(
       const std::string& extension_id,
-      const PermissionsCallback& callback) = 0;
+      GetPermissionsForExtensionCallback callback) = 0;
 
-  // Returns true if the user can grant any permission for
-  // |public_key_spki_derey_id| to extensions. |key_locations| must describe
-  // locations available to the user the private key is stored on.
-  virtual bool CanUserGrantPermissionFor(
+  // Determines if the user can grant any permission for |public_key_spki_der|
+  // to extensions. |callback| will be invoked with the result.
+  virtual void CanUserGrantPermissionForKey(
       const std::string& public_key_spki_der,
-      const std::vector<platform_keys::TokenId>& key_locations) const = 0;
+      CanUserGrantPermissionForKeyCallback callback) const = 0;
 
-  // Returns true if the key identified by |public_key_spki_der| that is
-  // located on |key_locations| is marked for corporate usage.
-  virtual bool IsCorporateKey(
-      const std::string& public_key_spki_der,
-      const std::vector<platform_keys::TokenId>& key_locations) const = 0;
+  // Determines if the key identified by |public_key_spki_der|is marked for
+  // corporate usage. |callback| will be invoked with the result.
+  virtual void IsCorporateKey(const std::string& public_key_spki_der,
+                              IsCorporateKeyCallback callback) const = 0;
 
   // Marks the key identified by |public_key_spki_der| as corporate usage.
-  // Accepts a single key location because this is intended for usage after key
-  // generation / import, when exactly one location is relevant.
+  // |callback| will be invoked with the resulting status.
   virtual void SetCorporateKey(const std::string& public_key_spki_der,
-                               platform_keys::TokenId key_location) const = 0;
+                               SetCorporateKeyCallback callback) const = 0;
 };
 
 }  // namespace platform_keys

chrome/browser/chromeos/platform_keys/key_permissions/key_permissions_service_factory.cc

@@ -50,7 +50,8 @@ KeyedService* KeyPermissionsServiceFactory::BuildServiceInstanceFor(
   return new KeyPermissionsServiceImpl(
       profile->GetProfilePolicyConnector()->IsManaged(), profile->GetPrefs(),
       profile->GetProfilePolicyConnector()->policy_service(),
-      extensions::ExtensionSystem::Get(profile)->state_store());
+      extensions::ExtensionSystem::Get(profile)->state_store(),
+      PlatformKeysServiceFactory::GetForBrowserContext(profile));
 }
 
 void KeyPermissionsServiceFactory::RegisterProfilePrefs(

chrome/browser/chromeos/platform_keys/key_permissions/key_permissions_service_impl.cc

@@ -6,16 +6,19 @@
 
 #include <memory>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
 #include "base/base64.h"
 #include "base/bind.h"
 #include "base/callback.h"
+#include "base/callback_forward.h"
 #include "base/logging.h"
 #include "base/stl_util.h"
 #include "base/values.h"
 #include "chrome/browser/chromeos/platform_keys/platform_keys.h"
+#include "chrome/browser/chromeos/platform_keys/platform_keys_service.h"
 #include "chrome/browser/policy/profile_policy_connector.h"
 #include "chrome/common/pref_names.h"
 #include "components/policy/core/common/policy_map.h"
@@ -197,11 +200,33 @@ KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
 KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
     ~PermissionsForExtensionImpl() = default;
 
-bool KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
+void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
     CanUseKeyForSigning(const std::string& public_key_spki_der,
-                        const std::vector<TokenId>& key_locations) {
-  if (key_locations.empty())
-    return false;
+                        CanUseKeyForSigningCallback callback) {
+  key_permissions_service_->platform_keys_service_->GetKeyLocations(
+      public_key_spki_der,
+      base::BindOnce(
+          &PermissionsForExtensionImpl::CanUseKeyForSigningWithLocations,
+          weak_factory_.GetWeakPtr(), public_key_spki_der,
+          std::move(callback)));
+}
+
+void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
+    CanUseKeyForSigningWithLocations(const std::string& public_key_spki_der,
+                                     CanUseKeyForSigningCallback callback,
+                                     const std::vector<TokenId>& key_locations,
+                                     Status key_locations_retrieval_status) {
+  if (key_locations_retrieval_status != Status::kSuccess) {
+    LOG(ERROR) << "PlatformKeysService error on requesting key locations: "
+               << StatusToString(key_locations_retrieval_status);
+    std::move(callback).Run(/*allowed=*/false);
+    return;
+  }
+
+  if (key_locations.empty()) {
+    std::move(callback).Run(/*allowed=*/false);
+    return;
+  }
 
   std::string public_key_spki_der_b64;
   base::Base64Encode(public_key_spki_der, &public_key_spki_der_b64);
@@ -215,49 +240,100 @@ bool KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
   // That means, once a certificate authority generated a certificate for the
   // key, the generating extension doesn't have access to the key anymore,
   // except if explicitly permitted by the administrator.
-  if (matching_entry->sign_once)
-    return true;
+  if (matching_entry->sign_once) {
+    std::move(callback).Run(/*allowed=*/true);
+    return;
+  }
 
+  auto bound_callback =
+      base::BindOnce(&KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
+                         CanUseKeyForSigningWithFlags,
+                     weak_factory_.GetWeakPtr(), std::move(callback),
+                     matching_entry->sign_unlimited);
+  key_permissions_service_->IsCorporateKeyWithLocations(
+      public_key_spki_der, std::move(bound_callback), key_locations,
+      key_locations_retrieval_status);
+}
+
+void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
+    CanUseKeyForSigningWithFlags(CanUseKeyForSigningCallback callback,
+                                 bool sign_unlimited_allowed,
+                                 bool is_corporate_key) {
   // Usage of corporate keys is solely determined by policy. The user must not
   // circumvent this decision.
-  if (key_permissions_service_->IsCorporateKey(public_key_spki_der,
-                                               key_locations))
-    return PolicyAllowsCorporateKeyUsage();
+  if (is_corporate_key) {
+    std::move(callback).Run(/*allowed=*/PolicyAllowsCorporateKeyUsage());
+    return;
+  }
 
   // Only permissions for keys that are not designated for corporate usage are
   // determined by user decisions.
-  return matching_entry->sign_unlimited;
+  std::move(callback).Run(sign_unlimited_allowed);
 }
 
 void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
     SetKeyUsedForSigning(const std::string& public_key_spki_der,
-                         const std::vector<TokenId>& key_locations) {
-  if (key_locations.empty())
+                         SetKeyUsedForSigningCallback callback) {
+  key_permissions_service_->platform_keys_service_->GetKeyLocations(
+      public_key_spki_der,
+      base::BindOnce(
+          &PermissionsForExtensionImpl::SetKeyUsedForSigningWithLocations,
+          weak_factory_.GetWeakPtr(), public_key_spki_der,
+          std::move(callback)));
+}
+
+void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
+    SetKeyUsedForSigningWithLocations(const std::string& public_key_spki_der,
+                                      SetKeyUsedForSigningCallback callback,
+                                      const std::vector<TokenId>& key_locations,
+                                      Status key_locations_retrieval_status) {
+  if (key_locations_retrieval_status != Status::kSuccess) {
+    std::move(callback).Run(key_locations_retrieval_status);
     return;
+  }
+
+  if (key_locations.empty()) {
+    std::move(callback).Run(Status::kErrorKeyNotFound);
+    return;
+  }
 
   std::string public_key_spki_der_b64;
   base::Base64Encode(public_key_spki_der, &public_key_spki_der_b64);
 
   KeyEntry* matching_entry = GetStateStoreEntry(public_key_spki_der_b64);
-
-  if (!matching_entry->sign_once) {
-    if (!CanUseKeyForSigning(public_key_spki_der, key_locations))
-      LOG(ERROR) << "Key was not allowed for signing.";
-    return;
-  }
-
   matching_entry->sign_once = false;
   WriteToStateStore();
+
+  std::move(callback).Run(Status::kSuccess);
 }
 
 void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
-    RegisterKeyForCorporateUsage(const std::string& public_key_spki_der,
-                                 const std::vector<TokenId>& key_locations) {
-  if (key_locations.empty()) {
-    NOTREACHED();
+    RegisterKeyForCorporateUsage(
+        const std::string& public_key_spki_der,
+        RegisterKeyForCorporateUsageCallback callback) {
+  key_permissions_service_->platform_keys_service_->GetKeyLocations(
+      public_key_spki_der,
+      base::BindOnce(&PermissionsForExtensionImpl::
+                         RegisterKeyForCorporateUsageWithLocations,
+                     weak_factory_.GetWeakPtr(), public_key_spki_der,
+                     std::move(callback)));
+}
+
+void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
+    RegisterKeyForCorporateUsageWithLocations(
+        const std::string& public_key_spki_der,
+        RegisterKeyForCorporateUsageCallback callback,
+        const std::vector<TokenId>& key_locations,
+        Status key_locations_retrieval_status) {
+  if (key_locations_retrieval_status != Status::kSuccess) {
+    std::move(callback).Run(key_locations_retrieval_status);
     return;
   }
 
+  if (key_locations.empty()) {
+    std::move(callback).Run(Status::kErrorKeyNotFound);
+    return;
+  }
   std::string public_key_spki_der_b64;
   base::Base64Encode(public_key_spki_der, &public_key_spki_der_b64);
 
@@ -265,6 +341,7 @@ void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
 
   if (matching_entry->sign_once) {
     VLOG(1) << "Key is already allowed for signing, skipping.";
+    std::move(callback).Run(Status::kSuccess);
     return;
   }
 
@@ -276,21 +353,57 @@ void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
   // stored the sign_once permission, so the enrolling extension in the same
   // profile can use the key for signing once in order to build a CSR even if it
   // doesn't have permission to use corporate keys.
-  if (!IsKeyOnUserSlot(key_locations))
+  if (!IsKeyOnUserSlot(key_locations)) {
+    std::move(callback).Run(Status::kSuccess);
     return;
+  }
 
-  key_permissions_service_->SetCorporateKey(public_key_spki_der,
-                                            TokenId::kUser);
+  key_permissions_service_->SetCorporateKeyWithLocations(
+      public_key_spki_der, std::move(callback), key_locations,
+      /*key_locations_retrieval_status=*/Status::kSuccess);
 }
 
 void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
     SetUserGrantedPermission(const std::string& public_key_spki_der,
-                             const std::vector<TokenId>& key_locations) {
-  if (!key_permissions_service_->CanUserGrantPermissionFor(public_key_spki_der,
-                                                           key_locations)) {
-    LOG(WARNING) << "Tried to grant permission for a key although prohibited "
-                    "(either key is a corporate key or this account is "
-                    "managed).";
+                             SetUserGrantedPermissionCallback callback) {
+  key_permissions_service_->platform_keys_service_->GetKeyLocations(
+      public_key_spki_der,
+      base::BindOnce(
+          &PermissionsForExtensionImpl::SetUserGrantedPermissionWithLocations,
+          weak_factory_.GetWeakPtr(), public_key_spki_der,
+          std::move(callback)));
+}
+
+void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
+    SetUserGrantedPermissionWithLocations(
+        const std::string& public_key_spki_der,
+        SetUserGrantedPermissionCallback callback,
+        const std::vector<TokenId>& key_locations,
+        Status key_locations_retrieval_status) {
+  key_permissions_service_->CanUserGrantPermissionForKeyWithLocations(
+      public_key_spki_der,
+      base::BindOnce(&PermissionsForExtensionImpl::
+                         SetUserGrantedPermissionWithLocationsAndFlag,
+                     weak_factory_.GetWeakPtr(), public_key_spki_der,
+                     std::move(callback), key_locations,
+                     key_locations_retrieval_status),
+      key_locations, key_locations_retrieval_status);
+}
+
+void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
+    SetUserGrantedPermissionWithLocationsAndFlag(
+        const std::string& public_key_spki_der,
+        SetUserGrantedPermissionCallback callback,
+        const std::vector<TokenId>& key_locations,
+        Status key_locations_retrieval_status,
+        bool can_user_grant_permission) {
+  if (key_locations_retrieval_status != Status::kSuccess) {
+    std::move(callback).Run(key_locations_retrieval_status);
+    return;
+  }
+
+  if (!can_user_grant_permission) {
+    std::move(callback).Run(Status::kErrorGrantKeyPermissionForExtension);
     return;
   }
 
@@ -305,11 +418,13 @@ void KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
 
   if (matching_entry->sign_unlimited) {
     VLOG(1) << "Key is already allowed for signing, skipping.";
+    std::move(callback).Run(Status::kSuccess);
     return;
   }
 
   matching_entry->sign_unlimited = true;
   WriteToStateStore();
+  std::move(callback).Run(Status::kSuccess);
 }
 
 bool KeyPermissionsServiceImpl::PermissionsForExtensionImpl::
@@ -399,13 +514,16 @@ KeyPermissionsServiceImpl::KeyPermissionsServiceImpl(
     bool profile_is_managed,
     PrefService* profile_prefs,
     policy::PolicyService* profile_policies,
-    extensions::StateStore* extensions_state_store)
+    extensions::StateStore* extensions_state_store,
+    PlatformKeysService* platform_keys_service)
     : profile_is_managed_(profile_is_managed),
       profile_prefs_(profile_prefs),
       profile_policies_(profile_policies),
-      extensions_state_store_(extensions_state_store) {
+      extensions_state_store_(extensions_state_store),
+      platform_keys_service_(platform_keys_service) {
   DCHECK(profile_prefs_);
   DCHECK(extensions_state_store_);
+  DCHECK(platform_keys_service_);
   DCHECK(!profile_is_managed_ || profile_policies_);
 }
 
@@ -413,68 +531,160 @@ KeyPermissionsServiceImpl::~KeyPermissionsServiceImpl() = default;
 
 void KeyPermissionsServiceImpl::GetPermissionsForExtension(
     const std::string& extension_id,
-    const PermissionsCallback& callback) {
+    GetPermissionsForExtensionCallback callback) {
   extensions_state_store_->GetExtensionValue(
       extension_id, kStateStorePlatformKeys,
       base::BindOnce(
           &KeyPermissionsServiceImpl::CreatePermissionObjectAndPassToCallback,
-          weak_factory_.GetWeakPtr(), extension_id, callback));
+          weak_factory_.GetWeakPtr(), extension_id, std::move(callback)));
 }
 
-bool KeyPermissionsServiceImpl::CanUserGrantPermissionFor(
+void KeyPermissionsServiceImpl::CanUserGrantPermissionForKey(
     const std::string& public_key_spki_der,
-    const std::vector<TokenId>& key_locations) const {
-  if (key_locations.empty())
-    return false;
+    CanUserGrantPermissionForKeyCallback callback) const {
+  platform_keys_service_->GetKeyLocations(
+      public_key_spki_der,
+      base::BindOnce(
+          &KeyPermissionsServiceImpl::CanUserGrantPermissionForKeyWithLocations,
+          weak_factory_.GetWeakPtr(), public_key_spki_der,
+          std::move(callback)));
+}
+
+void KeyPermissionsServiceImpl::CanUserGrantPermissionForKeyWithLocations(
+    const std::string& public_key_spki_der,
+    CanUserGrantPermissionForKeyCallback callback,
+    const std::vector<TokenId>& key_locations,
+    Status key_locations_retrieval_status) const {
+  auto bound_callback = base::BindOnce(
+      &KeyPermissionsServiceImpl::
+          CanUserGrantPermissionForKeyWithLocationsAndFlag,
+      weak_factory_.GetWeakPtr(), public_key_spki_der, std::move(callback),
+      key_locations, key_locations_retrieval_status);
+  IsCorporateKeyWithLocations(public_key_spki_der, std::move(bound_callback),
+                              key_locations, key_locations_retrieval_status);
+}
+
+void KeyPermissionsServiceImpl::
+    CanUserGrantPermissionForKeyWithLocationsAndFlag(
+        const std::string& public_key_spki_der,
+        CanUserGrantPermissionForKeyCallback callback,
+        const std::vector<TokenId>& key_locations,
+        Status status,
+        bool corporate_key) {
+  if (status != Status::kSuccess) {
+    std::move(callback).Run(/*allowed=*/false);
+    return;
+  }
+
+  if (key_locations.empty()) {
+    std::move(callback).Run(/*allowed=*/false);
+    return;
+  }
 
   // As keys cannot be tagged for non-corporate usage, the user can currently
   // not grant any permissions if the profile is managed.
-  if (profile_is_managed_)
-    return false;
+  if (profile_is_managed_) {
+    std::move(callback).Run(/*allowed=*/false);
+    return;
+  }
 
   // If this profile is not managed but we find a corporate key, don't allow
   // the user to grant permissions.
-  return !IsCorporateKey(public_key_spki_der, key_locations);
+  std::move(callback).Run(/*allowed=*/!corporate_key);
+}
+
+void KeyPermissionsServiceImpl::IsCorporateKey(
+    const std::string& public_key_spki_der,
+    IsCorporateKeyCallback callback) const {
+  platform_keys_service_->GetKeyLocations(
+      public_key_spki_der,
+      base::BindOnce(&KeyPermissionsServiceImpl::IsCorporateKeyWithLocations,
+                     weak_factory_.GetWeakPtr(), public_key_spki_der,
+                     std::move(callback)));
 }
 
-bool KeyPermissionsServiceImpl::IsCorporateKey(
+void KeyPermissionsServiceImpl::IsCorporateKeyWithLocations(
     const std::string& public_key_spki_der,
-    const std::vector<TokenId>& key_locations) const {
+    IsCorporateKeyCallback callback,
+    const std::vector<TokenId>& key_locations,
+    Status status) const {
+  if (status != Status::kSuccess) {
+    LOG(ERROR) << "Key locations retrieval failed: " << StatusToString(status);
+    std::move(callback).Run(/*corporate=*/false);
+    return;
+  }
+
   std::string public_key_spki_der_b64;
   base::Base64Encode(public_key_spki_der, &public_key_spki_der_b64);
 
   for (const auto key_location : key_locations) {
     switch (key_location) {
       case TokenId::kUser:
-        if (IsCorporateKeyForProfile(public_key_spki_der_b64, profile_prefs_))
-          return true;
+        if (IsCorporateKeyForProfile(public_key_spki_der_b64, profile_prefs_)) {
+          std::move(callback).Run(/*corporate=*/true);
+          return;
+        }
         break;
       case TokenId::kSystem:
-        return true;
+        std::move(callback).Run(/*corporate=*/true);
+        return;
     }
   }
-  return false;
+  std::move(callback).Run(/*corporate=*/false);
 }
 
 void KeyPermissionsServiceImpl::SetCorporateKey(
     const std::string& public_key_spki_der,
-    TokenId key_location) const {
-  if (key_location == TokenId::kSystem) {
-    // Nothing to do - all system-token keys are currently implicitly corporate.
+    SetCorporateKeyCallback callback) const {
+  platform_keys_service_->GetKeyLocations(
+      public_key_spki_der,
+      base::BindOnce(&KeyPermissionsServiceImpl::SetCorporateKeyWithLocations,
+                     weak_factory_.GetWeakPtr(), public_key_spki_der,
+                     std::move(callback)));
+}
+
+void KeyPermissionsServiceImpl::SetCorporateKeyWithLocations(
+    const std::string& public_key_spki_der,
+    SetCorporateKeyCallback callback,
+    const std::vector<TokenId>& key_locations,
+    Status key_locations_retrieval_status) const {
+  if (key_locations_retrieval_status != Status::kSuccess) {
+    std::move(callback).Run(key_locations_retrieval_status);
     return;
   }
 
-  std::string public_key_spki_der_b64;
-  base::Base64Encode(public_key_spki_der, &public_key_spki_der_b64);
-
-  DictionaryPrefUpdate update(profile_prefs_, prefs::kPlatformKeys);
-
-  std::unique_ptr<base::DictionaryValue> new_pref_entry(
-      new base::DictionaryValue);
-  new_pref_entry->SetKey(kPrefKeyUsage, base::Value(kPrefKeyUsageCorporate));
+  if (key_locations.empty()) {
+    std::move(callback).Run(Status::kErrorKeyNotFound);
+    return;
+  }
 
-  update->SetWithoutPathExpansion(public_key_spki_der_b64,
-                                  std::move(new_pref_entry));
+  // A single key location is expected because this is intended for usage after
+  // key generation / import, when exactly one location is relevant.
+  DCHECK_EQ(key_locations.size(), 1U);
+
+  switch (key_locations[0]) {
+    case TokenId::kSystem:
+      // Nothing to do - all system-token keys are currently implicitly
+      // corporate.
+      std::move(callback).Run(Status::kSuccess);
+      return;
+    case TokenId::kUser: {
+      std::string public_key_spki_der_b64;
+      base::Base64Encode(public_key_spki_der, &public_key_spki_der_b64);
+
+      DictionaryPrefUpdate update(profile_prefs_, prefs::kPlatformKeys);
+
+      std::unique_ptr<base::DictionaryValue> new_pref_entry(
+          new base::DictionaryValue);
+      new_pref_entry->SetKey(kPrefKeyUsage,
+                             base::Value(kPrefKeyUsageCorporate));
+
+      update->SetWithoutPathExpansion(public_key_spki_der_b64,
+                                      std::move(new_pref_entry));
+      std::move(callback).Run(Status::kSuccess);
+      return;
+    }
+  }
 }
 
 // static
@@ -519,9 +729,9 @@ KeyPermissionsServiceImpl::GetCorporateKeyUsageAllowedAppIds(
 
 void KeyPermissionsServiceImpl::CreatePermissionObjectAndPassToCallback(
     const std::string& extension_id,
-    const PermissionsCallback& callback,
+    GetPermissionsForExtensionCallback callback,
     std::unique_ptr<base::Value> value) {
-  callback.Run(std::make_unique<PermissionsForExtensionImpl>(
+  std::move(callback).Run(std::make_unique<PermissionsForExtensionImpl>(
       extension_id, std::move(value), profile_prefs_, profile_policies_, this));
 }
 

chrome/browser/chromeos/platform_keys/key_permissions/key_permissions_service_impl.h

@@ -32,6 +32,10 @@ class PolicyService;
 namespace chromeos {
 namespace platform_keys {
 
+class PlatformKeysService;
+
+// TODO(crbug.com/1130949): Convert KeyPermissionsServiceImpl operations into
+// classes.
 class KeyPermissionsServiceImpl : public KeyPermissionsService {
  public:
   // Implementation of PermissionsForExtension.
@@ -54,21 +58,19 @@ class KeyPermissionsServiceImpl : public KeyPermissionsService {
         const PermissionsForExtensionImpl& other) = delete;
     ~PermissionsForExtensionImpl() override;
 
-    bool CanUseKeyForSigning(
-        const std::string& public_key_spki_der,
-        const std::vector<platform_keys::TokenId>& key_locations) override;
+    void CanUseKeyForSigning(const std::string& public_key_spki_der,
+                             CanUseKeyForSigningCallback callback) override;
+
+    void SetKeyUsedForSigning(const std::string& public_key_spki_der,
+                              SetKeyUsedForSigningCallback callback) override;
 
     void RegisterKeyForCorporateUsage(
         const std::string& public_key_spki_der,
-        const std::vector<platform_keys::TokenId>& key_locations) override;
+        RegisterKeyForCorporateUsageCallback callback) override;
 
     void SetUserGrantedPermission(
         const std::string& public_key_spki_der,
-        const std::vector<platform_keys::TokenId>& key_locations) override;
-
-    void SetKeyUsedForSigning(
-        const std::string& public_key_spki_der,
-        const std::vector<platform_keys::TokenId>& key_locations) override;
+        SetUserGrantedPermissionCallback callback) override;
 
    private:
     struct KeyEntry;
@@ -95,11 +97,44 @@ class KeyPermissionsServiceImpl : public KeyPermissionsService {
 
     bool PolicyAllowsCorporateKeyUsage() const;
 
+    void CanUseKeyForSigningWithLocations(
+        const std::string& public_key_spki_der,
+        CanUseKeyForSigningCallback callback,
+        const std::vector<TokenId>& key_locations,
+        Status key_locations_retrieval_status);
+    void CanUseKeyForSigningWithFlags(CanUseKeyForSigningCallback callback,
+                                      bool sign_unlimited_allowed,
+                                      bool is_corporate_key);
+
+    void SetKeyUsedForSigningWithLocations(
+        const std::string& public_key_spki_der,
+        SetKeyUsedForSigningCallback callback,
+        const std::vector<TokenId>& key_locations,
+        Status key_locations_retrieval_status);
+    void RegisterKeyForCorporateUsageWithLocations(
+        const std::string& public_key_spki_der,
+        RegisterKeyForCorporateUsageCallback callback,
+        const std::vector<TokenId>& key_locations,
+        Status key_locations_retrieval_status);
+
+    void SetUserGrantedPermissionWithLocations(
+        const std::string& public_key_spki_der,
+        SetUserGrantedPermissionCallback callback,
+        const std::vector<TokenId>& key_locations,
+        Status key_locations_retrieval_status);
+    void SetUserGrantedPermissionWithLocationsAndFlag(
+        const std::string& public_key_spki_der,
+        SetUserGrantedPermissionCallback callback,
+        const std::vector<TokenId>& key_locations,
+        Status key_locations_retrieval_status,
+        bool can_user_grant_permission);
+
     const std::string extension_id_;
     std::vector<KeyEntry> state_store_entries_;
     PrefService* const profile_prefs_;
     policy::PolicyService* const profile_policies_;
     KeyPermissionsServiceImpl* const key_permissions_service_;
+    base::WeakPtrFactory<PermissionsForExtensionImpl> weak_factory_{this};
   };
 
   // |profile_prefs| and |extensions_state_store| must not be null and must
@@ -111,7 +146,8 @@ class KeyPermissionsServiceImpl : public KeyPermissionsService {
   KeyPermissionsServiceImpl(bool profile_is_managed,
                             PrefService* profile_prefs,
                             policy::PolicyService* profile_policies,
-                            extensions::StateStore* extensions_state_store);
+                            extensions::StateStore* extensions_state_store,
+                            PlatformKeysService* platform_keys_service);
 
   ~KeyPermissionsServiceImpl() override;
 
@@ -119,19 +155,19 @@ class KeyPermissionsServiceImpl : public KeyPermissionsService {
   KeyPermissionsServiceImpl& operator=(const KeyPermissionsServiceImpl& other) =
       delete;
 
-  void GetPermissionsForExtension(const std::string& extension_id,
-                                  const PermissionsCallback& callback) override;
+  void GetPermissionsForExtension(
+      const std::string& extension_id,
+      GetPermissionsForExtensionCallback callback) override;
 
-  bool CanUserGrantPermissionFor(
+  void CanUserGrantPermissionForKey(
       const std::string& public_key_spki_der,
-      const std::vector<platform_keys::TokenId>& key_locations) const override;
+      CanUserGrantPermissionForKeyCallback callback) const override;
 
-  bool IsCorporateKey(
-      const std::string& public_key_spki_der,
-      const std::vector<platform_keys::TokenId>& key_locations) const override;
+  void IsCorporateKey(const std::string& public_key_spki_der,
+                      IsCorporateKeyCallback callback) const override;
 
   void SetCorporateKey(const std::string& public_key_spki_der,
-                       platform_keys::TokenId key_location) const override;
+                       SetCorporateKeyCallback callback) const override;
 
   // Returns true if |public_key_spki_der_b64| is a corporate usage key.
   // TOOD(http://crbug.com/1127284): Remove this and migrate callers to
@@ -150,17 +186,41 @@ class KeyPermissionsServiceImpl : public KeyPermissionsService {
   // and passes the object to |callback|.
   void CreatePermissionObjectAndPassToCallback(
       const std::string& extension_id,
-      const PermissionsCallback& callback,
+      GetPermissionsForExtensionCallback callback,
       std::unique_ptr<base::Value> value);
 
   // Writes |value| to the state store of the extension with id |extension_id|.
   void SetPlatformKeysOfExtension(const std::string& extension_id,
                                   std::unique_ptr<base::Value> value);
 
+  void CanUserGrantPermissionForKeyWithLocations(
+      const std::string& public_key_spki_der,
+      CanUserGrantPermissionForKeyCallback callback,
+      const std::vector<TokenId>& key_locations,
+      Status key_locations_retrieval_status) const;
+  void CanUserGrantPermissionForKeyWithLocationsAndFlag(
+      const std::string& public_key_spki_der,
+      CanUserGrantPermissionForKeyCallback callback,
+      const std::vector<TokenId>& key_locations,
+      Status key_locations_retrieval_status,
+      bool is_corporate_key);
+
+  void IsCorporateKeyWithLocations(const std::string& public_key_spki_der,
+                                   IsCorporateKeyCallback callback,
+                                   const std::vector<TokenId>& key_locations,
+                                   Status key_locations_retrieval_status) const;
+
+  void SetCorporateKeyWithLocations(
+      const std::string& public_key_spki_der,
+      SetCorporateKeyCallback callback,
+      const std::vector<TokenId>& key_locations,
+      Status key_locations_retrieval_status) const;
+
   const bool profile_is_managed_;
   PrefService* const profile_prefs_;
   policy::PolicyService* const profile_policies_;
   extensions::StateStore* const extensions_state_store_;
+  PlatformKeysService* const platform_keys_service_;
   base::WeakPtrFactory<KeyPermissionsServiceImpl> weak_factory_{this};
 };
 

chrome/browser/chromeos/platform_keys/key_permissions/key_permissions_service_impl_unittest.cc

@@ -8,16 +8,89 @@
 
 #include "base/command_line.h"
 #include "base/files/file_path.h"
+#include "chrome/browser/chromeos/platform_keys/mock_platform_keys_service.h"
+#include "chrome/browser/chromeos/platform_keys/platform_keys_service.h"
 #include "chrome/browser/extensions/test_extension_system.h"
 #include "chrome/test/base/testing_profile.h"
 #include "components/policy/core/common/mock_policy_service.h"
 #include "components/prefs/pref_service.h"
 #include "content/public/test/browser_task_environment.h"
 #include "extensions/browser/state_store.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace chromeos {
 namespace platform_keys {
+namespace {
+
+// A helper that waits until execution of an asynchronous KeyPermissionsService
+// operation has finished and provides access to the results.
+template <typename... ResultCallbackArgs>
+class ExecutionWaiter {
+ public:
+  ExecutionWaiter() = default;
+  ~ExecutionWaiter() = default;
+  ExecutionWaiter(const ExecutionWaiter& other) = delete;
+  ExecutionWaiter& operator=(const ExecutionWaiter& other) = delete;
+
+  // Returns the callback to be passed to the KeyPermissionsService operation
+  // invocation.
+  base::OnceCallback<void(ResultCallbackArgs... result_callback_args)>
+  GetCallback() {
+    return base::BindOnce(&ExecutionWaiter::OnExecutionDone,
+                          weak_ptr_factory_.GetWeakPtr());
+  }
+
+  // Waits until the callback returned by GetCallback() has been called.
+  void Wait() { run_loop_.Run(); }
+
+ protected:
+  // A std::tuple that is capable of storing the arguments passed to the result
+  // callback.
+  using ResultCallbackArgsTuple =
+      std::tuple<std::decay_t<ResultCallbackArgs>...>;
+
+  // Access to the arguments passed to the callback.
+  const ResultCallbackArgsTuple& result_callback_args() const {
+    EXPECT_TRUE(done_);
+    return result_callback_args_;
+  }
+
+ private:
+  void OnExecutionDone(ResultCallbackArgs... result_callback_args) {
+    EXPECT_FALSE(done_);
+    done_ = true;
+    result_callback_args_ = ResultCallbackArgsTuple(
+        std::forward<ResultCallbackArgs>(result_callback_args)...);
+    run_loop_.Quit();
+  }
+
+  base::RunLoop run_loop_;
+  bool done_ = false;
+  ResultCallbackArgsTuple result_callback_args_;
+
+  base::WeakPtrFactory<ExecutionWaiter> weak_ptr_factory_{this};
+};
+
+// Supports waiting for the result of KeyPermissionsService::IsCorporateKey.
+class IsCorporateKeyExecutionWaiter : public ExecutionWaiter<bool> {
+ public:
+  IsCorporateKeyExecutionWaiter() = default;
+  ~IsCorporateKeyExecutionWaiter() = default;
+
+  bool corporate() const { return std::get<0>(result_callback_args()); }
+};
+
+// Supports waiting for the result of KeyPermissionsService::SetCorporateKey.
+class SetCorporateKeyExecutionWaiter : public ExecutionWaiter<Status> {
+ public:
+  SetCorporateKeyExecutionWaiter() = default;
+  ~SetCorporateKeyExecutionWaiter() = default;
+
+  Status status() const { return std::get<0>(result_callback_args()); }
+};
+
+}  // namespace
 
 class KeyPermissionsServiceImplTest : public ::testing::Test {
  public:
@@ -43,12 +116,39 @@ class KeyPermissionsServiceImplTest : public ::testing::Test {
         /*autoupdate_enabled=*/false);
     extensions_state_store_ = extension_system->state_store();
 
+    platform_keys_service_ = std::make_unique<MockPlatformKeysService>();
     key_permissions_service_ = std::make_unique<KeyPermissionsServiceImpl>(
         /*profile_is_managed=*/true, profile_->GetPrefs(), policy_service_,
-        extensions_state_store_);
+        extensions_state_store_, platform_keys_service_.get());
   }
 
  protected:
+  void SetKeyLocations(const std::string& public_key,
+                       const std::vector<TokenId>& key_locations) {
+    ON_CALL(*platform_keys_service_, GetKeyLocations(public_key, testing::_))
+        .WillByDefault(testing::Invoke(
+            [key_locations](const std::string& public_key_spki_der,
+                            GetKeyLocationsCallback callback) {
+              std::move(callback).Run(std::move(key_locations),
+                                      Status::kSuccess);
+            }));
+  }
+
+  bool IsCorporateKey(const std::string& public_key) const {
+    IsCorporateKeyExecutionWaiter is_corporate_key_waiter;
+    key_permissions_service_->IsCorporateKey(
+        public_key, is_corporate_key_waiter.GetCallback());
+    is_corporate_key_waiter.Wait();
+    return is_corporate_key_waiter.corporate();
+  }
+
+  void SetCorporateKey(const std::string& public_key) {
+    SetCorporateKeyExecutionWaiter set_corporate_key_waiter;
+    key_permissions_service_->SetCorporateKey(
+        public_key, set_corporate_key_waiter.GetCallback());
+    set_corporate_key_waiter.Wait();
+  }
+
   content::BrowserTaskEnvironment task_environment_;
   std::unique_ptr<TestingProfile> profile_;
   // Owned by |profile_|.
@@ -57,30 +157,33 @@ class KeyPermissionsServiceImplTest : public ::testing::Test {
   extensions::StateStore* extensions_state_store_ = nullptr;
 
   std::unique_ptr<KeyPermissionsServiceImpl> key_permissions_service_;
+  std::unique_ptr<platform_keys::MockPlatformKeysService>
+      platform_keys_service_;
 };
 
 TEST_F(KeyPermissionsServiceImplTest, SystemTokenKeyIsImplicitlyCorporate) {
-  EXPECT_TRUE(key_permissions_service_->IsCorporateKey("some_public_key",
-                                                       {TokenId::kSystem}));
-  EXPECT_TRUE(key_permissions_service_->IsCorporateKey(
-      "some_public_key", {TokenId::kUser, TokenId::kSystem}));
+  const std::string kPublicKey = "test_public_key";
+
+  SetKeyLocations(kPublicKey, {TokenId::kSystem});
+  EXPECT_TRUE(IsCorporateKey(kPublicKey));
+
+  SetKeyLocations(kPublicKey, {TokenId::kUser, TokenId::kSystem});
+  EXPECT_TRUE(IsCorporateKey(kPublicKey));
 }
 
 TEST_F(KeyPermissionsServiceImplTest, CorporateRoundTrip) {
-  // By default, user-token keys are not corporate.
-  EXPECT_FALSE(key_permissions_service_->IsCorporateKey("some_public_key",
-                                                        {TokenId::kUser}));
+  const std::string kPublicKey = "test_public_key";
 
-  key_permissions_service_->SetCorporateKey("some_public_key", TokenId::kUser);
+  // By default, user-token keys are not corporate.
+  SetKeyLocations(kPublicKey, {TokenId::kUser});
+  EXPECT_FALSE(IsCorporateKey(kPublicKey));
 
-  EXPECT_TRUE(key_permissions_service_->IsCorporateKey("some_public_key",
-                                                       {TokenId::kUser}));
+  SetCorporateKey(kPublicKey);
+  EXPECT_TRUE(IsCorporateKey(kPublicKey));
 
   // Check that a repeated call doesn't corrupt the stored state.
-  key_permissions_service_->SetCorporateKey("some_public_key", TokenId::kUser);
-
-  EXPECT_TRUE(key_permissions_service_->IsCorporateKey("some_public_key",
-                                                       {TokenId::kUser}));
+  SetCorporateKey(kPublicKey);
+  EXPECT_TRUE(IsCorporateKey(kPublicKey));
 }
 
 }  // namespace platform_keys

chrome/browser/chromeos/platform_keys/key_permissions/mock_key_permissions_service.h

@@ -29,25 +29,25 @@ class MockKeyPermissionsService : public KeyPermissionsService {
   MOCK_METHOD(void,
               GetPermissionsForExtension,
               (const std::string& extension_id,
-               const PermissionsCallback& callback),
+               GetPermissionsForExtensionCallback callback),
               (override));
 
-  MOCK_METHOD(bool,
-              CanUserGrantPermissionFor,
+  MOCK_METHOD(void,
+              CanUserGrantPermissionForKey,
               (const std::string& public_key_spki_der,
-               const std::vector<platform_keys::TokenId>& key_locations),
+               CanUserGrantPermissionForKeyCallback callback),
               (const override));
 
-  MOCK_METHOD(bool,
+  MOCK_METHOD(void,
               IsCorporateKey,
               (const std::string& public_key_spki_der_b64,
-               const std::vector<platform_keys::TokenId>& key_locations),
+               IsCorporateKeyCallback callback),
               (const override));
 
   MOCK_METHOD(void,
               SetCorporateKey,
               (const std::string& public_key_spki_der_b64,
-               platform_keys::TokenId key_location),
+               SetCorporateKeyCallback callback),
               (const override));
 };
 

chrome/browser/chromeos/platform_keys/platform_keys.cc

@@ -55,6 +55,9 @@ std::string StatusToString(Status status) {
       return "Algorithm not supported.";
     case Status::kErrorCertificateNotFound:
       return "Certificate could not be found.";
+    case Status::kErrorGrantKeyPermissionForExtension:
+      return "Tried to grant permission for a key although prohibited (either "
+             "key is a corporate key or this account is managed).";
     case Status::kErrorInternal:
       return "Internal Error.";
     case Status::kErrorKeyAttributeRetrievalFailed:

chrome/browser/chromeos/platform_keys/platform_keys.h

@@ -43,6 +43,7 @@ enum class Status {
   kSuccess,
   kErrorAlgorithmNotSupported,
   kErrorCertificateNotFound,
+  kErrorGrantKeyPermissionForExtension,
   kErrorInternal,
   kErrorKeyAttributeRetrievalFailed,
   kErrorKeyAttributeSettingFailed,

chrome/browser/extensions/api/platform_keys/platform_keys_apitest_nss.cc

@@ -164,15 +164,27 @@ class PlatformKeysTest : public PlatformKeysTestBase {
     crypto::SetPrivateSoftwareSlotForChromeOSUserForTesting(std::move(slot));
   }
 
+  void OnKeyRegisteredForCorporateUsage(
+      std::unique_ptr<chromeos::platform_keys::KeyPermissionsService::
+                          PermissionsForExtension> permissions_for_ext,
+      const base::Closure& done_callback,
+      chromeos::platform_keys::Status status) {
+    ASSERT_EQ(status, chromeos::platform_keys::Status::kSuccess);
+    done_callback.Run();
+  }
+
   void GotPermissionsForExtension(
       const base::Closure& done_callback,
       std::unique_ptr<chromeos::platform_keys::KeyPermissionsService::
                           PermissionsForExtension> permissions_for_ext) {
+    auto* permissions_for_ext_unowned = permissions_for_ext.get();
     std::string client_cert1_spki =
         chromeos::platform_keys::GetSubjectPublicKeyInfo(client_cert1_);
-    permissions_for_ext->RegisterKeyForCorporateUsage(
-        client_cert1_spki, {chromeos::platform_keys::TokenId::kUser});
-    done_callback.Run();
+    permissions_for_ext_unowned->RegisterKeyForCorporateUsage(
+        client_cert1_spki,
+        base::BindOnce(&PlatformKeysTest::OnKeyRegisteredForCorporateUsage,
+                       base::Unretained(this), std::move(permissions_for_ext),
+                       done_callback));
   }
 
   void SetupTestCerts(const base::Closure& done_callback,