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/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.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,