Add support to GCMKeyStore for multiple keys per app_id

components/gcm_driver/crypto/gcm_key_store.cc

Index: components/gcm_driver/crypto/gcm_key_store.cc
diff --git a/components/gcm_driver/crypto/gcm_key_store.cc b/components/gcm_driver/crypto/gcm_key_store.cc
index c24af360e83dcfcf53b26844875bede52e01054c..2b3a576351b3e8cc3bc0030f578007dc14c7460a 100644
--- a/components/gcm_driver/crypto/gcm_key_store.cc
+++ b/components/gcm_driver/crypto/gcm_key_store.cc
@@ -8,6 +8,8 @@
 
 #include <utility>
 
+#include "base/bind_helpers.h"
+#include "base/callback.h"
 #include "base/logging.h"
 #include "base/metrics/histogram_macros.h"
 #include "components/gcm_driver/crypto/p256_key_util.h"
@@ -16,6 +18,20 @@
 
 namespace gcm {
 
+namespace {
+
+std::string DatabaseKey(const std::string& app_id,
+                        const std::string& instance_id_authorized_entity) {
+  DCHECK_EQ(std::string::npos, app_id.find(','));
+  DCHECK_EQ(std::string::npos, instance_id_authorized_entity.find(','));
+  DCHECK_NE("*", instance_id_authorized_entity) << "Don't store wildcards";
+  return instance_id_authorized_entity.empty()
+             ? app_id  // No comma, for compatibility with existing keys.
+             : app_id + ',' + instance_id_authorized_entity;
+}
+
+}  // namespace
+
 // Statistics are logged to UMA with this string as part of histogram name. They
 // can all be found under LevelDB.*.GCMKeyStore. Changing this needs to
 // synchronize with histograms.xml, AND will also become incompatible with older
@@ -46,46 +62,73 @@ GCMKeyStore::GCMKeyStore(
 GCMKeyStore::~GCMKeyStore() {}
 
 void GCMKeyStore::GetKeys(const std::string& app_id,
+                          const std::string& instance_id_authorized_entity,
+                          bool fallback_to_empty_authorized_entity,
                           const KeysCallback& callback) {
   LazyInitialize(base::Bind(&GCMKeyStore::GetKeysAfterInitialize,
-                            weak_factory_.GetWeakPtr(), app_id, callback));
+                            weak_factory_.GetWeakPtr(), app_id,
+                            instance_id_authorized_entity,
+                            fallback_to_empty_authorized_entity, callback));
 }
 
-void GCMKeyStore::GetKeysAfterInitialize(const std::string& app_id,
-                                         const KeysCallback& callback) {
+void GCMKeyStore::GetKeysAfterInitialize(
+    const std::string& app_id,
+    const std::string& instance_id_authorized_entity,
+    bool fallback_to_empty_authorized_entity,
+    const KeysCallback& callback) {
   DCHECK(state_ == State::INITIALIZED || state_ == State::FAILED);
-  const auto& iter = key_pairs_.find(app_id);
+  bool success = false;
+
+  if (state_ == State::INITIALIZED) {
+    auto outer_iter = key_data_.find(app_id);
+    if (outer_iter != key_data_.end()) {
+      const auto& inner_map = outer_iter->second;
+      auto inner_iter = inner_map.find(instance_id_authorized_entity);
+      if (fallback_to_empty_authorized_entity && inner_iter == inner_map.end())
+        inner_iter = inner_map.find(std::string());
+      if (inner_iter != inner_map.end()) {
+        const KeyPairAndAuthSecret& key_and_auth = inner_iter->second;
+        callback.Run(key_and_auth.first, key_and_auth.second);
+        success = true;
+      }
+    }
+  }
 
-  const bool success = state_ == State::INITIALIZED && iter != key_pairs_.end();
   UMA_HISTOGRAM_BOOLEAN("GCM.Crypto.GetKeySuccessRate", success);
-
-  if (!success) {
+  if (!success)
     callback.Run(KeyPair(), std::string() /* auth_secret */);
-    return;
-  }
-
-  const auto& auth_secret_iter = auth_secrets_.find(app_id);
-  DCHECK(auth_secret_iter != auth_secrets_.end());
-
-  callback.Run(iter->second, auth_secret_iter->second);
 }
 
 void GCMKeyStore::CreateKeys(const std::string& app_id,
+                             const std::string& instance_id_authorized_entity,
                              const KeysCallback& callback) {
   LazyInitialize(base::Bind(&GCMKeyStore::CreateKeysAfterInitialize,
-                            weak_factory_.GetWeakPtr(), app_id, callback));
+                            weak_factory_.GetWeakPtr(), app_id,
+                            instance_id_authorized_entity, callback));
 }
 
-void GCMKeyStore::CreateKeysAfterInitialize(const std::string& app_id,
-                                            const KeysCallback& callback) {
+void GCMKeyStore::CreateKeysAfterInitialize(
+    const std::string& app_id,
+    const std::string& instance_id_authorized_entity,
+    const KeysCallback& callback) {
   DCHECK(state_ == State::INITIALIZED || state_ == State::FAILED);
   if (state_ != State::INITIALIZED) {
     callback.Run(KeyPair(), std::string() /* auth_secret */);
     return;
   }
 
-  // Only allow creating new keys if no keys currently exist.
-  DCHECK_EQ(0u, key_pairs_.count(app_id));
+  // Only allow creating new keys if no keys currently exist. Multiple Instance
+  // ID tokens can share an app_id (with different authorized entities), but
+  // Instance ID tokens cannot share an app_id with a legacy GCM registration.
+  // This invariant is necessary for the fallback_to_empty_authorized_entity
+  // mode of GetKey (needed by GCMEncryptionProvider::DecryptMessage, which
+  // can't distinguish Instance ID tokens from legacy GCM registrations).
+  DCHECK(!key_data_.count(app_id) ||
+         (!instance_id_authorized_entity.empty() &&
+          !key_data_[app_id].count(instance_id_authorized_entity) &&
+          !key_data_[app_id].count(std::string())))
+      << "Instance ID tokens cannot share an app_id with a legacy GCM "
+         "registration";
 
   std::string private_key, public_key_x509, public_key;
   if (!CreateP256KeyPair(&private_key, &public_key_x509, &public_key)) {
@@ -105,6 +148,10 @@ void GCMKeyStore::CreateKeysAfterInitialize(const std::string& app_id,
   // Store the keys in a new EncryptionData object.
   EncryptionData encryption_data;
   encryption_data.set_app_id(app_id);
+  if (!instance_id_authorized_entity.empty()) {
+    encryption_data.set_instance_id_authorized_entity(
+        instance_id_authorized_entity);
+  }
   encryption_data.set_auth_secret(auth_secret);
 
   KeyPair* pair = encryption_data.add_keys();
@@ -115,8 +162,7 @@ void GCMKeyStore::CreateKeysAfterInitialize(const std::string& app_id,
 
   // Write them immediately to our cache, so subsequent calls to
   // {Get/Create/Remove}Keys can see them.
-  key_pairs_[app_id] = *pair;
-  auth_secrets_[app_id] = auth_secret;
+  key_data_[app_id][instance_id_authorized_entity] = {*pair, auth_secret};
 
   using EntryVectorType =
       leveldb_proto::ProtoDatabase<EncryptionData>::KeyEntryVector;
@@ -125,7 +171,8 @@ void GCMKeyStore::CreateKeysAfterInitialize(const std::string& app_id,
   std::unique_ptr<std::vector<std::string>> keys_to_remove(
       new std::vector<std::string>());
 
-  entries_to_save->push_back(std::make_pair(app_id, encryption_data));
+  entries_to_save->push_back(std::make_pair(
+      DatabaseKey(app_id, instance_id_authorized_entity), encryption_data));
 
   database_->UpdateEntries(
       std::move(entries_to_save), std::move(keys_to_remove),
@@ -153,31 +200,56 @@ void GCMKeyStore::DidStoreKeys(const KeyPair& pair,
 }
 
 void GCMKeyStore::RemoveKeys(const std::string& app_id,
+                             const std::string& instance_id_authorized_entity,
                              const base::Closure& callback) {
   LazyInitialize(base::Bind(&GCMKeyStore::RemoveKeysAfterInitialize,
-                            weak_factory_.GetWeakPtr(), app_id, callback));
+                            weak_factory_.GetWeakPtr(), app_id,
+                            instance_id_authorized_entity, callback));
 }
 
-void GCMKeyStore::RemoveKeysAfterInitialize(const std::string& app_id,
-                                            const base::Closure& callback) {
+void GCMKeyStore::RemoveKeysAfterInitialize(
+    const std::string& app_id,
+    const std::string& instance_id_authorized_entity,
+    const base::Closure& callback) {
   DCHECK(state_ == State::INITIALIZED || state_ == State::FAILED);
-  const auto iter = key_pairs_.find(app_id);
-  if (iter == key_pairs_.end() || state_ != State::INITIALIZED) {
+
+  const auto& outer_iter = key_data_.find(app_id);
+  if (outer_iter == key_data_.end() || state_ != State::INITIALIZED) {
     callback.Run();
     return;
   }
 
-  // Clear them immediately from our cache, so subsequent calls to
-  // {Get/Create/Remove}Keys don't see them.
-  key_pairs_.erase(app_id);
-  auth_secrets_.erase(app_id);
-
   using EntryVectorType =
       leveldb_proto::ProtoDatabase<EncryptionData>::KeyEntryVector;
 
   std::unique_ptr<EntryVectorType> entries_to_save(new EntryVectorType());
   std::unique_ptr<std::vector<std::string>> keys_to_remove(
-      new std::vector<std::string>(1, app_id));
+      new std::vector<std::string>());
+
+  bool had_keys = false;
+  auto& inner_map = outer_iter->second;
+  for (auto it = inner_map.begin(); it != inner_map.end();) {
+    // Wildcard "*" matches all non-empty authorized entities (not legacy GCM).
+    if (instance_id_authorized_entity == "*"
+            ? !it->first.empty()
+            : it->first == instance_id_authorized_entity) {
+      had_keys = true;
+
+      keys_to_remove->push_back(DatabaseKey(app_id, it->first));
+
+      // Clear keys immediately from our cache, so subsequent calls to
+      // {Get/Create/Remove}Keys don't see them.
+      it = inner_map.erase(it);
+    } else {
+      ++it;
+    }
+  }
+  if (!had_keys) {
+    callback.Run();
+    return;
+  }
+  if (inner_map.empty())
+    key_data_.erase(app_id);
 
   database_->UpdateEntries(std::move(entries_to_save),
                            std::move(keys_to_remove),
@@ -247,8 +319,11 @@ void GCMKeyStore::DidLoadKeys(
   for (const EncryptionData& entry : *entries) {
     DCHECK_EQ(1, entry.keys_size());
 
-    key_pairs_[entry.app_id()] = entry.keys(0);
-    auth_secrets_[entry.app_id()] = entry.auth_secret();
+    std::string instance_id_authorized_entity;
+    if (entry.has_instance_id_authorized_entity())
+      instance_id_authorized_entity = entry.instance_id_authorized_entity();
+    key_data_[entry.app_id()][instance_id_authorized_entity] = {
+        entry.keys(0), entry.auth_secret()};
   }
 
   state_ = State::INITIALIZED;