Add support to GCMKeyStore for multiple keys per app_id

components/gcm_driver/crypto/gcm_key_store.h

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #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"
 #include "base/files/file_path.h"
 #include "base/macros.h"
 #include "base/memory/ref_counted.h"
 #include "base/memory/weak_ptr.h"
 #include "components/gcm_driver/crypto/proto/gcm_encryption_data.pb.h"
 #include "components/gcm_driver/gcm_delayed_task_controller.h"
 
 namespace base {
 class SequencedTaskRunner;
 }
 
 namespace leveldb_proto {
 template <typename T>
 class ProtoDatabase;
 }
 
 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
 // rather than returning the result. Do not rely on the timing of the callbacks.
 class GCMKeyStore {
  public:
   using KeysCallback = base::Callback<void(const KeyPair& pair,
                                            const std::string& auth_secret)>;
 
   GCMKeyStore(
       const base::FilePath& key_store_path,
       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);
+  // 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 |app_id|, and invokes
-  // |callback| when they are available, or when an error occurred.
-  void CreateKeys(const std::string& app_id, 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 |app_id|, and invokes |callback| when
-  // the operation has finished.
-  void RemoveKeys(const std::string& app_id, const base::Closure& 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.
   void LazyInitialize(const base::Closure& done_closure);
 
   void DidInitialize(bool success);
   void DidLoadKeys(bool success,
                    std::unique_ptr<std::vector<EncryptionData>> entries);
 
   void DidStoreKeys(const KeyPair& pair,
                     const std::string& auth_secret,
                     const KeysCallback& callback,
                     bool success);
 
   void DidRemoveKeys(const base::Closure& callback, bool success);
 
   // Private implementations of the API that will be executed when the database
   // 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_;
 
   // Blocking task runner which the database will do I/O operations on.
   scoped_refptr<base::SequencedTaskRunner> blocking_task_runner_;
 
   // Instance of the ProtoDatabase backing the key store.
   std::unique_ptr<leveldb_proto::ProtoDatabase<EncryptionData>> database_;
 
   enum class State;
 
   // The current state of the database. It has to be initialized before use.
   State state_;
 
   // Controller for tasks that should be executed once the key store has
   // 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_;
 
   DISALLOW_COPY_AND_ASSIGN(GCMKeyStore);
 };
 
 }  // namespace gcm
 
 #endif  // COMPONENTS_GCM_DRIVER_CRYPTO_GCM_KEY_STORE_H_