// Copyright 2026 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/sync/model/crypto/agile_symmetric_key_set.h"

#include <string>
#include <vector>

#include "base/base64.h"
#include "components/sync/model/crypto/agile_symmetric_key.h"
#include "components/sync/model/crypto/nigori.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace syncer {
namespace {

using ::testing::Optional;

TEST(AgileSymmetricKeySetTest, ShouldCreateEmpty) {
  const std::unique_ptr<AgileSymmetricKeySet> keyset =
      AgileSymmetricKeySet::CreateEmpty();
  ASSERT_NE(keyset, nullptr);
  EXPECT_EQ(keyset->size(), 0u);
  EXPECT_EQ(keyset->primary_key_id(), 0u);

  // Encryption and decryption return nullopt on empty keyset.
  EXPECT_EQ(keyset->Encrypt({1, 2, 3}), std::nullopt);
  sync_pb::EncryptedData empty_proto;
  EXPECT_EQ(keyset->Decrypt(empty_proto), std::nullopt);
}

TEST(AgileSymmetricKeySetTest, ShouldRotateAndEncryptDecryptAEAD) {
  const std::unique_ptr<AgileSymmetricKeySet> keyset =
      AgileSymmetricKeySet::CreateEmpty();
  ASSERT_NE(keyset, nullptr);

  // Rotate to generate primary key.
  const uint32_t key_id1 = keyset->RotatePrimaryToNewlyGeneratedRandomKey();
  EXPECT_NE(key_id1, 0u);
  EXPECT_EQ(keyset->size(), 1u);
  EXPECT_EQ(keyset->primary_key_id(), key_id1);

  // Encrypt.
  const std::vector<uint8_t> plaintext = {1, 2, 3, 4};
  const std::optional<sync_pb::EncryptedData> encrypted =
      keyset->Encrypt(plaintext);
  ASSERT_TRUE(encrypted.has_value());

  // Verified modern fields are populated, legacy are empty.
  EXPECT_FALSE(encrypted->has_key_name());
  EXPECT_FALSE(encrypted->has_blob());
  EXPECT_TRUE(encrypted->has_blob_v2());
  EXPECT_TRUE(encrypted->has_key_id_v2());
  EXPECT_EQ(encrypted->key_id_v2(), key_id1);

  // Decrypt.
  EXPECT_THAT(keyset->Decrypt(*encrypted), Optional(plaintext));

  // Rotate again.
  const uint32_t key_id2 = keyset->RotatePrimaryToNewlyGeneratedRandomKey();
  EXPECT_NE(key_id2, 0u);
  EXPECT_NE(key_id1, key_id2);
  EXPECT_EQ(keyset->size(), 2u);
  EXPECT_EQ(keyset->primary_key_id(), key_id2);

  // Encrypt under new key.
  const std::optional<sync_pb::EncryptedData> encrypted2 =
      keyset->Encrypt(plaintext);
  ASSERT_TRUE(encrypted2.has_value());
  EXPECT_EQ(encrypted2->key_id_v2(), key_id2);

  // Decrypt under new key -> succeeds.
  EXPECT_THAT(keyset->Decrypt(*encrypted2), Optional(plaintext));

  // Decrypt the old payload encrypted *before* rotation -> still succeeds.
  EXPECT_THAT(keyset->Decrypt(*encrypted), Optional(plaintext));
}

TEST(AgileSymmetricKeySetTest, ShouldSupportLegacyNigoriCompatibility) {
  const std::string kUserKey = "1234567890123456";
  const std::string kEncryptionKey = "abcdefghijklmnop";
  const std::string kMacKey = "qrstuvwxyz123456";

  // Create legacy Nigori class instance.
  const std::unique_ptr<Nigori> raw_nigori = Nigori::CreateByImport(
      NigoriPassKey::ForTesting(), kUserKey, kEncryptionKey, kMacKey);
  ASSERT_NE(raw_nigori, nullptr);
  const std::string expected_legacy_name = raw_nigori->GetKeyName();

  // 1. Manually build a proto containing a legacy key to verify FromProto.
  sync_pb::AgileSymmetricKeySet keyset_proto;
  keyset_proto.set_primary_key_id(42u);  // Arbitrary ID.
  sync_pb::AgileSymmetricKeySet::Key* key_entry = keyset_proto.add_key();
  key_entry->set_key_id(42u);

  sync_pb::AgileSymmetricKey* agile_key_proto = key_entry->mutable_key_data();
  sync_pb::NigoriKey* legacy_proto = agile_key_proto->mutable_legacy_nigori();
  legacy_proto->set_deprecated_user_key(kUserKey);
  legacy_proto->set_encryption_key(kEncryptionKey);
  legacy_proto->set_mac_key(kMacKey);

  const std::unique_ptr<AgileSymmetricKeySet> keyset =
      AgileSymmetricKeySet::FromProto(keyset_proto);
  ASSERT_NE(keyset, nullptr);
  EXPECT_EQ(keyset->size(), 1u);
  EXPECT_EQ(keyset->primary_key_id(), 42u);

  // 2. Encrypt under legacy key -> verified legacy-only fields.
  const std::vector<uint8_t> plaintext = {'h', 'e', 'l', 'l', 'o'};
  const std::optional<sync_pb::EncryptedData> encrypted =
      keyset->Encrypt(plaintext);
  ASSERT_TRUE(encrypted.has_value());

  EXPECT_TRUE(encrypted->has_key_name());
  EXPECT_EQ(encrypted->key_name(), expected_legacy_name);
  EXPECT_TRUE(encrypted->has_blob());  // Legacy Base64 blob.
  EXPECT_FALSE(encrypted->has_blob_v2());
  EXPECT_FALSE(encrypted->has_key_id_v2());

  // Modern decryption works.
  EXPECT_THAT(keyset->Decrypt(*encrypted), Optional(plaintext));

  // 3. Decrypt payload generated by an *older non-agile client* (only legacy
  // fields).
  sync_pb::EncryptedData legacy_only_encrypted;
  legacy_only_encrypted.set_key_name(expected_legacy_name);

  // Encrypt using raw Nigori (outputs Base64 ciphertext).
  const std::string raw_plaintext = "hello";
  const std::string raw_base64_ciphertext = raw_nigori->Encrypt(raw_plaintext);
  legacy_only_encrypted.set_blob(raw_base64_ciphertext);

  // AgileSymmetricKeySet decrypts this legacy-only payload using O(N) lookup
  // fallback.
  EXPECT_THAT(keyset->Decrypt(legacy_only_encrypted), Optional(plaintext));

  // 4. Rotate the keyset containing a legacy key to a new modern key.
  const uint32_t modern_key_id =
      keyset->RotatePrimaryToNewlyGeneratedRandomKey();
  EXPECT_NE(modern_key_id, 42u);
  EXPECT_EQ(keyset->size(), 2u);
  EXPECT_EQ(keyset->primary_key_id(), modern_key_id);

  // Encrypt under newly rotated key -> modern-only, no legacy double-write.
  const std::optional<sync_pb::EncryptedData> encrypted_modern =
      keyset->Encrypt(plaintext);
  ASSERT_TRUE(encrypted_modern.has_value());
  EXPECT_FALSE(encrypted_modern->has_key_name());
  EXPECT_FALSE(encrypted_modern->has_blob());
  EXPECT_TRUE(encrypted_modern->has_blob_v2());
  EXPECT_EQ(encrypted_modern->key_id_v2(), modern_key_id);

  // Decrypt new modern payload -> succeeds.
  EXPECT_THAT(keyset->Decrypt(*encrypted_modern), Optional(plaintext));

  // Decrypt old legacy payload -> still succeeds (old key is kept active in
  // keyset).
  EXPECT_THAT(keyset->Decrypt(legacy_only_encrypted), Optional(plaintext));
}

TEST(AgileSymmetricKeySetTest, ShouldRoundTripProto) {
  const std::unique_ptr<AgileSymmetricKeySet> original =
      AgileSymmetricKeySet::CreateEmpty();
  const uint32_t key_id = original->RotatePrimaryToNewlyGeneratedRandomKey();

  const sync_pb::AgileSymmetricKeySet proto = original->ToProto();
  EXPECT_EQ(proto.primary_key_id(), key_id);
  ASSERT_EQ(proto.key().size(), 1);
  EXPECT_EQ(proto.key(0).key_id(), key_id);
  EXPECT_TRUE(proto.key(0).key_data().has_aes_256_gcm());

  const std::unique_ptr<AgileSymmetricKeySet> restored =
      AgileSymmetricKeySet::FromProto(proto);
  ASSERT_NE(restored, nullptr);
  EXPECT_EQ(restored->size(), 1u);
  EXPECT_EQ(restored->primary_key_id(), key_id);

  const std::vector<uint8_t> plaintext = {1, 2, 3, 4};
  const std::optional<sync_pb::EncryptedData> encrypted =
      original->Encrypt(plaintext);
  ASSERT_TRUE(encrypted.has_value());

  EXPECT_THAT(restored->Decrypt(*encrypted), Optional(plaintext));
}

TEST(AgileSymmetricKeySetTest, FromProtoShouldRejectInvalidInvariants) {
  sync_pb::AgileSymmetricKeySet proto;

  // 1. Primary key ID is nonzero but keyset is empty.
  proto.set_primary_key_id(42u);
  EXPECT_EQ(AgileSymmetricKeySet::FromProto(proto), nullptr);

  // 2. Keyset has keys but primary key ID is unset/zero.
  proto.clear_primary_key_id();
  sync_pb::AgileSymmetricKeySet::Key* key_entry = proto.add_key();
  key_entry->set_key_id(42u);
  key_entry->mutable_key_data()->mutable_aes_256_gcm()->set_key(
      std::string(32, '1'));
  EXPECT_EQ(AgileSymmetricKeySet::FromProto(proto), nullptr);

  // 3. Keyset has keys, primary is nonzero, but matches no key in the set.
  proto.set_primary_key_id(43u);  // Matches nothing.
  EXPECT_EQ(AgileSymmetricKeySet::FromProto(proto), nullptr);

  // 4. Duplicate key IDs in proto.
  proto.set_primary_key_id(42u);
  sync_pb::AgileSymmetricKeySet::Key* key_entry2 = proto.add_key();
  key_entry2->set_key_id(42u);  // Duplicate ID.
  key_entry2->mutable_key_data()->mutable_chacha20_poly1305()->set_key(
      std::string(32, '2'));
  EXPECT_EQ(AgileSymmetricKeySet::FromProto(proto), nullptr);
}

TEST(AgileSymmetricKeySetTest, FromProtoShouldTolerateMalformedNonPrimaryKeys) {
  sync_pb::AgileSymmetricKeySet proto;
  proto.set_primary_key_id(42u);

  // Key 1: ID 42 (primary, valid GCM-256 key, size 32).
  sync_pb::AgileSymmetricKeySet::Key* key1 = proto.add_key();
  key1->set_key_id(42u);
  key1->mutable_key_data()->mutable_aes_256_gcm()->set_key(
      std::string(32, '1'));

  // Key 2: ID 43 (non-primary, malformed key - invalid size 15).
  sync_pb::AgileSymmetricKeySet::Key* key2 = proto.add_key();
  key2->set_key_id(43u);
  key2->mutable_key_data()->mutable_aes_256_gcm()->set_key(
      std::string(15, '2'));

  // 1. FromProto works and ignores Key 2.
  const std::unique_ptr<AgileSymmetricKeySet> keyset =
      AgileSymmetricKeySet::FromProto(proto);
  ASSERT_NE(keyset, nullptr);
  EXPECT_EQ(keyset->size(), 1u);  // Key 2 skipped.
  const std::vector<uint8_t> plaintext = {9, 9, 9};
  const std::optional<sync_pb::EncryptedData> encrypted =
      keyset->Encrypt(plaintext);
  ASSERT_TRUE(encrypted.has_value());
  EXPECT_EQ(encrypted->key_id_v2(), 42u);

  EXPECT_THAT(keyset->Decrypt(*encrypted), Optional(plaintext));

  // 2. If the malformed key is the primary, FromProto fails.
  proto.set_primary_key_id(43u);  // Make the malformed key the primary.
  EXPECT_EQ(AgileSymmetricKeySet::FromProto(proto), nullptr);
}

TEST(AgileSymmetricKeySetTest, ShouldDecryptModernBlobWithoutKeyId) {
  const std::unique_ptr<AgileSymmetricKeySet> keyset =
      AgileSymmetricKeySet::CreateEmpty();

  const uint32_t key_id1 = keyset->RotatePrimaryToNewlyGeneratedRandomKey();
  const std::vector<uint8_t> plaintext = {1, 2, 3, 4};

  // Encrypt under key1 (primary).
  std::optional<sync_pb::EncryptedData> encrypted = keyset->Encrypt(plaintext);
  ASSERT_TRUE(encrypted.has_value());
  ASSERT_TRUE(encrypted->has_key_id_v2());

  // Clear key_id_v2.
  encrypted->clear_key_id_v2();

  // Decrypt should still succeed (primary key tried first).
  EXPECT_THAT(keyset->Decrypt(*encrypted), Optional(plaintext));

  // Rotate to key2 (key1 is now non-primary).
  const uint32_t key_id2 = keyset->RotatePrimaryToNewlyGeneratedRandomKey();
  ASSERT_NE(key_id1, key_id2);

  // Encrypt under key2 (new primary) to get a new ciphertext.
  std::optional<sync_pb::EncryptedData> encrypted2 = keyset->Encrypt(plaintext);
  ASSERT_TRUE(encrypted2.has_value());
  encrypted2->clear_key_id_v2();

  // Decrypt of encrypted2 should succeed (primary key tried first).
  EXPECT_THAT(keyset->Decrypt(*encrypted2), Optional(plaintext));

  // Decrypt of old encrypted (under key1, now non-primary) should also succeed
  // (fallback to other keys).
  EXPECT_THAT(keyset->Decrypt(*encrypted), Optional(plaintext));
}

}  // namespace
}  // namespace syncer