| // Copyright 2017 The Fuchsia Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| #include "src/lib/util/encrypted_message_util.h" |
| |
| #include <string> |
| #include <utility> |
| |
| #include <gtest/gtest.h> |
| |
| #include "src/pb/encrypted_message.pb.h" |
| #include "src/pb/envelope.pb.h" |
| #include "src/pb/key.pb.h" |
| #include "src/pb/observation.pb.h" |
| #include "third_party/tink/cc/cleartext_keyset_handle.h" |
| #include "third_party/tink/cc/hybrid_config.h" |
| #include "third_party/tink/cc/hybrid_decrypt_factory.h" |
| #include "third_party/tink/cc/hybrid_encrypt_factory.h" |
| #include "third_party/tink/cc/hybrid_key_templates.h" |
| #include "third_party/tink/cc/keyset_handle.h" |
| #include "third_party/tink/proto/tink.pb.h" |
| |
| namespace cobalt::util { |
| |
| constexpr char kShufflerContextInfo[] = "cobalt-1.0-shuffler"; |
| constexpr char kAnalyzerContextInfo[] = "cobalt-1.0-analyzer"; |
| |
| std::string MakeCobaltEncryptionKeyBytes(const std::string& key_bytes, uint32_t key_index, |
| CobaltEncryptionKey::KeyPurpose purpose) { |
| CobaltEncryptionKey key; |
| key.set_serialized_key(key_bytes); |
| key.set_key_index(key_index); |
| key.set_purpose(purpose); |
| std::string cobalt_key_bytes; |
| EXPECT_TRUE(key.SerializeToString(&cobalt_key_bytes)); |
| return cobalt_key_bytes; |
| } |
| |
| // Tests the use of the no-encryption option. |
| TEST(EncryptedMessageUtilTest, NoEncryption) { |
| // Make a dummy observation. |
| std::string obs_id = "obs_id"; |
| Observation observation; |
| observation.set_random_id(obs_id); |
| |
| // Make an EncryptedMessageMaker that outputs plaintext for testing purposes. |
| util::PinnedUniquePtr<EncryptedMessageMaker> maker(EncryptedMessageMaker::MakeUnencrypted()); |
| // Encrypt the dummy observation. |
| EncryptedMessage encrypted_message; |
| EXPECT_TRUE(maker->Encrypt(observation, &encrypted_message)); |
| |
| // Decrypt and check. |
| observation.Clear(); |
| EXPECT_TRUE(observation.random_id().empty()); |
| EXPECT_TRUE(observation.ParseFromString(encrypted_message.ciphertext())); |
| EXPECT_EQ(obs_id, observation.random_id()); |
| } |
| |
| class EncryptedMessageMakerTest : public ::testing::Test { |
| public: |
| // Get a serialized Keyset proto for a public key. |
| [[nodiscard]] std::string GetPublicKeysetBytes() const { |
| // Get the keyset protobuf message itself. |
| const google::crypto::tink::Keyset public_keyset = |
| ::crypto::tink::CleartextKeysetHandle::GetKeyset(*GetPublicKeysetHandle()); |
| EXPECT_EQ(1, public_keyset.key_size()); |
| |
| // Serialize and encode the public keyset. |
| std::string public_keyset_bytes; |
| EXPECT_TRUE(public_keyset.SerializeToString(&public_keyset_bytes)); |
| return public_keyset_bytes; |
| } |
| |
| // Decrypt the specified ciphertext. |
| std::string Decrypt(const std::string& ciphertext, const std::string& context_info) { |
| // Obtain a decrypter to be able to check the encrypted dummy observation. |
| auto decrypter_result = keyset_handle_->GetPrimitive<::crypto::tink::HybridDecrypt>(); |
| EXPECT_TRUE(decrypter_result.ok()); |
| auto decrypter = std::move(decrypter_result.ValueOrDie()); |
| auto decrypted_result = decrypter->Decrypt(ciphertext, context_info); |
| EXPECT_TRUE(decrypted_result.ok()); |
| return decrypted_result.ValueOrDie(); |
| } |
| |
| protected: |
| void SetUp() override { |
| auto status = ::crypto::tink::HybridConfig::Register(); |
| EXPECT_TRUE(status.ok()); |
| |
| // Create a new key pair for testing. |
| auto keyset_handle_result = ::crypto::tink::KeysetHandle::GenerateNew( |
| ::crypto::tink::HybridKeyTemplates::EciesP256HkdfHmacSha256Aes128Gcm()); |
| EXPECT_TRUE(keyset_handle_result.ok()); |
| keyset_handle_ = std::move(keyset_handle_result.ValueOrDie()); |
| } |
| |
| private: |
| [[nodiscard]] std::unique_ptr<::crypto::tink::KeysetHandle> GetPublicKeysetHandle() const { |
| auto public_keyset_handle_result = keyset_handle_->GetPublicKeysetHandle(); |
| EXPECT_TRUE(public_keyset_handle_result.ok()); |
| return std::move(public_keyset_handle_result.ValueOrDie()); |
| } |
| |
| std::unique_ptr<::crypto::tink::KeysetHandle> keyset_handle_; |
| }; |
| |
| // Try to roundtrip an observation through a message encrypter. |
| TEST_F(EncryptedMessageMakerTest, EncryptObservation) { |
| uint32_t key_index = 1; |
| auto key_bytes = MakeCobaltEncryptionKeyBytes(GetPublicKeysetBytes(), key_index, |
| CobaltEncryptionKey::ANALYZER); |
| auto encrypted_message_maker_or_status = EncryptedMessageMaker::MakeForObservations(key_bytes); |
| EXPECT_TRUE(encrypted_message_maker_or_status.ok()); |
| util::PinnedUniquePtr<EncryptedMessageMaker> maker( |
| std::move(encrypted_message_maker_or_status).value()); |
| |
| // Build an observation. |
| std::string obs_id = "obs_id"; |
| Observation observation; |
| observation.set_random_id(obs_id); |
| |
| // Encrypt the observation. |
| EncryptedMessage encrypted_message; |
| EXPECT_TRUE(maker->Encrypt(observation, &encrypted_message)); |
| |
| EXPECT_EQ(key_index, encrypted_message.key_index()); |
| |
| // Decrypt the observation. |
| std::string decrypted = Decrypt(encrypted_message.ciphertext(), kAnalyzerContextInfo); |
| observation.Clear(); |
| |
| // Check that the observation was correctly round-tripped. |
| EXPECT_TRUE(observation.random_id().empty()); |
| EXPECT_TRUE(observation.ParseFromString(decrypted)); |
| EXPECT_EQ(obs_id, observation.random_id()); |
| } |
| |
| // Try to roundtrip an envelope through a message encrypter. |
| TEST_F(EncryptedMessageMakerTest, EncryptEnvelope) { |
| uint32_t key_index = 1; |
| auto key_bytes = MakeCobaltEncryptionKeyBytes(GetPublicKeysetBytes(), key_index, |
| CobaltEncryptionKey::SHUFFLER); |
| auto encrypted_message_maker_or_status = EncryptedMessageMaker::MakeForEnvelopes(key_bytes); |
| EXPECT_TRUE(encrypted_message_maker_or_status.ok()); |
| util::PinnedUniquePtr<EncryptedMessageMaker> maker( |
| std::move(encrypted_message_maker_or_status).value()); |
| |
| // Build an envelope proto. |
| uint32_t metric_id = 25; |
| Envelope envelope; |
| envelope.add_batch()->mutable_meta_data()->set_metric_id(metric_id); |
| |
| // Encrypt the envelope. |
| EncryptedMessage encrypted_message; |
| EXPECT_TRUE(maker->Encrypt(envelope, &encrypted_message)); |
| EXPECT_EQ(key_index, encrypted_message.key_index()); |
| |
| // Decrypt the envelope. |
| std::string decrypted = Decrypt(encrypted_message.ciphertext(), kShufflerContextInfo); |
| envelope.Clear(); |
| EXPECT_EQ(0, envelope.batch_size()); |
| EXPECT_TRUE(envelope.ParseFromString(decrypted)); |
| |
| // Check to see that the envelope rountripped correctly. |
| EXPECT_EQ(1, envelope.batch_size()); |
| EXPECT_EQ(metric_id, envelope.batch(0).meta_data().metric_id()); |
| } |
| |
| // Expect an error if the key_index field is set to 0 or unset. |
| TEST_F(EncryptedMessageMakerTest, ZeroKeyIndex) { |
| auto key_bytes = |
| MakeCobaltEncryptionKeyBytes(GetPublicKeysetBytes(), 0, CobaltEncryptionKey::SHUFFLER); |
| auto result = EncryptedMessageMaker::MakeForEnvelopes(key_bytes); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| |
| key_bytes = |
| MakeCobaltEncryptionKeyBytes(GetPublicKeysetBytes(), 0, CobaltEncryptionKey::ANALYZER); |
| result = EncryptedMessageMaker::MakeForObservations(key_bytes); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| } |
| |
| // Expect an error if the purpose field is unset. |
| TEST_F(EncryptedMessageMakerTest, PurposeUnset) { |
| auto key_bytes = |
| MakeCobaltEncryptionKeyBytes(GetPublicKeysetBytes(), 1, CobaltEncryptionKey::UNSET); |
| auto result = EncryptedMessageMaker::MakeForEnvelopes(key_bytes); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| |
| key_bytes = MakeCobaltEncryptionKeyBytes(GetPublicKeysetBytes(), 1, CobaltEncryptionKey::UNSET); |
| result = EncryptedMessageMaker::MakeForObservations(key_bytes); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| } |
| |
| // Expect an error if trying to use an analyzer key to encrypt envelopes or |
| // trying to use a shuffler key to encrypt observations. |
| TEST_F(EncryptedMessageMakerTest, WrongPurpose) { |
| auto key_bytes = |
| MakeCobaltEncryptionKeyBytes(GetPublicKeysetBytes(), 1, CobaltEncryptionKey::ANALYZER); |
| auto result = EncryptedMessageMaker::MakeForEnvelopes(key_bytes); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| |
| key_bytes = |
| MakeCobaltEncryptionKeyBytes(GetPublicKeysetBytes(), 1, CobaltEncryptionKey::SHUFFLER); |
| result = EncryptedMessageMaker::MakeForObservations(key_bytes); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| } |
| |
| // Expect an error if the string purported to be a serialized |
| // CobaltEncryptionKey is actually not. |
| TEST_F(EncryptedMessageMakerTest, NotASerializedCobaltEncryptionKey) { |
| auto result = EncryptedMessageMaker::MakeForEnvelopes("hello"); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| |
| result = EncryptedMessageMaker::MakeForObservations("hello"); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| } |
| |
| // Expect an error if the serialized key is invalid. |
| TEST_F(EncryptedMessageMakerTest, NotRealSerializedKey) { |
| auto key_bytes = MakeCobaltEncryptionKeyBytes("not key bytes", 1, CobaltEncryptionKey::SHUFFLER); |
| auto result = EncryptedMessageMaker::MakeForEnvelopes(key_bytes); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| |
| key_bytes = MakeCobaltEncryptionKeyBytes("not key bytes", 1, CobaltEncryptionKey::ANALYZER); |
| result = EncryptedMessageMaker::MakeForObservations(key_bytes); |
| EXPECT_EQ(StatusCode::INVALID_ARGUMENT, result.status().error_code()); |
| } |
| |
| } // namespace cobalt::util |