[Sync] Add support for keystore key rotation.

Key rotation will trigger a full re-encryption of all sync data with the newest
keystore key. Previous keys will be added to the keybag as well. We detect key
rotation by checking whether we have multiple keystore keys and the nigori's
keybag is not encrypted with the current keystore key. In that case, we no
longer support backwards compatibility with non-keystore supporting versions,
and re-encrypt using the newest keystore key.

This change also fixes two latent issues: lack of encryption of the keystore
bootstrap and not properly posting OnPassphraseAccepted when we resolve pending
keys.

BUG=163744


Review URL: https://chromiumcodereview.appspot.com/11434070

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@173830 0039d316-1c4b-4281-b951-d872f2087c98

3 files changed, 848 insertions, 89 deletions

diff --git a/sync/internal_api/sync_encryption_handler_impl.cc b/sync/internal_api/sync_encryption_handler_impl.cc
index 22f5028..852f5b6 100644
--- a/sync/internal_api/sync_encryption_handler_impl.cc
+++ b/sync/internal_api/sync_encryption_handler_impl.cc
@@ -9,6 +9,7 @@
 
 #include "base/base64.h"
 #include "base/bind.h"
+#include "base/json/json_string_value_serializer.h"
 #include "base/message_loop.h"
 #include "base/time.h"
 #include "base/tracked_objects.h"
@@ -28,6 +29,7 @@
 #include "sync/syncable/entry.h"
 #include "sync/syncable/nigori_util.h"
 #include "sync/util/cryptographer.h"
+#include "sync/util/encryptor.h"
 #include "sync/util/time.h"
 
 namespace syncer {
@@ -41,6 +43,19 @@ namespace {
 // should be, in which case they might ping pong (see crbug.com/119207).
 static const int kNigoriOverwriteLimit = 10;
 
+// Enumeration of nigori keystore migration results (for use in UMA stats).
+enum NigoriMigrationResult {
+  FAILED_TO_SET_DEFAULT_KEYSTORE,
+  FAILED_TO_SET_NONDEFAULT_KEYSTORE,
+  FAILED_TO_EXTRACT_DECRYPTOR,
+  FAILED_TO_EXTRACT_KEYBAG,
+  MIGRATION_SUCCESS_KEYSTORE_DEFAULT,
+  MIGRATION_SUCCESS_KEYSTORE_NONDEFAULT,
+  MIGRATION_SUCCESS_FROZEN_IMPLICIT,
+  MIGRATION_SUCCESS_CUSTOM,
+  MIGRATION_RESULT_SIZE,
+};
+
 // The new passphrase state is sufficient to determine whether a nigori node
 // is migrated to support keystore encryption. In addition though, we also
 // want to verify the conditions for proper keystore encryption functionality.
@@ -105,6 +120,71 @@ bool IsExplicitPassphrase(PassphraseType type) {
   return type == CUSTOM_PASSPHRASE || type == FROZEN_IMPLICIT_PASSPHRASE;
 }
 
+// Keystore Bootstrap Token helper methods.
+// The bootstrap is a base64 encoded, encrypted, ListValue of keystore key
+// strings, with the current keystore key as the last value in the list.
+std::string PackKeystoreBootstrapToken(
+    const std::vector<std::string>& old_keystore_keys,
+    const std::string& current_keystore_key,
+    Encryptor* encryptor) {
+  if (current_keystore_key.empty())
+    return "";
+
+  base::ListValue keystore_key_values;
+  for (size_t i = 0; i < old_keystore_keys.size(); ++i)
+    keystore_key_values.AppendString(old_keystore_keys[i]);
+  keystore_key_values.AppendString(current_keystore_key);
+
+  // Update the bootstrap token.
+  // The bootstrap is a base64 encoded, encrypted, ListValue of keystore key
+  // strings, with the current keystore key as the last value in the list.
+  std::string serialized_keystores;
+  JSONStringValueSerializer json(&serialized_keystores);
+  json.Serialize(keystore_key_values);
+  std::string encrypted_keystores;
+  encryptor->EncryptString(serialized_keystores,
+                           &encrypted_keystores);
+  std::string keystore_bootstrap;
+  base::Base64Encode(encrypted_keystores, &keystore_bootstrap);
+  return keystore_bootstrap;
+}
+
+bool UnpackKeystoreBootstrapToken(
+    const std::string& keystore_bootstrap_token,
+    Encryptor* encryptor,
+    std::vector<std::string>* old_keystore_keys,
+    std::string* current_keystore_key) {
+  if (keystore_bootstrap_token.empty())
+    return false;
+  std::string base64_decoded_keystore_bootstrap;
+  if (!base::Base64Decode(keystore_bootstrap_token,
+                          &base64_decoded_keystore_bootstrap)) {
+    return false;
+  }
+  std::string decrypted_keystore_bootstrap;
+  if (!encryptor->DecryptString(base64_decoded_keystore_bootstrap,
+                                &decrypted_keystore_bootstrap)) {
+    return false;
+  }
+  JSONStringValueSerializer json(&decrypted_keystore_bootstrap);
+  scoped_ptr<base::Value> deserialized_keystore_keys(
+      json.Deserialize(NULL, NULL));
+  if (!deserialized_keystore_keys.get())
+    return false;
+  base::ListValue* internal_list_value = NULL;
+  if (!deserialized_keystore_keys->GetAsList(&internal_list_value))
+    return false;
+  int number_of_keystore_keys = internal_list_value->GetSize();
+  if (!internal_list_value->GetString(number_of_keystore_keys - 1,
+                                      current_keystore_key)) {
+    return false;
+  }
+  old_keystore_keys->resize(number_of_keystore_keys - 1);
+  for (int i = 0; i < number_of_keystore_keys - 1; ++i)
+    internal_list_value->GetString(i, &(*old_keystore_keys)[i]);
+  return true;
+}
+
 }  // namespace
 
 SyncEncryptionHandlerImpl::Vault::Vault(
@@ -127,11 +207,18 @@ SyncEncryptionHandlerImpl::SyncEncryptionHandlerImpl(
       vault_unsafe_(encryptor, SensitiveTypes()),
       encrypt_everything_(false),
       passphrase_type_(IMPLICIT_PASSPHRASE),
-      keystore_key_(restored_keystore_key_for_bootstrapping),
       nigori_overwrite_count_(0) {
-  // We only bootstrap the user provided passphrase. The keystore key is handled
-  // at Init time once we're sure the nigori is downloaded.
+  // Restore the cryptographer's previous keys. Note that we don't add the
+  // keystore keys into the cryptographer here, in case a migration was pending.
   vault_unsafe_.cryptographer.Bootstrap(restored_key_for_bootstrapping);
+
+  // If this fails, we won't have a valid keystore key, and will simply request
+  // new ones from the server on the next DownloadUpdates.
+  UnpackKeystoreBootstrapToken(
+      restored_keystore_key_for_bootstrapping,
+      encryptor,
+      &old_keystore_keys_,
+      &keystore_key_);
 }
 
 SyncEncryptionHandlerImpl::~SyncEncryptionHandlerImpl() {}
@@ -528,45 +615,71 @@ bool SyncEncryptionHandlerImpl::NeedKeystoreKey(
   return keystore_key_.empty();
 }
 
-bool SyncEncryptionHandlerImpl::SetKeystoreKey(
-    const std::string& key,
+bool SyncEncryptionHandlerImpl::SetKeystoreKeys(
+    const google::protobuf::RepeatedPtrField<google::protobuf::string>& keys,
     syncable::BaseTransaction* const trans) {
   DCHECK(thread_checker_.CalledOnValidThread());
-  if (!keystore_key_.empty() || key.empty())
+  if (keys.size() == 0)
     return false;
-  // Base64 encode so we can persist the keystore key directly via preferences.
-  if (!base::Base64Encode(key, &keystore_key_)) {
-    LOG(ERROR) << "Failed to base64 encode keystore key.";
-    keystore_key_ = "";
+  // The last key in the vector is the current keystore key. The others are kept
+  // around for decryption only.
+  const std::string& raw_keystore_key = keys.Get(keys.size() - 1);
+  if (raw_keystore_key.empty())
     return false;
-  }
 
-  DVLOG(1) << "Keystore bootstrap token updated.";
+  // Note: in order to Pack the keys, they must all be base64 encoded (else
+  // JSON serialization fails).
+  if (!base::Base64Encode(raw_keystore_key, &keystore_key_))
+    return false;
+
+  // Go through and save the old keystore keys. We always persist all keystore
+  // keys the server sends us.
+  old_keystore_keys_.resize(keys.size() - 1);
+  for (int i = 0; i < keys.size() - 1; ++i)
+    base::Base64Encode(keys.Get(i), &old_keystore_keys_[i]);
+
+  Cryptographer* cryptographer = &UnlockVaultMutable(trans)->cryptographer;
+
+  // Update the bootstrap token. If this fails, we persist an empty string,
+  // which will force us to download the keystore keys again on the next
+  // restart.
+  std::string keystore_bootstrap = PackKeystoreBootstrapToken(
+      old_keystore_keys_,
+      keystore_key_,
+      cryptographer->encryptor());
+  DCHECK_EQ(keystore_bootstrap.empty(), keystore_key_.empty());
   FOR_EACH_OBSERVER(SyncEncryptionHandler::Observer, observers_,
-                    OnBootstrapTokenUpdated(keystore_key_,
+                    OnBootstrapTokenUpdated(keystore_bootstrap,
                                             KEYSTORE_BOOTSTRAP_TOKEN));
+  DVLOG(1) << "Keystore bootstrap token updated.";
 
-  Cryptographer* cryptographer = &UnlockVaultMutable(trans)->cryptographer;
+  // If this is a first time sync, we get the encryption keys before we process
+  // the nigori node. Just return for now, ApplyNigoriUpdate will be invoked
+  // once we have the nigori node.
   syncable::Entry entry(trans, syncable::GET_BY_SERVER_TAG, kNigoriTag);
-  if (entry.good()) {
-    const sync_pb::NigoriSpecifics& nigori =
-        entry.Get(syncable::SPECIFICS).nigori();
-    if (cryptographer->has_pending_keys() &&
-        IsNigoriMigratedToKeystore(nigori) &&
-        !nigori.keystore_decryptor_token().blob().empty()) {
-      // If the nigori is already migrated and we have pending keys, we might
-      // be able to decrypt them using the keystore decryptor token.
-      DecryptPendingKeysWithKeystoreKey(keystore_key_,
-                                        nigori.keystore_decryptor_token(),
-                                        cryptographer);
-    } else if (ShouldTriggerMigration(nigori, *cryptographer)) {
-      // We call rewrite nigori to attempt to trigger migration.
-      // Need to post a task to open a new sync_api transaction.
-      MessageLoop::current()->PostTask(
-          FROM_HERE,
-          base::Bind(&SyncEncryptionHandlerImpl::RewriteNigori,
-                     weak_ptr_factory_.GetWeakPtr()));
-    }
+  if (!entry.good())
+    return true;
+
+  const sync_pb::NigoriSpecifics& nigori =
+      entry.Get(syncable::SPECIFICS).nigori();
+  if (cryptographer->has_pending_keys() &&
+      IsNigoriMigratedToKeystore(nigori) &&
+      !nigori.keystore_decryptor_token().blob().empty()) {
+    // If the nigori is already migrated and we have pending keys, we might
+    // be able to decrypt them using either the keystore decryptor token
+    // or the existing keystore keys.
+    DecryptPendingKeysWithKeystoreKey(keystore_key_,
+                                      nigori.keystore_decryptor_token(),
+                                      cryptographer);
+  }
+
+  // Note that triggering migration will have no effect if we're already
+  // properly migrated with the newest keystore keys.
+  if (ShouldTriggerMigration(nigori, *cryptographer)) {
+    MessageLoop::current()->PostTask(
+        FROM_HERE,
+        base::Bind(&SyncEncryptionHandlerImpl::RewriteNigori,
+                   weak_ptr_factory_.GetWeakPtr()));
   }
 
   return true;
@@ -1150,9 +1263,21 @@ bool SyncEncryptionHandlerImpl::ShouldTriggerMigration(
       // We need to overwrite the keybag. This might involve overwriting the
       // keystore decryptor too.
       return true;
-    } else {
-      return false;
+    } else if (old_keystore_keys_.size() > 0 && !keystore_key_.empty()) {
+      // Check to see if a server key rotation has happened, but the nigori
+      // node's keys haven't been rotated yet, and hence we should re-migrate.
+      // Note that once a key rotation has been performed, we no longer
+      // preserve backwards compatibility, and the keybag will therefore be
+      // encrypted with the current keystore key.
+      Cryptographer temp_cryptographer(cryptographer.encryptor());
+      KeyParams keystore_params = {"localhost", "dummy", keystore_key_};
+      temp_cryptographer.AddKey(keystore_params);
+      if (!temp_cryptographer.CanDecryptUsingDefaultKey(
+              nigori.encryption_keybag())) {
+        return true;
+      }
     }
+    return false;
   } else if (keystore_key_.empty()) {
     // If we haven't already migrated, we don't want to do anything unless
     // a keystore key is available (so that those clients without keystore
@@ -1202,9 +1327,43 @@ bool SyncEncryptionHandlerImpl::AttemptToMigrateNigoriToKeystore(
 
   if (!keystore_key_.empty()) {
     KeyParams key_params = {"localhost", "dummy", keystore_key_};
-    if (!cryptographer->AddNonDefaultKey(key_params)) {
-      LOG(ERROR) << "Failed to add keystore key as non-default key.";
-      return false;
+    if (old_keystore_keys_.size() > 0 &&
+        new_passphrase_type == KEYSTORE_PASSPHRASE) {
+      // At least one key rotation has been performed, so we no longer care
+      // about backwards compatibility. Ensure the keystore key is the default
+      // key.
+      DVLOG(1) << "Migrating keybag to keystore key.";
+      if (!cryptographer->AddKey(key_params)) {
+        LOG(ERROR) << "Failed to add keystore key as default key";
+        UMA_HISTOGRAM_ENUMERATION("Sync.AttemptNigoriMigration",
+                                  FAILED_TO_SET_DEFAULT_KEYSTORE,
+                                  MIGRATION_RESULT_SIZE);
+        return false;
+      }
+    } else {
+      // We're in backwards compatible mode -- either the account has an
+      // explicit passphrase, or we want to preserve the current GAIA-based key
+      // as the default because we can (there have been no key rotations since
+      // the migration).
+      DVLOG(1) << "Migrating keybag while preserving old key";
+      if (!cryptographer->AddNonDefaultKey(key_params)) {
+        LOG(ERROR) << "Failed to add keystore key as non-default key.";
+        UMA_HISTOGRAM_ENUMERATION("Sync.AttemptNigoriMigration",
+                                  FAILED_TO_SET_NONDEFAULT_KEYSTORE,
+                                  MIGRATION_RESULT_SIZE);
+        return false;
+      }
+    }
+  }
+  if (!old_keystore_keys_.empty()) {
+    // Go through and add all the old keystore keys as non default keys, so
+    // they'll be preserved in the encryption_keybag when we next write the
+    // nigori node.
+    for (std::vector<std::string>::const_iterator iter =
+             old_keystore_keys_.begin(); iter != old_keystore_keys_.end();
+         ++iter) {
+      KeyParams key_params = {"localhost", "dummy", *iter};
+      cryptographer->AddNonDefaultKey(key_params);
     }
   }
   if (new_passphrase_type == KEYSTORE_PASSPHRASE &&
@@ -1213,10 +1372,16 @@ bool SyncEncryptionHandlerImpl::AttemptToMigrateNigoriToKeystore(
           keystore_key_,
           migrated_nigori.mutable_keystore_decryptor_token())) {
     LOG(ERROR) << "Failed to extract keystore decryptor token.";
+    UMA_HISTOGRAM_ENUMERATION("Sync.AttemptNigoriMigration",
+                              FAILED_TO_EXTRACT_DECRYPTOR,
+                              MIGRATION_RESULT_SIZE);
     return false;
   }
   if (!cryptographer->GetKeys(migrated_nigori.mutable_encryption_keybag())) {
     LOG(ERROR) << "Failed to extract encryption keybag.";
+    UMA_HISTOGRAM_ENUMERATION("Sync.AttemptNigoriMigration",
+                              FAILED_TO_EXTRACT_KEYBAG,
+                              MIGRATION_RESULT_SIZE);
     return false;
   }
 
@@ -1229,9 +1394,6 @@ bool SyncEncryptionHandlerImpl::AttemptToMigrateNigoriToKeystore(
         TimeToProtoTime(custom_passphrase_time_));
   }
 
-  DVLOG(1) << "Completing nigori migration to keystore support.";
-  nigori_node->SetNigoriSpecifics(migrated_nigori);
-
   FOR_EACH_OBSERVER(
       SyncEncryptionHandler::Observer,
       observers_,
@@ -1247,6 +1409,40 @@ bool SyncEncryptionHandlerImpl::AttemptToMigrateNigoriToKeystore(
   if (new_encrypt_everything && !encrypt_everything_) {
     EnableEncryptEverythingImpl(trans->GetWrappedTrans());
     ReEncryptEverything(trans);
+  } else if (!cryptographer->CanDecryptUsingDefaultKey(
+                 old_nigori.encryption_keybag())) {
+    DVLOG(1) << "Rencrypting everything due to key rotation.";
+    ReEncryptEverything(trans);
+  }
+
+  DVLOG(1) << "Completing nigori migration to keystore support.";
+  nigori_node->SetNigoriSpecifics(migrated_nigori);
+
+  switch (new_passphrase_type) {
+    case KEYSTORE_PASSPHRASE:
+      if (old_keystore_keys_.size() > 0) {
+        UMA_HISTOGRAM_ENUMERATION("Sync.AttemptNigoriMigration",
+                                  MIGRATION_SUCCESS_KEYSTORE_DEFAULT,
+                                  MIGRATION_RESULT_SIZE);
+      } else {
+        UMA_HISTOGRAM_ENUMERATION("Sync.AttemptNigoriMigration",
+                                  MIGRATION_SUCCESS_KEYSTORE_NONDEFAULT,
+                                  MIGRATION_RESULT_SIZE);
+      }
+      break;
+    case FROZEN_IMPLICIT_PASSPHRASE:
+      UMA_HISTOGRAM_ENUMERATION("Sync.AttemptNigoriMigration",
+                                MIGRATION_SUCCESS_FROZEN_IMPLICIT,
+                                MIGRATION_RESULT_SIZE);
+      break;
+    case CUSTOM_PASSPHRASE:
+      UMA_HISTOGRAM_ENUMERATION("Sync.AttemptNigoriMigration",
+                                MIGRATION_SUCCESS_CUSTOM,
+                                MIGRATION_RESULT_SIZE);
+      break;
+    default:
+      NOTREACHED();
+      break;
   }
   return true;
 }
@@ -1307,6 +1503,16 @@ bool SyncEncryptionHandlerImpl::DecryptPendingKeysWithKeystoreKey(
   if (keystore_decryptor_token.blob().empty())
     return false;
   Cryptographer temp_cryptographer(cryptographer->encryptor());
+
+  // First, go through and all all the old keystore keys to the temporary
+  // cryptographer.
+  for (size_t i = 0; i < old_keystore_keys_.size(); ++i) {
+    KeyParams old_key_params = {"localhost", "dummy", old_keystore_keys_[i]};
+    temp_cryptographer.AddKey(old_key_params);
+  }
+
+  // Then add the current keystore key as the default key and see if we can
+  // decrypt.
   KeyParams keystore_params = {"localhost", "dummy", keystore_key_};
   if (temp_cryptographer.AddKey(keystore_params) &&
       temp_cryptographer.CanDecrypt(keystore_decryptor_token)) {
@@ -1317,16 +1523,32 @@ bool SyncEncryptionHandlerImpl::DecryptPendingKeysWithKeystoreKey(
     // The keystore decryptor token is a keystore key encrypted blob containing
     // the current serialized default encryption key (and as such should be
     // able to decrypt the nigori node's encryption keybag).
+    // Note: it's possible a key rotation has happened since the migration, and
+    // we're decrypting using an old keystore key. In that case we need to
+    // ensure we re-encrypt using the newest key.
     DVLOG(1) << "Attempting to decrypt pending keys using "
              << "keystore decryptor token.";
     std::string serialized_nigori =
         temp_cryptographer.DecryptToString(keystore_decryptor_token);
+
     // This will decrypt the pending keys and add them if possible. The key
     // within |serialized_nigori| will be the default after.
     cryptographer->ImportNigoriKey(serialized_nigori);
-    // Theoretically the encryption keybag should already contain the keystore
-    // key. We explicitly add it as a safety measure.
-    cryptographer->AddNonDefaultKey(keystore_params);
+
+    if (!temp_cryptographer.CanDecryptUsingDefaultKey(
+            keystore_decryptor_token)) {
+      // The keystore decryptor token was derived from an old keystore key.
+      // A key rotation is necessary, so set the current keystore key as the
+      // default key (which will trigger a re-migration).
+      DVLOG(1) << "Pending keys based on old keystore key. Setting newest "
+               << "keystore key as default.";
+      cryptographer->AddKey(keystore_params);
+    } else {
+      // Theoretically the encryption keybag should already contain the keystore
+      // key. We explicitly add it as a safety measure.
+      DVLOG(1) << "Pending keys based on newest keystore key.";
+      cryptographer->AddNonDefaultKey(keystore_params);
+    }
     if (cryptographer->is_ready()) {
       std::string bootstrap_token;
       cryptographer->GetBootstrapToken(&bootstrap_token);
@@ -1334,6 +1556,10 @@ bool SyncEncryptionHandlerImpl::DecryptPendingKeysWithKeystoreKey(
       FOR_EACH_OBSERVER(
           SyncEncryptionHandler::Observer,
           observers_,
+          OnPassphraseAccepted());
+      FOR_EACH_OBSERVER(
+          SyncEncryptionHandler::Observer,
+          observers_,
           OnBootstrapTokenUpdated(bootstrap_token,
                                   PASSPHRASE_BOOTSTRAP_TOKEN));
       FOR_EACH_OBSERVER(
diff --git a/sync/internal_api/sync_encryption_handler_impl.h b/sync/internal_api/sync_encryption_handler_impl.h
index d93db15..a308503 100644
--- a/sync/internal_api/sync_encryption_handler_impl.h
+++ b/sync/internal_api/sync_encryption_handler_impl.h
@@ -73,8 +73,8 @@ class SyncEncryptionHandlerImpl
       syncable::BaseTransaction* const trans) const OVERRIDE;
   virtual bool NeedKeystoreKey(
       syncable::BaseTransaction* const trans) const OVERRIDE;
-  virtual bool SetKeystoreKey(
-      const std::string& key,
+  virtual bool SetKeystoreKeys(
+      const google::protobuf::RepeatedPtrField<google::protobuf::string>& keys,
       syncable::BaseTransaction* const trans) OVERRIDE;
   // Can be called from any thread.
   virtual ModelTypeSet GetEncryptedTypes(
@@ -90,6 +90,7 @@ class SyncEncryptionHandlerImpl
   base::Time custom_passphrase_time() const;
 
  private:
+  friend class SyncEncryptionHandlerImplTest;
   FRIEND_TEST_ALL_PREFIXES(SyncEncryptionHandlerImplTest,
                            NigoriEncryptionTypes);
   FRIEND_TEST_ALL_PREFIXES(SyncEncryptionHandlerImplTest,
@@ -216,8 +217,8 @@ class SyncEncryptionHandlerImpl
   // triggered or not.
   // Conditions for triggering migration:
   // 1. Cryptographer has no pending keys
-  // 2. Nigori node isn't already properly migrated.
-  // 3. Keystore key is available (if we are not migrated yet).
+  // 2. Nigori node isn't already properly migrated or we need to rotate keys.
+  // 3. Keystore key is available.
   // Note: if the nigori node is migrated but has an invalid state, will return
   // true (e.g. node has KEYSTORE_PASSPHRASE, local is CUSTOM_PASSPHRASE).
   bool ShouldTriggerMigration(const sync_pb::NigoriSpecifics& nigori,
@@ -283,9 +284,15 @@ class SyncEncryptionHandlerImpl
   // keys stored in the nigori node.
   PassphraseType passphrase_type_;
 
-  // The keystore key provided by the server.
+  // The current keystore key provided by the server.
   std::string keystore_key_;
 
+  // The set of old keystore keys. Every time a key rotation occurs, the server
+  // sends down all previous keystore keys as well as the new key. We preserve
+  // the old keys so that when we re-encrypt we can ensure they're all added to
+  // the keybag (and to detect that a key rotation has occurred).
+  std::vector<std::string> old_keystore_keys_;
+
   // The number of times we've automatically (i.e. not via SetPassphrase or
   // conflict resolver) updated the nigori's encryption keys in this chrome
   // instantiation.
diff --git a/sync/internal_api/sync_encryption_handler_impl_unittest.cc b/sync/internal_api/sync_encryption_handler_impl_unittest.cc
index 65fc964..74bf0e42 100644
--- a/sync/internal_api/sync_encryption_handler_impl_unittest.cc
+++ b/sync/internal_api/sync_encryption_handler_impl_unittest.cc
@@ -7,6 +7,7 @@
 #include <string>
 
 #include "base/base64.h"
+#include "base/json/json_string_value_serializer.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/message_loop.h"
 #include "base/tracked_objects.h"
@@ -33,6 +34,7 @@ namespace {
 
 using ::testing::_;
 using ::testing::AnyNumber;
+using ::testing::AtLeast;
 using ::testing::Mock;
 using ::testing::SaveArg;
 using ::testing::StrictMock;
@@ -59,6 +61,13 @@ class SyncEncryptionHandlerObserverMock
                                              base::Time));  // NOLINT
 };
 
+google::protobuf::RepeatedPtrField<google::protobuf::string>
+BuildEncryptionKeyProto(std::string encryption_key) {
+  google::protobuf::RepeatedPtrField<google::protobuf::string> keys;
+  keys.Add()->assign(encryption_key);
+  return keys;
+}
+
 }  // namespace
 
 class SyncEncryptionHandlerImplTest : public ::testing::Test {
@@ -177,6 +186,149 @@ class SyncEncryptionHandlerImplTest : public ::testing::Test {
         nigori.encryption_keybag()));
   }
 
+  sync_pb::NigoriSpecifics BuildMigratedNigori(
+      PassphraseType passphrase_type,
+      int64 migration_time,
+      const std::string& default_passphrase,
+      const std::string& keystore_key) {
+    DCHECK_NE(passphrase_type, IMPLICIT_PASSPHRASE);
+    Cryptographer other_cryptographer(GetCryptographer()->encryptor());
+
+    std::string default_key = default_passphrase;
+    if (default_key.empty()) {
+      default_key = keystore_key;
+    } else {
+      KeyParams keystore_params = {"localhost", "dummy", keystore_key};
+      other_cryptographer.AddKey(keystore_params);
+    }
+    KeyParams params = {"localhost", "dummy", default_key};
+    other_cryptographer.AddKey(params);
+    EXPECT_TRUE(other_cryptographer.is_ready());
+
+    sync_pb::NigoriSpecifics nigori;
+    other_cryptographer.GetKeys(nigori.mutable_encryption_keybag());
+    nigori.set_keybag_is_frozen(true);
+    nigori.set_keystore_migration_time(migration_time);
+
+    if (passphrase_type == KEYSTORE_PASSPHRASE) {
+      sync_pb::EncryptedData keystore_decryptor_token;
+      EXPECT_TRUE(encryption_handler()->GetKeystoreDecryptor(
+          other_cryptographer,
+          keystore_key,
+          &keystore_decryptor_token));
+      nigori.mutable_keystore_decryptor_token()->CopyFrom(
+          keystore_decryptor_token);
+      nigori.set_passphrase_type(sync_pb::NigoriSpecifics::KEYSTORE_PASSPHRASE);
+    } else {
+      nigori.set_encrypt_everything(true);
+      nigori.set_passphrase_type(
+          passphrase_type == CUSTOM_PASSPHRASE ?
+              sync_pb::NigoriSpecifics::CUSTOM_PASSPHRASE :
+              sync_pb::NigoriSpecifics::FROZEN_IMPLICIT_PASSPHRASE);
+    }
+    return nigori;
+  }
+
+  // Build a migrated nigori node with the specified default passphrase
+  // and keystore key and initialize the encryption handler with it.
+  void InitKeystoreMigratedNigori(int64 migration_time,
+                                  const std::string& default_passphrase,
+                                  const std::string& keystore_key) {
+    {
+      WriteTransaction trans(FROM_HERE, user_share());
+      WriteNode nigori_node(&trans);
+      ASSERT_EQ(nigori_node.InitByTagLookup(kNigoriTag), BaseNode::INIT_OK);
+      sync_pb::NigoriSpecifics nigori = BuildMigratedNigori(
+          KEYSTORE_PASSPHRASE,
+          migration_time,
+          default_passphrase,
+          keystore_key);
+      nigori_node.SetNigoriSpecifics(nigori);
+    }
+
+    EXPECT_CALL(*observer(),
+                OnPassphraseTypeChanged(KEYSTORE_PASSPHRASE, _));
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, PASSPHRASE_BOOTSTRAP_TOKEN));
+    EXPECT_CALL(*observer(),
+                OnCryptographerStateChanged(_)).Times(AtLeast(1));
+    EXPECT_CALL(*observer(),
+                OnEncryptedTypesChanged(_, false));
+    EXPECT_CALL(*observer(),
+                OnEncryptionComplete()).Times(AtLeast(1));
+    encryption_handler()->Init();
+    EXPECT_TRUE(encryption_handler()->MigratedToKeystore());
+    EXPECT_EQ(encryption_handler()->GetPassphraseType(), KEYSTORE_PASSPHRASE);
+    EXPECT_FALSE(encryption_handler()->EncryptEverythingEnabled());
+    Mock::VerifyAndClearExpectations(observer());
+  }
+
+  // Build a migrated nigori node with the specified default passphrase
+  // as a custom passphrase.
+  void InitCustomPassMigratedNigori(int64 migration_time,
+                                    const std::string& default_passphrase) {
+    {
+      WriteTransaction trans(FROM_HERE, user_share());
+      WriteNode nigori_node(&trans);
+      ASSERT_EQ(nigori_node.InitByTagLookup(kNigoriTag), BaseNode::INIT_OK);
+      sync_pb::NigoriSpecifics nigori = BuildMigratedNigori(
+          CUSTOM_PASSPHRASE,
+          migration_time,
+          default_passphrase,
+          kKeystoreKey);
+      nigori_node.SetNigoriSpecifics(nigori);
+    }
+
+    EXPECT_CALL(*observer(),
+                OnPassphraseTypeChanged(CUSTOM_PASSPHRASE, _));
+    EXPECT_CALL(*observer(),
+                OnCryptographerStateChanged(_)).Times(AtLeast(1));
+    EXPECT_CALL(*observer(),
+                OnEncryptedTypesChanged(_, true)).Times(AtLeast(1));
+    EXPECT_CALL(*observer(),
+                OnEncryptionComplete()).Times(AtLeast(1));
+    encryption_handler()->Init();
+    EXPECT_TRUE(encryption_handler()->MigratedToKeystore());
+    EXPECT_EQ(encryption_handler()->GetPassphraseType(), CUSTOM_PASSPHRASE);
+    EXPECT_TRUE(encryption_handler()->EncryptEverythingEnabled());
+    Mock::VerifyAndClearExpectations(observer());
+  }
+
+  // Build an unmigrated nigori node with the specified passphrase and type and
+  // initialize the encryption handler with it.
+  void InitUnmigratedNigori(const std::string& default_passphrase,
+                            PassphraseType passphrase_type) {
+    DCHECK_NE(passphrase_type, FROZEN_IMPLICIT_PASSPHRASE);
+    Cryptographer other_cryptographer(GetCryptographer()->encryptor());
+    KeyParams default_key = {"localhost", "dummy", default_passphrase};
+    other_cryptographer.AddKey(default_key);
+    EXPECT_TRUE(other_cryptographer.is_ready());
+
+    {
+      WriteTransaction trans(FROM_HERE, user_share());
+      WriteNode nigori_node(&trans);
+      ASSERT_EQ(nigori_node.InitByTagLookup(kNigoriTag), BaseNode::INIT_OK);
+      sync_pb::NigoriSpecifics nigori;
+      other_cryptographer.GetKeys(nigori.mutable_encryption_keybag());
+      nigori.set_keybag_is_frozen(passphrase_type == CUSTOM_PASSPHRASE);
+      nigori_node.SetNigoriSpecifics(nigori);
+    }
+
+    if (passphrase_type != IMPLICIT_PASSPHRASE) {
+      EXPECT_CALL(*observer(),
+                  OnPassphraseTypeChanged(passphrase_type, _));
+    }
+    EXPECT_CALL(*observer(),
+                OnCryptographerStateChanged(_)).Times(AtLeast(1));
+    EXPECT_CALL(*observer(),
+                OnEncryptedTypesChanged(_, false));
+    encryption_handler()->Init();
+    EXPECT_FALSE(encryption_handler()->MigratedToKeystore());
+    EXPECT_EQ(encryption_handler()->GetPassphraseType(), passphrase_type);
+    EXPECT_FALSE(encryption_handler()->EncryptEverythingEnabled());
+    Mock::VerifyAndClearExpectations(observer());
+  }
+
  protected:
   TestUserShare test_user_share_;
   FakeEncryptor encryptor_;
@@ -450,27 +602,73 @@ TEST_F(SyncEncryptionHandlerImplTest, ReceiveOldNigori) {
 }
 
 // Ensure setting the keystore key works, updates the bootstrap token, and
-// doesn't modify the cryptographer.
-TEST_F(SyncEncryptionHandlerImplTest, SetKeystoreUpdatedBoostrapToken) {
+// doesn't modify the cryptographer. Then verify that the bootstrap token
+// can be correctly parsed by the encryption handler at startup time.
+TEST_F(SyncEncryptionHandlerImplTest, SetKeystoreUpdatesBootstrap) {
   WriteTransaction trans(FROM_HERE, user_share());
+
+  // Passing no keys should do nothing.
   EXPECT_CALL(*observer(), OnBootstrapTokenUpdated(_, _)).Times(0);
   EXPECT_FALSE(GetCryptographer()->is_initialized());
   EXPECT_TRUE(encryption_handler()->NeedKeystoreKey(trans.GetWrappedTrans()));
-  EXPECT_FALSE(encryption_handler()->SetKeystoreKey("",
-                                                    trans.GetWrappedTrans()));
+  EXPECT_FALSE(
+      encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(""),
+                                            trans.GetWrappedTrans()));
   EXPECT_TRUE(encryption_handler()->NeedKeystoreKey(trans.GetWrappedTrans()));
   Mock::VerifyAndClearExpectations(observer());
 
+  // Build a set of keystore keys.
+  const char kRawOldKeystoreKey[] = "old_keystore_key";
+  std::string old_keystore_key;
+  base::Base64Encode(kRawOldKeystoreKey, &old_keystore_key);
+  google::protobuf::RepeatedPtrField<google::protobuf::string> keys;
+  keys.Add()->assign(kRawOldKeystoreKey);
+  keys.Add()->assign(kRawKeystoreKey);
+
+  // Pass them to the encryption handler, triggering a bootstrap token update.
   std::string encoded_key;
-  ASSERT_TRUE(base::Base64Encode(kRawKeystoreKey, &encoded_key));
-  ASSERT_EQ(kKeystoreKey, encoded_key);
-  EXPECT_CALL(*observer(),
-              OnBootstrapTokenUpdated(kKeystoreKey,
-                                      KEYSTORE_BOOTSTRAP_TOKEN));
-  EXPECT_TRUE(encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                                   trans.GetWrappedTrans()));
+  std::string keystore_bootstrap;
+  EXPECT_CALL(*observer(),
+              OnBootstrapTokenUpdated(_,
+                                      KEYSTORE_BOOTSTRAP_TOKEN)).
+      WillOnce(SaveArg<0>(&keystore_bootstrap));
+  EXPECT_TRUE(
+      encryption_handler()->SetKeystoreKeys(
+          keys,
+          trans.GetWrappedTrans()));
   EXPECT_FALSE(encryption_handler()->NeedKeystoreKey(trans.GetWrappedTrans()));
   EXPECT_FALSE(GetCryptographer()->is_initialized());
+
+  // Ensure the bootstrap is encoded properly (a base64 encoded encrypted blob
+  // of list values containing the keystore keys).
+  std::string decoded_bootstrap;
+  ASSERT_TRUE(base::Base64Decode(keystore_bootstrap, &decoded_bootstrap));
+  std::string decrypted_bootstrap;
+  ASSERT_TRUE(
+      GetCryptographer()->encryptor()->DecryptString(decoded_bootstrap,
+                                                     &decrypted_bootstrap));
+  JSONStringValueSerializer json(decrypted_bootstrap);
+  scoped_ptr<base::Value> deserialized_keystore_keys(
+      json.Deserialize(NULL, NULL));
+  ASSERT_TRUE(deserialized_keystore_keys.get());
+  base::ListValue* keystore_list = NULL;
+  deserialized_keystore_keys->GetAsList(&keystore_list);
+  ASSERT_TRUE(keystore_list);
+  ASSERT_EQ(2U, keystore_list->GetSize());
+  std::string test_string;
+  keystore_list->GetString(0, &test_string);
+  ASSERT_EQ(old_keystore_key, test_string);
+  keystore_list->GetString(1, &test_string);
+  ASSERT_EQ(kKeystoreKey, test_string);
+
+
+  // Now make sure a new encryption handler can correctly parse the bootstrap
+  // token.
+  SyncEncryptionHandlerImpl handler2(user_share(),
+                                     &encryptor_,
+                                     "",  // Cryptographer bootstrap.
+                                     keystore_bootstrap);
+  EXPECT_FALSE(handler2.NeedKeystoreKey(trans.GetWrappedTrans()));
 }
 
 // Ensure GetKeystoreDecryptor only updates the keystore decryptor token if it
@@ -504,8 +702,9 @@ TEST_F(SyncEncryptionHandlerImplTest, MigrateOnDecryptImplicitPass) {
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
     ReadTransaction trans(FROM_HERE, user_share());
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
     Mock::VerifyAndClearExpectations(observer());
   }
   EXPECT_FALSE(encryption_handler()->MigratedToKeystore());
@@ -560,8 +759,9 @@ TEST_F(SyncEncryptionHandlerImplTest, MigrateOnDecryptCustomPass) {
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
     ReadTransaction trans(FROM_HERE, user_share());
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
     Mock::VerifyAndClearExpectations(observer());
   }
   EXPECT_FALSE(encryption_handler()->MigratedToKeystore());
@@ -632,8 +832,9 @@ TEST_F(SyncEncryptionHandlerImplTest, MigrateOnKeystoreKeyAvailableImplicit) {
                 OnCryptographerStateChanged(_)).Times(AnyNumber());
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   EXPECT_CALL(*observer(),
               OnPassphraseTypeChanged(KEYSTORE_PASSPHRASE, _));
@@ -676,8 +877,9 @@ TEST_F(SyncEncryptionHandlerImplTest,
                 OnCryptographerStateChanged(_)).Times(AnyNumber());
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   EXPECT_CALL(*observer(),
               OnPassphraseTypeChanged(FROZEN_IMPLICIT_PASSPHRASE, _));
@@ -729,8 +931,9 @@ TEST_F(SyncEncryptionHandlerImplTest,
                 OnCryptographerStateChanged(_)).Times(AnyNumber());
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   // The actual migration gets posted, so run all pending tasks.
   PumpLoop();
@@ -773,8 +976,9 @@ TEST_F(SyncEncryptionHandlerImplTest,
                 OnCryptographerStateChanged(_)).Times(AnyNumber());
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   EXPECT_CALL(*observer(),
               OnEncryptedTypesChanged(_, true));
@@ -822,6 +1026,7 @@ TEST_F(SyncEncryptionHandlerImplTest, ReceiveMigratedNigoriKeystorePass) {
     nigori.set_keystore_migration_time(1);
     nigori.set_passphrase_type(sync_pb::NigoriSpecifics::KEYSTORE_PASSPHRASE);
 
+    EXPECT_CALL(*observer(), OnPassphraseAccepted());
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, PASSPHRASE_BOOTSTRAP_TOKEN));
     EXPECT_CALL(*observer(),
@@ -830,8 +1035,9 @@ TEST_F(SyncEncryptionHandlerImplTest, ReceiveMigratedNigoriKeystorePass) {
                 OnPassphraseTypeChanged(KEYSTORE_PASSPHRASE, _));
     EXPECT_CALL(*observer(),
                 OnCryptographerStateChanged(_)).Times(AnyNumber());
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
     encryption_handler()->ApplyNigoriUpdate(nigori, trans.GetWrappedTrans());
     nigori_node.SetNigoriSpecifics(nigori);
   }
@@ -874,8 +1080,9 @@ TEST_F(SyncEncryptionHandlerImplTest, ReceiveMigratedNigoriFrozenImplicitPass) {
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
     ReadTransaction trans(FROM_HERE, user_share());
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   EXPECT_FALSE(encryption_handler()->MigratedToKeystore());
 
@@ -953,8 +1160,9 @@ TEST_F(SyncEncryptionHandlerImplTest, ReceiveMigratedNigoriCustomPass) {
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
     ReadTransaction trans(FROM_HERE, user_share());
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   EXPECT_FALSE(encryption_handler()->MigratedToKeystore());
 
@@ -1060,8 +1268,9 @@ TEST_F(SyncEncryptionHandlerImplTest, ReceiveUnmigratedNigoriAfterMigration) {
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
     ReadTransaction trans(FROM_HERE, user_share());
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   Mock::VerifyAndClearExpectations(observer());
 
@@ -1069,6 +1278,7 @@ TEST_F(SyncEncryptionHandlerImplTest, ReceiveUnmigratedNigoriAfterMigration) {
   // properly overwrite it with the migrated + encrypt everything state.
   EXPECT_CALL(*observer(),
               OnCryptographerStateChanged(_)).Times(AnyNumber());
+  EXPECT_CALL(*observer(), OnEncryptionComplete());
   {
     Cryptographer other_cryptographer(GetCryptographer()->encryptor());
     other_cryptographer.AddKey(old_key);
@@ -1138,8 +1348,9 @@ TEST_F(SyncEncryptionHandlerImplTest, ReceiveOldMigratedNigori) {
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
     ReadTransaction trans(FROM_HERE, user_share());
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   Mock::VerifyAndClearExpectations(observer());
 
@@ -1147,6 +1358,7 @@ TEST_F(SyncEncryptionHandlerImplTest, ReceiveOldMigratedNigori) {
   // properly overwrite it with the migrated + encrypt everything state.
   EXPECT_CALL(*observer(),
               OnCryptographerStateChanged(_)).Times(AnyNumber());
+  EXPECT_CALL(*observer(), OnEncryptionComplete());
   {
     WriteTransaction trans(FROM_HERE, user_share());
     WriteNode nigori_node(&trans);
@@ -1225,6 +1437,7 @@ TEST_F(SyncEncryptionHandlerImplTest, SetKeystoreAfterReceivingMigratedNigori) {
   EXPECT_FALSE(encryption_handler()->EncryptEverythingEnabled());
   Mock::VerifyAndClearExpectations(observer());
 
+  EXPECT_CALL(*observer(), OnPassphraseAccepted());
   EXPECT_CALL(*observer(),
               OnCryptographerStateChanged(_)).Times(AnyNumber());
   EXPECT_CALL(*observer(),
@@ -1233,8 +1446,9 @@ TEST_F(SyncEncryptionHandlerImplTest, SetKeystoreAfterReceivingMigratedNigori) {
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
     ReadTransaction trans(FROM_HERE, user_share());
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   PumpLoop();
   EXPECT_TRUE(encryption_handler()->MigratedToKeystore());
@@ -1289,10 +1503,12 @@ TEST_F(SyncEncryptionHandlerImplTest, SetCustomPassAfterMigration) {
     nigori_node.SetNigoriSpecifics(nigori);
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
 
+  EXPECT_CALL(*observer(), OnPassphraseAccepted());
   EXPECT_CALL(*observer(),
               OnPassphraseTypeChanged(KEYSTORE_PASSPHRASE, _));
   EXPECT_CALL(*observer(),
@@ -1680,11 +1896,12 @@ TEST_F(SyncEncryptionHandlerImplTest,
     EXPECT_CALL(*observer(),
                 OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
     ReadTransaction trans(FROM_HERE, user_share());
-    encryption_handler()->SetKeystoreKey(kRawKeystoreKey,
-                                         trans.GetWrappedTrans());
+    encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                              kRawKeystoreKey),
+                                          trans.GetWrappedTrans());
   }
   EXPECT_CALL(*observer(),
-            OnPassphraseTypeChanged(KEYSTORE_PASSPHRASE, _));
+              OnPassphraseTypeChanged(KEYSTORE_PASSPHRASE, _));
   PumpLoop();
   Mock::VerifyAndClearExpectations(observer());
   EXPECT_TRUE(encryption_handler()->MigratedToKeystore());
@@ -1694,6 +1911,7 @@ TEST_F(SyncEncryptionHandlerImplTest,
   // Now build an old keystore passphrase nigori node.
   EXPECT_CALL(*observer(),
               OnCryptographerStateChanged(_)).Times(AnyNumber());
+  EXPECT_CALL(*observer(), OnEncryptionComplete());
   {
     WriteTransaction trans(FROM_HERE, user_share());
     WriteNode nigori_node(&trans);
@@ -1723,4 +1941,312 @@ TEST_F(SyncEncryptionHandlerImplTest,
   VerifyMigratedNigori(KEYSTORE_PASSPHRASE, kCurKey);
 }
 
+// Trigger a key rotation upon receiving new keys if we already had a keystore
+// migrated nigori with the gaia key as the default (still in backwards
+// compatible mode).
+TEST_F(SyncEncryptionHandlerImplTest, RotateKeysGaiaDefault) {
+  const char kOldGaiaKey[] = "old_gaia_key";
+  const char kRawOldKeystoreKey[] = "old_keystore_key";
+  std::string old_keystore_key;
+  base::Base64Encode(kRawOldKeystoreKey, &old_keystore_key);
+  {
+    ReadTransaction trans(FROM_HERE, user_share());
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
+      encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                                kRawOldKeystoreKey),
+                                            trans.GetWrappedTrans());
+  }
+  PumpLoop();
+  Mock::VerifyAndClearExpectations(observer());
+
+  EXPECT_CALL(*observer(), OnPassphraseAccepted());
+  InitKeystoreMigratedNigori(1, kOldGaiaKey, old_keystore_key);
+
+  // Now set some new keystore keys.
+  EXPECT_CALL(*observer(), OnCryptographerStateChanged(_)).Times(AnyNumber());
+  EXPECT_CALL(*observer(), OnEncryptionComplete());
+  {
+    google::protobuf::RepeatedPtrField<google::protobuf::string> keys;
+    keys.Add()->assign(kRawOldKeystoreKey);
+    keys.Add()->assign(kRawKeystoreKey);
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
+    ReadTransaction trans(FROM_HERE, user_share());
+    encryption_handler()->SetKeystoreKeys(keys,
+                                          trans.GetWrappedTrans());
+  }
+  // Pump for any posted tasks.
+  PumpLoop();
+  Mock::VerifyAndClearExpectations(observer());
+
+  // Verify we're still migrated and have proper encryption state. We should
+  // have rotated the keybag so that it's now encrypted with the newest keystore
+  // key (instead of the old gaia key).
+  EXPECT_TRUE(encryption_handler()->MigratedToKeystore());
+  EXPECT_TRUE(GetCryptographer()->is_ready());
+  EXPECT_EQ(encryption_handler()->GetPassphraseType(), KEYSTORE_PASSPHRASE);
+  EXPECT_FALSE(encryption_handler()->EncryptEverythingEnabled());
+  VerifyMigratedNigori(KEYSTORE_PASSPHRASE, kKeystoreKey);
+}
+
+// Trigger a key rotation upon receiving new keys if we already had a keystore
+// migrated nigori with the keystore key as the default.
+TEST_F(SyncEncryptionHandlerImplTest, RotateKeysKeystoreDefault) {
+  const char kRawOldKeystoreKey[] = "old_keystore_key";
+  std::string old_keystore_key;
+  base::Base64Encode(kRawOldKeystoreKey, &old_keystore_key);
+  {
+    ReadTransaction trans(FROM_HERE, user_share());
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
+      encryption_handler()->SetKeystoreKeys(BuildEncryptionKeyProto(
+                                                kRawOldKeystoreKey),
+                                            trans.GetWrappedTrans());
+  }
+  PumpLoop();
+  Mock::VerifyAndClearExpectations(observer());
+
+  EXPECT_CALL(*observer(), OnPassphraseAccepted());
+  InitKeystoreMigratedNigori(1, old_keystore_key, old_keystore_key);
+
+  // Now set some new keystore keys.
+  EXPECT_CALL(*observer(), OnCryptographerStateChanged(_)).Times(AnyNumber());
+  EXPECT_CALL(*observer(), OnEncryptionComplete());
+  {
+    google::protobuf::RepeatedPtrField<google::protobuf::string> keys;
+    keys.Add()->assign(kRawOldKeystoreKey);
+    keys.Add()->assign(kRawKeystoreKey);
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
+    ReadTransaction trans(FROM_HERE, user_share());
+    encryption_handler()->SetKeystoreKeys(keys,
+                                          trans.GetWrappedTrans());
+  }
+  // Pump for any posted tasks.
+  PumpLoop();
+  Mock::VerifyAndClearExpectations(observer());
+
+  // Verify we're still migrated and have proper encryption state. We should
+  // have rotated the keybag so that it's now encrypted with the newest keystore
+  // key (instead of the old gaia key).
+  EXPECT_TRUE(encryption_handler()->MigratedToKeystore());
+  EXPECT_TRUE(GetCryptographer()->is_ready());
+  EXPECT_EQ(encryption_handler()->GetPassphraseType(), KEYSTORE_PASSPHRASE);
+  EXPECT_FALSE(encryption_handler()->EncryptEverythingEnabled());
+  VerifyMigratedNigori(KEYSTORE_PASSPHRASE, kKeystoreKey);
+}
+
+// Trigger a key rotation upon when a pending gaia passphrase is resolved.
+TEST_F(SyncEncryptionHandlerImplTest, RotateKeysAfterPendingGaiaResolved) {
+  const char kOldGaiaKey[] = "old_gaia_key";
+  const char kRawOldKeystoreKey[] = "old_keystore_key";
+
+  EXPECT_CALL(*observer(), OnPassphraseRequired(_, _));
+  InitUnmigratedNigori(kOldGaiaKey, IMPLICIT_PASSPHRASE);
+
+  {
+    // Pass multiple keystore keys, signaling a rotation has happened.
+    google::protobuf::RepeatedPtrField<google::protobuf::string> keys;
+    keys.Add()->assign(kRawOldKeystoreKey);
+    keys.Add()->assign(kRawKeystoreKey);
+    ReadTransaction trans(FROM_HERE, user_share());
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
+    encryption_handler()->SetKeystoreKeys(keys,
+                                          trans.GetWrappedTrans());
+  }
+  PumpLoop();
+  Mock::VerifyAndClearExpectations(observer());
+
+  // Resolve the pending keys. This should trigger the key rotation.
+  EXPECT_CALL(*observer(),
+              OnCryptographerStateChanged(_)).Times(AnyNumber());
+  EXPECT_CALL(*observer(),
+              OnPassphraseAccepted());
+  EXPECT_CALL(*observer(),
+              OnPassphraseTypeChanged(KEYSTORE_PASSPHRASE, _));
+  EXPECT_CALL(*observer(),
+              OnBootstrapTokenUpdated(_, PASSPHRASE_BOOTSTRAP_TOKEN));
+  EXPECT_CALL(*observer(),
+              OnEncryptionComplete()).Times(AtLeast(1));
+  EXPECT_FALSE(encryption_handler()->MigratedToKeystore());
+  encryption_handler()->SetDecryptionPassphrase(kOldGaiaKey);
+  EXPECT_TRUE(encryption_handler()->MigratedToKeystore());
+  EXPECT_EQ(KEYSTORE_PASSPHRASE, encryption_handler()->GetPassphraseType());
+  VerifyMigratedNigori(KEYSTORE_PASSPHRASE, kKeystoreKey);
+}
+
+// When signing in for the first time, make sure we can rotate keys if we
+// already have a keystore migrated nigori.
+TEST_F(SyncEncryptionHandlerImplTest, RotateKeysGaiaDefaultOnInit) {
+  // Destroy the existing nigori node so we init without a nigori node.
+  TearDown();
+  test_user_share_.SetUp();
+  SetUpEncryption();
+
+  const char kOldGaiaKey[] = "old_gaia_key";
+  const char kRawOldKeystoreKey[] = "old_keystore_key";
+  std::string old_keystore_key;
+  base::Base64Encode(kRawOldKeystoreKey, &old_keystore_key);
+
+  // Set two keys, signaling that a rotation has been performed. No nigori
+  // node is present yet, so we can't rotate.
+  {
+    google::protobuf::RepeatedPtrField<google::protobuf::string> keys;
+    keys.Add()->assign(kRawOldKeystoreKey);
+    keys.Add()->assign(kRawKeystoreKey);
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
+    ReadTransaction trans(FROM_HERE, user_share());
+    encryption_handler()->SetKeystoreKeys(keys,
+                                          trans.GetWrappedTrans());
+  }
+
+  // Then init the nigori node with an old set of keys.
+  CreateRootForType(NIGORI);
+  EXPECT_CALL(*observer(), OnPassphraseAccepted());
+  InitKeystoreMigratedNigori(1, kOldGaiaKey, old_keystore_key);
+  PumpLoop();
+  Mock::VerifyAndClearExpectations(observer());
+
+  // Verify we're still migrated and have proper encryption state. We should
+  // have rotated the keybag so that it's now encrypted with the newest keystore
+  // key (instead of the old gaia key).
+  EXPECT_TRUE(encryption_handler()->MigratedToKeystore());
+  EXPECT_TRUE(GetCryptographer()->is_ready());
+  EXPECT_EQ(encryption_handler()->GetPassphraseType(), KEYSTORE_PASSPHRASE);
+  EXPECT_FALSE(encryption_handler()->EncryptEverythingEnabled());
+  VerifyMigratedNigori(KEYSTORE_PASSPHRASE, kKeystoreKey);
+}
+
+// Trigger a key rotation when a migrated nigori (with an old keystore key) is
+// applied.
+TEST_F(SyncEncryptionHandlerImplTest, RotateKeysWhenMigratedNigoriArrives) {
+  const char kOldGaiaKey[] = "old_gaia_key";
+  const char kRawOldKeystoreKey[] = "old_keystore_key";
+  std::string old_keystore_key;
+  base::Base64Encode(kRawOldKeystoreKey, &old_keystore_key);
+
+  EXPECT_CALL(*observer(), OnPassphraseRequired(_, _));
+  InitUnmigratedNigori(kOldGaiaKey, IMPLICIT_PASSPHRASE);
+
+  {
+    // Pass multiple keystore keys, signaling a rotation has happened.
+    google::protobuf::RepeatedPtrField<google::protobuf::string> keys;
+    keys.Add()->assign(kRawOldKeystoreKey);
+    keys.Add()->assign(kRawKeystoreKey);
+    ReadTransaction trans(FROM_HERE, user_share());
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
+    encryption_handler()->SetKeystoreKeys(keys,
+                                          trans.GetWrappedTrans());
+  }
+  PumpLoop();
+  Mock::VerifyAndClearExpectations(observer());
+
+  // Now simulate downloading a nigori node that was migrated before the
+  // keys were rotated, and hence still encrypt with the old gaia key.
+  EXPECT_CALL(*observer(),
+              OnCryptographerStateChanged(_)).Times(AnyNumber());
+  EXPECT_CALL(*observer(),
+              OnPassphraseAccepted());
+  EXPECT_CALL(*observer(),
+              OnPassphraseTypeChanged(KEYSTORE_PASSPHRASE, _));
+  EXPECT_CALL(*observer(),
+              OnBootstrapTokenUpdated(_, PASSPHRASE_BOOTSTRAP_TOKEN));
+  EXPECT_CALL(*observer(),
+              OnEncryptionComplete()).Times(AtLeast(1));
+  {
+    sync_pb::NigoriSpecifics nigori = BuildMigratedNigori(
+        KEYSTORE_PASSPHRASE,
+        1,
+        kOldGaiaKey,
+        old_keystore_key);
+    // Update the encryption handler.
+    WriteTransaction trans(FROM_HERE, user_share());
+    encryption_handler()->ApplyNigoriUpdate(
+        nigori,
+        trans.GetWrappedTrans());
+  }
+  EXPECT_FALSE(encryption_handler()->MigratedToKeystore());
+  PumpLoop();
+
+  EXPECT_TRUE(encryption_handler()->MigratedToKeystore());
+  EXPECT_EQ(KEYSTORE_PASSPHRASE, encryption_handler()->GetPassphraseType());
+  VerifyMigratedNigori(KEYSTORE_PASSPHRASE, kKeystoreKey);
+}
+
+// Verify that performing a migration while having more than one keystore key
+// preserves a custom passphrase.
+TEST_F(SyncEncryptionHandlerImplTest, RotateKeysUnmigratedCustomPassphrase) {
+  const char kCustomPass[] = "custom_passphrase";
+  const char kRawOldKeystoreKey[] = "old_keystore_key";
+
+  EXPECT_CALL(*observer(), OnPassphraseRequired(_, _));
+  InitUnmigratedNigori(kCustomPass, CUSTOM_PASSPHRASE);
+
+  {
+    // Pass multiple keystore keys, signaling a rotation has happened.
+    google::protobuf::RepeatedPtrField<google::protobuf::string> keys;
+    keys.Add()->assign(kRawOldKeystoreKey);
+    keys.Add()->assign(kRawKeystoreKey);
+    ReadTransaction trans(FROM_HERE, user_share());
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
+    encryption_handler()->SetKeystoreKeys(keys,
+                                          trans.GetWrappedTrans());
+  }
+  PumpLoop();
+  Mock::VerifyAndClearExpectations(observer());
+
+  // Pass the decryption passphrase. This will also trigger the migration,
+  // but should not overwrite the default key.
+  EXPECT_CALL(*observer(),
+              OnCryptographerStateChanged(_)).Times(AnyNumber());
+  EXPECT_CALL(*observer(),
+              OnPassphraseAccepted());
+  EXPECT_CALL(*observer(),
+              OnEncryptedTypesChanged(_, true));
+  EXPECT_CALL(*observer(),
+              OnEncryptionComplete()).Times(AnyNumber());
+  EXPECT_CALL(*observer(),
+              OnBootstrapTokenUpdated(_, PASSPHRASE_BOOTSTRAP_TOKEN));
+  encryption_handler()->SetDecryptionPassphrase(kCustomPass);
+  Mock::VerifyAndClearExpectations(observer());
+
+  VerifyMigratedNigori(CUSTOM_PASSPHRASE, kCustomPass);
+}
+
+// Verify that a key rotation done after we've migrated a custom passphrase
+// nigori node preserves the custom passphrase.
+TEST_F(SyncEncryptionHandlerImplTest, RotateKeysMigratedCustomPassphrase) {
+  const char kCustomPass[] = "custom_passphrase";
+  const char kRawOldKeystoreKey[] = "old_keystore_key";
+
+  KeyParams custom_key = {"localhost", "dummy", kCustomPass};
+  GetCryptographer()->AddKey(custom_key);
+
+  InitCustomPassMigratedNigori(1, kCustomPass);
+  VerifyMigratedNigoriWithTimestamp(1, CUSTOM_PASSPHRASE, kCustomPass);
+
+  {
+    // Pass multiple keystore keys, signaling a rotation has happened.
+    google::protobuf::RepeatedPtrField<google::protobuf::string> keys;
+    keys.Add()->assign(kRawOldKeystoreKey);
+    keys.Add()->assign(kRawKeystoreKey);
+    ReadTransaction trans(FROM_HERE, user_share());
+    EXPECT_CALL(*observer(),
+                OnBootstrapTokenUpdated(_, KEYSTORE_BOOTSTRAP_TOKEN));
+    EXPECT_CALL(*observer(),
+                OnCryptographerStateChanged(_)).Times(AnyNumber());
+    encryption_handler()->SetKeystoreKeys(keys,
+                                          trans.GetWrappedTrans());
+  }
+  PumpLoop();
+  Mock::VerifyAndClearExpectations(observer());
+
+  VerifyMigratedNigoriWithTimestamp(1, CUSTOM_PASSPHRASE, kCustomPass);
+}
+
 }  // namespace syncer