Add support to GCMKeyStore for multiple keys per app_id

components/gcm_driver/crypto/gcm_key_store.h

Index: components/gcm_driver/crypto/gcm_key_store.h
diff --git a/components/gcm_driver/crypto/gcm_key_store.h b/components/gcm_driver/crypto/gcm_key_store.h
index 05efc034ad096d7b48bf32708306c5e8ca594c83..b83c3b9b1f933eb6be6eda33a65bd859b5ff4c69 100644
--- a/components/gcm_driver/crypto/gcm_key_store.h
+++ b/components/gcm_driver/crypto/gcm_key_store.h
@@ -5,9 +5,10 @@
 #ifndef COMPONENTS_GCM_DRIVER_CRYPTO_GCM_KEY_STORE_H_
 #define COMPONENTS_GCM_DRIVER_CRYPTO_GCM_KEY_STORE_H_
 
-#include <map>
 #include <memory>
 #include <string>
+#include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "base/callback_forward.h"
@@ -31,7 +32,7 @@ namespace gcm {
 
 // Key storage for use with encrypted messages received from Google Cloud
 // Messaging. It provides the ability to create and store a key-pair for a given
-// app id, as well as retrieving and deleting key-pairs.
+// app id + authorized entity pair, and to retrieve and delete key-pairs.
 //
 // This class is backed by a proto database and might end up doing file I/O on
 // a background task runner. For this reason, all public APIs take a callback
@@ -46,17 +47,38 @@ class GCMKeyStore {
       const scoped_refptr<base::SequencedTaskRunner>& blocking_task_runner);
   ~GCMKeyStore();
 
-  // Retrieves the public/private key-pair associated with |app_id|, and
-  // invokes |callback| when they are available, or when an error occurred.
-  void GetKeys(const std::string& app_id, const KeysCallback& callback);
-
-  // Creates a new public/private key-pair for |app_id|, and invokes
-  // |callback| when they are available, or when an error occurred.
-  void CreateKeys(const std::string& app_id, const KeysCallback& callback);
-
-  // Removes the keys associated with |app_id|, and invokes |callback| when
-  // the operation has finished.
-  void RemoveKeys(const std::string& app_id, const base::Closure& callback);
+  // Retrieves the public/private key-pair associated with the |app_id| +
+  // authorized entity pair, and invokes |callback| when they are available, or
+  // when an error occurred.
+  // |instance_id_authorized_entity|: pass InstanceID token's authorized_entity
+  //                                  or "" for legacy GCM registrations.
+  // |fallback_to_empty_authorized_entity|: if true and the keys are not found,
+  //                                        will try again with empty authorized
+  //                                        entity (use this when you're not
+  //                                        sure if you have an Instance ID).
+  void GetKeys(const std::string& app_id,
+               const std::string& instance_id_authorized_entity,
+               bool fallback_to_empty_authorized_entity,
+               const KeysCallback& callback);
+
+  // Creates a new public/private key-pair for the |app_id| + authorized entity
+  // pair, and invokes |callback| when they are available, or when an error
+  // occurred. Simultaneously using the same |app_id| for both a legacy GCM
+  // registration and one or more InstanceID tokens is not supported.
+  // |instance_id_authorized_entity|: pass InstanceID token's authorized_entity
+  //                                  or "" for legacy GCM registrations.
+  void CreateKeys(const std::string& app_id,
+                  const std::string& instance_id_authorized_entity,
+                  const KeysCallback& callback);
+
+  // Removes the keys associated with the |app_id| + authorized entity pair, and
+  // invokes |callback| when the operation has finished.
+  // |instance_id_authorized_entity|: pass InstanceID token's authorized_entity
+  //                                  or "*" to remove for all InstanceID tokens
+  //                                  or "" for legacy GCM registrations.
+  void RemoveKeys(const std::string& app_id,
+                  const std::string& instance_id_authorized_entity,
+                  const base::Closure& callback);
 
  private:
   // Initializes the database if necessary, and runs |done_closure| when done.
@@ -77,11 +99,17 @@ class GCMKeyStore {
   // has either been successfully loaded, or failed to load.
 
   void GetKeysAfterInitialize(const std::string& app_id,
+                              const std::string& instance_id_authorized_entity,
+                              bool fallback_to_empty_authorized_entity,
                               const KeysCallback& callback);
-  void CreateKeysAfterInitialize(const std::string& app_id,
-                                 const KeysCallback& callback);
-  void RemoveKeysAfterInitialize(const std::string& app_id,
-                                 const base::Closure& callback);
+  void CreateKeysAfterInitialize(
+      const std::string& app_id,
+      const std::string& instance_id_authorized_entity,
+      const KeysCallback& callback);
+  void RemoveKeysAfterInitialize(
+      const std::string& app_id,
+      const std::string& instance_id_authorized_entity,
+      const base::Closure& callback);
 
   // Path in which the key store database will be saved.
   base::FilePath key_store_path_;
@@ -101,10 +129,12 @@ class GCMKeyStore {
   // finished initializing.
   GCMDelayedTaskController delayed_task_controller_;
 
-  // Mapping of an app id to the loaded key pair and authentication secrets.
-  // TODO(peter): Switch these to std::unordered_map<> once allowed.
-  std::map<std::string, KeyPair> key_pairs_;
-  std::map<std::string, std::string> auth_secrets_;
+  // Nested map from app_id to a map from instance_id_authorized_entity to the
+  // loaded key pair and authentication secrets.
+  using KeyPairAndAuthSecret = std::pair<KeyPair, std::string>;
+  std::unordered_map<std::string,

[Peter Beverloo, 2016/05/09 14:10:10]

Why can't we use a multimap from |app_id| to a tuple of (|authorized_entity|,
|key_pair|, |auth_secret|)?

It seems to me like that could help remove the outer_iter/inner_iter code from
the GetKeys() and RemoveKeys() paths. This operation is so infrequent, and
operates on (generally) such small amounts of data, that losing the O(1)ish
performance of the unordered_maps doesn't really matter here, and we can
optimize for maintainability instead.

[johnme, 2016/05/09 18:15:55]

I considered that, but I'm not convinced a multimap would be more maintainable:

- GetKeysAfterInitialize would be at least as complex, since it would need to
loop though the equal_range of keys, then loop through a second time for
fallback (whereas fallback is currently a one-liner).

- CreateKeysAfterInitialize would be more complex, since it would no longer be
possible to express conditions like !key_data_[app_id].count(authorized_entity)
in the DCHECK (except perhaps using lambda functions, but that wouldn't improve
legibility).

- RemoveKeysAfterInitialize would be about equally complex (looping through
equal_range instead of looping through inner map).

+                     std::unordered_map<std::string, KeyPairAndAuthSecret>>
+      key_data_;
 
   base::WeakPtrFactory<GCMKeyStore> weak_factory_;