| // Copyright 2021 The Chromium OS 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 "missive/encryption/encryption_module.h" |
| |
| #include <memory> |
| #include <string> |
| #include <utility> |
| #include <vector> |
| |
| #include <base/bind.h> |
| #include <base/containers/flat_map.h> |
| #include <base/hash/hash.h> |
| #include <base/rand_util.h> |
| #include <base/strings/strcat.h> |
| #include <base/synchronization/waitable_event.h> |
| #include <base/feature_list.h> |
| #include <base/task/thread_pool.h> |
| #include <base/test/task_environment.h> |
| #include <base/time/time.h> |
| #include <gmock/gmock.h> |
| #include <gtest/gtest.h> |
| |
| #include "missive/encryption/decryption.h" |
| #include "missive/encryption/encryption.h" |
| #include "missive/encryption/encryption_module_interface.h" |
| #include "missive/encryption/primitives.h" |
| #include "missive/encryption/scoped_encryption_feature.h" |
| #include "missive/encryption/testing_primitives.h" |
| #include "missive/proto/record.pb.h" |
| #include "missive/util/status.h" |
| #include "missive/util/status_macros.h" |
| #include "missive/util/statusor.h" |
| #include "missive/util/test_support_callbacks.h" |
| |
| using ::testing::Eq; |
| using ::testing::Ne; |
| |
| namespace reporting { |
| namespace { |
| |
| class EncryptionModuleTest : public ::testing::Test { |
| protected: |
| EncryptionModuleTest() = default; |
| |
| void SetUp() override { |
| encryption_module_ = EncryptionModule::Create(); |
| |
| auto decryptor_result = test::Decryptor::Create(); |
| ASSERT_OK(decryptor_result.status()) << decryptor_result.status(); |
| decryptor_ = std::move(decryptor_result.ValueOrDie()); |
| } |
| |
| StatusOr<EncryptedRecord> EncryptSync(base::StringPiece data) { |
| test::TestEvent<StatusOr<EncryptedRecord>> encrypt_record; |
| encryption_module_->EncryptRecord(data, encrypt_record.cb()); |
| return encrypt_record.result(); |
| } |
| |
| StatusOr<std::string> DecryptSync( |
| std::pair<std::string /*shared_secret*/, std::string /*encrypted_data*/> |
| encrypted) { |
| test::TestEvent<StatusOr<test::Decryptor::Handle*>> open_decrypt; |
| decryptor_->OpenRecord(encrypted.first, open_decrypt.cb()); |
| auto open_decrypt_result = open_decrypt.result(); |
| RETURN_IF_ERROR(open_decrypt_result.status()); |
| test::Decryptor::Handle* const dec_handle = |
| open_decrypt_result.ValueOrDie(); |
| |
| test::TestEvent<Status> add_decrypt; |
| dec_handle->AddToRecord(encrypted.second, add_decrypt.cb()); |
| RETURN_IF_ERROR(add_decrypt.result()); |
| |
| std::string decrypted_string; |
| test::TestEvent<Status> close_decrypt; |
| dec_handle->CloseRecord(base::BindOnce( |
| [](std::string* decrypted_string, |
| base::OnceCallback<void(Status)> close_cb, |
| StatusOr<base::StringPiece> result) { |
| if (!result.ok()) { |
| std::move(close_cb).Run(result.status()); |
| return; |
| } |
| *decrypted_string = std::string(result.ValueOrDie()); |
| std::move(close_cb).Run(Status::StatusOK()); |
| }, |
| base::Unretained(&decrypted_string), close_decrypt.cb())); |
| RETURN_IF_ERROR(close_decrypt.result()); |
| return decrypted_string; |
| } |
| |
| StatusOr<std::string> DecryptMatchingSecret( |
| Encryptor::PublicKeyId public_key_id, base::StringPiece encrypted_key) { |
| // Retrieve private key that matches public key hash. |
| test::TestEvent<StatusOr<std::string>> retrieve_private_key; |
| decryptor_->RetrieveMatchingPrivateKey(public_key_id, |
| retrieve_private_key.cb()); |
| ASSIGN_OR_RETURN(std::string private_key, retrieve_private_key.result()); |
| // Decrypt symmetric key with that private key and peer public key. |
| ASSIGN_OR_RETURN(std::string shared_secret, |
| decryptor_->DecryptSecret(private_key, encrypted_key)); |
| return shared_secret; |
| } |
| |
| Status AddNewKeyPair() { |
| // Generate new pair of private key and public value. |
| uint8_t out_public_value[kKeySize]; |
| uint8_t out_private_key[kKeySize]; |
| test::GenerateEncryptionKeyPair(out_private_key, out_public_value); |
| |
| test::TestEvent<StatusOr<Encryptor::PublicKeyId>> record_keys; |
| decryptor_->RecordKeyPair( |
| std::string(reinterpret_cast<const char*>(out_private_key), kKeySize), |
| std::string(reinterpret_cast<const char*>(out_public_value), kKeySize), |
| record_keys.cb()); |
| ASSIGN_OR_RETURN(Encryptor::PublicKeyId new_public_key_id, |
| record_keys.result()); |
| test::TestEvent<Status> set_public_key; |
| encryption_module_->UpdateAsymmetricKey( |
| std::string(reinterpret_cast<const char*>(out_public_value), kKeySize), |
| new_public_key_id, set_public_key.cb()); |
| RETURN_IF_ERROR(set_public_key.result()); |
| return Status::StatusOK(); |
| } |
| |
| base::test::TaskEnvironment task_environment_{ |
| base::test::TaskEnvironment::TimeSource::MOCK_TIME}; |
| |
| scoped_refptr<EncryptionModuleInterface> encryption_module_; |
| scoped_refptr<test::Decryptor> decryptor_; |
| |
| private: |
| test::ScopedEncryptionFeature encryption_feature_{/*enable=*/true}; |
| }; |
| |
| TEST_F(EncryptionModuleTest, EncryptAndDecrypt) { |
| constexpr char kTestString[] = "ABCDEF"; |
| |
| // Register new pair of private key and public value. |
| ASSERT_OK(AddNewKeyPair()); |
| |
| // Encrypt the test string using the last public value. |
| const auto encrypted_result = EncryptSync(kTestString); |
| ASSERT_OK(encrypted_result.status()) << encrypted_result.status(); |
| |
| // Decrypt shared secret with private asymmetric key. |
| auto decrypt_secret_result = DecryptMatchingSecret( |
| encrypted_result.ValueOrDie().encryption_info().public_key_id(), |
| encrypted_result.ValueOrDie().encryption_info().encryption_key()); |
| ASSERT_OK(decrypt_secret_result.status()) << decrypt_secret_result.status(); |
| |
| // Decrypt back. |
| const auto decrypted_result = DecryptSync( |
| std::make_pair(decrypt_secret_result.ValueOrDie(), |
| encrypted_result.ValueOrDie().encrypted_wrapped_record())); |
| ASSERT_OK(decrypted_result.status()) << decrypted_result.status(); |
| |
| EXPECT_THAT(decrypted_result.ValueOrDie(), ::testing::StrEq(kTestString)); |
| } |
| |
| TEST_F(EncryptionModuleTest, EncryptionDisabled) { |
| constexpr char kTestString[] = "ABCDEF"; |
| |
| // Disable encryption for this test. |
| test::ScopedEncryptionFeature encryption_feature_{/*enable=*/false}; |
| |
| // Encrypt the test string. |
| const auto encrypted_result = EncryptSync(kTestString); |
| ASSERT_OK(encrypted_result.status()); |
| |
| // Expect the result to be identical to the original record, |
| // and have no encryption_info. |
| EXPECT_EQ(encrypted_result.ValueOrDie().encrypted_wrapped_record(), |
| kTestString); |
| EXPECT_FALSE(encrypted_result.ValueOrDie().has_encryption_info()); |
| } |
| |
| TEST_F(EncryptionModuleTest, PublicKeyUpdate) { |
| constexpr char kTestString[] = "ABCDEF"; |
| |
| // No key yet, attempt to encrypt the test string. |
| ASSERT_FALSE(encryption_module_->has_encryption_key()); |
| ASSERT_TRUE(encryption_module_->need_encryption_key()); |
| auto encrypted_result = EncryptSync(kTestString); |
| EXPECT_EQ(encrypted_result.status().error_code(), error::NOT_FOUND); |
| |
| // Register new pair of private key and public value. |
| ASSERT_OK(AddNewKeyPair()); |
| ASSERT_TRUE(encryption_module_->has_encryption_key()); |
| ASSERT_FALSE(encryption_module_->need_encryption_key()); |
| |
| // Encrypt the test string using the last public value. |
| encrypted_result = EncryptSync(kTestString); |
| ASSERT_OK(encrypted_result.status()) << encrypted_result.status(); |
| |
| // Simulate short wait. Key is still available and not needed. |
| task_environment_.FastForwardBy(base::TimeDelta::FromHours(8)); |
| ASSERT_TRUE(encryption_module_->has_encryption_key()); |
| ASSERT_FALSE(encryption_module_->need_encryption_key()); |
| encrypted_result = EncryptSync(kTestString); |
| ASSERT_OK(encrypted_result.status()) << encrypted_result.status(); |
| |
| // Simulate long wait. Key is still available, but is needed now. |
| task_environment_.FastForwardBy(base::TimeDelta::FromDays(1)); |
| ASSERT_TRUE(encryption_module_->has_encryption_key()); |
| ASSERT_TRUE(encryption_module_->need_encryption_key()); |
| encrypted_result = EncryptSync(kTestString); |
| ASSERT_OK(encrypted_result.status()) << encrypted_result.status(); |
| |
| // Register one more pair of private key and public value. |
| ASSERT_OK(AddNewKeyPair()); |
| ASSERT_TRUE(encryption_module_->has_encryption_key()); |
| ASSERT_FALSE(encryption_module_->need_encryption_key()); |
| encrypted_result = EncryptSync(kTestString); |
| ASSERT_OK(encrypted_result.status()) << encrypted_result.status(); |
| } |
| |
| TEST_F(EncryptionModuleTest, EncryptAndDecryptMultiple) { |
| constexpr const char* kTestStrings[] = {"Rec1", "Rec22", "Rec333", |
| "Rec4444", "Rec55555", "Rec666666"}; |
| // Encrypted records. |
| std::vector<EncryptedRecord> encrypted_records; |
| |
| // 1. Register first key pair. |
| ASSERT_OK(AddNewKeyPair()); |
| |
| // 2. Encrypt 3 test strings. |
| for (const char* test_string : |
| {kTestStrings[0], kTestStrings[1], kTestStrings[2]}) { |
| const auto encrypted_result = EncryptSync(test_string); |
| ASSERT_OK(encrypted_result.status()) << encrypted_result.status(); |
| encrypted_records.emplace_back(encrypted_result.ValueOrDie()); |
| } |
| |
| // 3. Register second key pair. |
| ASSERT_OK(AddNewKeyPair()); |
| |
| // 4. Encrypt 2 test strings. |
| for (const char* test_string : {kTestStrings[3], kTestStrings[4]}) { |
| const auto encrypted_result = EncryptSync(test_string); |
| ASSERT_OK(encrypted_result.status()) << encrypted_result.status(); |
| encrypted_records.emplace_back(encrypted_result.ValueOrDie()); |
| } |
| |
| // 3. Register third key pair. |
| ASSERT_OK(AddNewKeyPair()); |
| |
| // 4. Encrypt one more test strings. |
| for (const char* test_string : {kTestStrings[5]}) { |
| const auto encrypted_result = EncryptSync(test_string); |
| ASSERT_OK(encrypted_result.status()) << encrypted_result.status(); |
| encrypted_records.emplace_back(encrypted_result.ValueOrDie()); |
| } |
| |
| // For every encrypted record: |
| for (size_t i = 0; i < encrypted_records.size(); ++i) { |
| // Decrypt encrypted_key with private asymmetric key. |
| auto decrypt_secret_result = DecryptMatchingSecret( |
| encrypted_records[i].encryption_info().public_key_id(), |
| encrypted_records[i].encryption_info().encryption_key()); |
| ASSERT_OK(decrypt_secret_result.status()) << decrypt_secret_result.status(); |
| |
| // Decrypt back. |
| const auto decrypted_result = DecryptSync( |
| std::make_pair(decrypt_secret_result.ValueOrDie(), |
| encrypted_records[i].encrypted_wrapped_record())); |
| ASSERT_OK(decrypted_result.status()) << decrypted_result.status(); |
| |
| // Verify match. |
| EXPECT_THAT(decrypted_result.ValueOrDie(), |
| ::testing::StrEq(kTestStrings[i])); |
| } |
| } |
| |
| TEST_F(EncryptionModuleTest, EncryptAndDecryptMultipleParallel) { |
| // Context of single encryption. Self-destructs upon completion or failure. |
| class SingleEncryptionContext { |
| public: |
| SingleEncryptionContext( |
| base::StringPiece test_string, |
| base::StringPiece public_key, |
| Encryptor::PublicKeyId public_key_id, |
| scoped_refptr<EncryptionModuleInterface> encryption_module, |
| base::OnceCallback<void(StatusOr<EncryptedRecord>)> response) |
| : test_string_(test_string), |
| public_key_(public_key), |
| public_key_id_(public_key_id), |
| encryption_module_(encryption_module), |
| response_(std::move(response)) {} |
| |
| SingleEncryptionContext(const SingleEncryptionContext& other) = delete; |
| SingleEncryptionContext& operator=(const SingleEncryptionContext& other) = |
| delete; |
| |
| ~SingleEncryptionContext() { |
| DCHECK(!response_) << "Self-destruct without prior response"; |
| } |
| |
| void Start() { |
| base::ThreadPool::PostTask( |
| FROM_HERE, base::BindOnce(&SingleEncryptionContext::SetPublicKey, |
| base::Unretained(this))); |
| } |
| |
| private: |
| void Respond(StatusOr<EncryptedRecord> result) { |
| std::move(response_).Run(result); |
| delete this; |
| } |
| void SetPublicKey() { |
| encryption_module_->UpdateAsymmetricKey( |
| public_key_, public_key_id_, |
| base::BindOnce( |
| [](SingleEncryptionContext* self, Status status) { |
| if (!status.ok()) { |
| self->Respond(status); |
| return; |
| } |
| base::ThreadPool::PostTask( |
| FROM_HERE, |
| base::BindOnce(&SingleEncryptionContext::EncryptRecord, |
| base::Unretained(self))); |
| }, |
| base::Unretained(this))); |
| } |
| void EncryptRecord() { |
| encryption_module_->EncryptRecord( |
| test_string_, |
| base::BindOnce( |
| [](SingleEncryptionContext* self, |
| StatusOr<EncryptedRecord> encryption_result) { |
| base::ThreadPool::PostTask( |
| FROM_HERE, |
| base::BindOnce(&SingleEncryptionContext::Respond, |
| base::Unretained(self), encryption_result)); |
| }, |
| base::Unretained(this))); |
| } |
| |
| private: |
| const std::string test_string_; |
| const std::string public_key_; |
| const Encryptor::PublicKeyId public_key_id_; |
| const scoped_refptr<EncryptionModuleInterface> encryption_module_; |
| base::OnceCallback<void(StatusOr<EncryptedRecord>)> response_; |
| }; |
| |
| // Context of single decryption. Self-destructs upon completion or failure. |
| class SingleDecryptionContext { |
| public: |
| SingleDecryptionContext( |
| const EncryptedRecord& encrypted_record, |
| scoped_refptr<test::Decryptor> decryptor, |
| base::OnceCallback<void(StatusOr<base::StringPiece>)> response) |
| : encrypted_record_(encrypted_record), |
| decryptor_(decryptor), |
| response_(std::move(response)) {} |
| |
| SingleDecryptionContext(const SingleDecryptionContext& other) = delete; |
| SingleDecryptionContext& operator=(const SingleDecryptionContext& other) = |
| delete; |
| |
| ~SingleDecryptionContext() { |
| DCHECK(!response_) << "Self-destruct without prior response"; |
| } |
| |
| void Start() { |
| base::ThreadPool::PostTask( |
| FROM_HERE, |
| base::BindOnce(&SingleDecryptionContext::RetrieveMatchingPrivateKey, |
| base::Unretained(this))); |
| } |
| |
| private: |
| void Respond(StatusOr<base::StringPiece> result) { |
| std::move(response_).Run(result); |
| delete this; |
| } |
| |
| void RetrieveMatchingPrivateKey() { |
| // Retrieve private key that matches public key hash. |
| decryptor_->RetrieveMatchingPrivateKey( |
| encrypted_record_.encryption_info().public_key_id(), |
| base::BindOnce( |
| [](SingleDecryptionContext* self, |
| StatusOr<std::string> private_key_result) { |
| if (!private_key_result.ok()) { |
| self->Respond(private_key_result.status()); |
| return; |
| } |
| base::ThreadPool::PostTask( |
| FROM_HERE, |
| base::BindOnce( |
| &SingleDecryptionContext::DecryptSharedSecret, |
| base::Unretained(self), |
| private_key_result.ValueOrDie())); |
| }, |
| base::Unretained(this))); |
| } |
| |
| void DecryptSharedSecret(base::StringPiece private_key) { |
| // Decrypt shared secret from private key and peer public key. |
| auto shared_secret_result = decryptor_->DecryptSecret( |
| private_key, encrypted_record_.encryption_info().encryption_key()); |
| if (!shared_secret_result.ok()) { |
| Respond(shared_secret_result.status()); |
| return; |
| } |
| base::ThreadPool::PostTask( |
| FROM_HERE, base::BindOnce(&SingleDecryptionContext::OpenRecord, |
| base::Unretained(this), |
| shared_secret_result.ValueOrDie())); |
| } |
| |
| void OpenRecord(base::StringPiece shared_secret) { |
| decryptor_->OpenRecord( |
| shared_secret, |
| base::BindOnce( |
| [](SingleDecryptionContext* self, |
| StatusOr<test::Decryptor::Handle*> handle_result) { |
| if (!handle_result.ok()) { |
| self->Respond(handle_result.status()); |
| return; |
| } |
| base::ThreadPool::PostTask( |
| FROM_HERE, |
| base::BindOnce( |
| &SingleDecryptionContext::AddToRecord, |
| base::Unretained(self), |
| base::Unretained(handle_result.ValueOrDie()))); |
| }, |
| base::Unretained(this))); |
| } |
| |
| void AddToRecord(test::Decryptor::Handle* handle) { |
| handle->AddToRecord( |
| encrypted_record_.encrypted_wrapped_record(), |
| base::BindOnce( |
| [](SingleDecryptionContext* self, test::Decryptor::Handle* handle, |
| Status status) { |
| if (!status.ok()) { |
| self->Respond(status); |
| return; |
| } |
| base::ThreadPool::PostTask( |
| FROM_HERE, |
| base::BindOnce(&SingleDecryptionContext::CloseRecord, |
| base::Unretained(self), |
| base::Unretained(handle))); |
| }, |
| base::Unretained(this), base::Unretained(handle))); |
| } |
| |
| void CloseRecord(test::Decryptor::Handle* handle) { |
| handle->CloseRecord(base::BindOnce( |
| [](SingleDecryptionContext* self, |
| StatusOr<base::StringPiece> decryption_result) { |
| self->Respond(decryption_result); |
| }, |
| base::Unretained(this))); |
| } |
| |
| private: |
| const EncryptedRecord encrypted_record_; |
| const scoped_refptr<test::Decryptor> decryptor_; |
| base::OnceCallback<void(StatusOr<base::StringPiece>)> response_; |
| }; |
| |
| constexpr std::array<const char*, 6> kTestStrings = { |
| "Rec1", "Rec22", "Rec333", "Rec4444", "Rec55555", "Rec666666"}; |
| |
| // Public and private key pairs in this test are reversed strings. |
| std::vector<std::string> private_key_strings; |
| std::vector<std::string> public_value_strings; |
| std::vector<Encryptor::PublicKeyId> public_value_ids; |
| for (size_t i = 0; i < 3; ++i) { |
| // Generate new pair of private key and public value. |
| uint8_t out_public_value[kKeySize]; |
| uint8_t out_private_key[kKeySize]; |
| test::GenerateEncryptionKeyPair(out_private_key, out_public_value); |
| private_key_strings.emplace_back( |
| reinterpret_cast<const char*>(out_private_key), kKeySize); |
| public_value_strings.emplace_back( |
| reinterpret_cast<const char*>(out_public_value), kKeySize); |
| } |
| |
| // Register all key pairs for decryption. |
| std::vector<test::TestEvent<StatusOr<Encryptor::PublicKeyId>>> record_results( |
| public_value_strings.size()); |
| for (size_t i = 0; i < public_value_strings.size(); ++i) { |
| base::ThreadPool::PostTask( |
| FROM_HERE, |
| base::BindOnce( |
| [](base::StringPiece private_key_string, |
| base::StringPiece public_key_string, |
| scoped_refptr<test::Decryptor> decryptor, |
| base::OnceCallback<void(StatusOr<Encryptor::PublicKeyId>)> |
| done_cb) { |
| decryptor->RecordKeyPair(private_key_string, public_key_string, |
| std::move(done_cb)); |
| }, |
| private_key_strings[i], public_value_strings[i], decryptor_, |
| record_results[i].cb())); |
| } |
| // Verify registration success. |
| for (auto& record_result : record_results) { |
| const auto result = record_result.result(); |
| ASSERT_OK(result.status()) << result.status(); |
| public_value_ids.push_back(result.ValueOrDie()); |
| } |
| |
| // Encrypt all records in parallel. |
| std::vector<test::TestEvent<StatusOr<EncryptedRecord>>> results( |
| kTestStrings.size()); |
| for (size_t i = 0; i < kTestStrings.size(); ++i) { |
| // Choose random key pair. |
| size_t i_key_pair = base::RandInt(0, public_value_strings.size() - 1); |
| (new SingleEncryptionContext( |
| kTestStrings[i], public_value_strings[i_key_pair], |
| public_value_ids[i_key_pair], encryption_module_, results[i].cb())) |
| ->Start(); |
| } |
| |
| // Decrypt all records in parallel. |
| std::vector<test::TestEvent<StatusOr<std::string>>> decryption_results( |
| kTestStrings.size()); |
| for (size_t i = 0; i < results.size(); ++i) { |
| // Verify encryption success. |
| const auto result = results[i].result(); |
| ASSERT_OK(result.status()) << result.status(); |
| // Decrypt and compare encrypted_record. |
| (new SingleDecryptionContext( |
| result.ValueOrDie(), decryptor_, |
| base::BindOnce( |
| [](base::OnceCallback<void(StatusOr<std::string>)> |
| decryption_result, |
| StatusOr<base::StringPiece> result) { |
| if (!result.ok()) { |
| std::move(decryption_result).Run(result.status()); |
| return; |
| } |
| std::move(decryption_result) |
| .Run(std::string(result.ValueOrDie())); |
| }, |
| decryption_results[i].cb()))) |
| ->Start(); |
| } |
| |
| // Verify decryption results. |
| for (size_t i = 0; i < decryption_results.size(); ++i) { |
| const auto decryption_result = decryption_results[i].result(); |
| ASSERT_OK(decryption_result.status()) << decryption_result.status(); |
| // Verify data match. |
| EXPECT_THAT(decryption_result.ValueOrDie(), |
| ::testing::StrEq(kTestStrings[i])); |
| } |
| } |
| } // namespace |
| } // namespace reporting |
