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..f67f59415f94bfa7f7a13df79dd753dfc0909d16 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,35 @@ 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. |authorized_entity| should be the InstanceID
+  // token's authorized entity, or "" for non-InstanceID GCM registrations. If
+  // |fallback_to_empty_authorized_entity| is true and the keys are not found,
+  // GetKeys will try again with an empty authorized entity; this can be used
+  // when it's not known whether or not the |app_id| is for an InstanceID.
+  void GetKeys(const std::string& app_id,
+               const std::string& 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. |authorized_entity| should be the InstanceID
+  // token's authorized entity, or "" for non-InstanceID GCM registrations.
+  // Simultaneously using the same |app_id| for both a non-InstanceID GCM
+  // registration and one or more InstanceID tokens is not supported.
+  void CreateKeys(const std::string& app_id,
+                  const std::string& authorized_entity,
+                  const KeysCallback& callback);
+
+  // Removes the keys associated with the |app_id| + |authorized_entity| pair,
+  // and invokes |callback| when the operation has finished. |authorized_entity|
+  // should be the InstanceID token's authorized entity, or "*" to remove for
+  // all InstanceID tokens, or "" for non-InstanceID GCM registrations.
+  void RemoveKeys(const std::string& app_id,
+                  const std::string& authorized_entity,
+                  const base::Closure& callback);
 
  private:
   // Initializes the database if necessary, and runs |done_closure| when done.
@@ -77,10 +96,14 @@ class GCMKeyStore {
   // has either been successfully loaded, or failed to load.
 
   void GetKeysAfterInitialize(const std::string& app_id,
+                              const std::string& authorized_entity,
+                              bool fallback_to_empty_authorized_entity,
                               const KeysCallback& callback);
   void CreateKeysAfterInitialize(const std::string& app_id,
+                                 const std::string& authorized_entity,
                                  const KeysCallback& callback);
   void RemoveKeysAfterInitialize(const std::string& app_id,
+                                 const std::string& authorized_entity,
                                  const base::Closure& callback);
 
   // Path in which the key store database will be saved.
@@ -101,10 +124,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 authorized_entity to the loaded key
+  // pair and authentication secrets.
+  using KeyPairAndAuthSecret = std::pair<KeyPair, std::string>;
+  std::unordered_map<std::string,
+                     std::unordered_map<std::string, KeyPairAndAuthSecret>>
+      key_data_;
 
   base::WeakPtrFactory<GCMKeyStore> weak_factory_;