diff options
Diffstat (limited to 'tde2e/test/e2e.cpp')
| -rw-r--r-- | tde2e/test/e2e.cpp | 914 |
1 files changed, 914 insertions, 0 deletions
diff --git a/tde2e/test/e2e.cpp b/tde2e/test/e2e.cpp new file mode 100644 index 000000000..6ad625d06 --- /dev/null +++ b/tde2e/test/e2e.cpp @@ -0,0 +1,914 @@ +// +// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2025 +// +// Distributed under the Boost Software License, Version 1.0. (See accompanying +// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt) +// +#include "td/e2e/BitString.h" +#include "td/e2e/Blockchain.h" +#include "td/e2e/Call.h" +#include "td/e2e/CheckSharedSecret.h" +#include "td/e2e/Container.h" +#include "td/e2e/DecryptedKey.h" +#include "td/e2e/e2e_api.h" +#include "td/e2e/EncryptedKey.h" +#include "td/e2e/EncryptedStorage.h" +#include "td/e2e/MessageEncryption.h" +#include "td/e2e/Mnemonic.h" +#include "td/e2e/QRHandshake.h" +#include "td/e2e/TestBlockchain.h" +#include "td/e2e/Trie.h" + +#include "td/telegram/e2e_api.h" + +#include "td/utils/base64.h" +#include "td/utils/benchmark.h" +#include "td/utils/common.h" +#include "td/utils/crypto.h" +#include "td/utils/Ed25519.h" +#include "td/utils/FlatHashMap.h" +#include "td/utils/format.h" +#include "td/utils/logging.h" +#include "td/utils/overloaded.h" +#include "td/utils/Random.h" +#include "td/utils/SharedSlice.h" +#include "td/utils/simple_tests.h" +#include "td/utils/Slice.h" +#include "td/utils/SliceBuilder.h" +#include "td/utils/Status.h" +#include "td/utils/tests.h" +#include "td/utils/tl_helpers.h" +#include "td/utils/UInt.h" + +#include <map> +#include <memory> +#include <optional> +#include <utility> + +using namespace tde2e_core; + +namespace api = tde2e_api; + +template <class T> +static td::Status expect_error(td::Result<T> result) { + if (result.is_ok()) { + return td::Status::Error("Receive Ok instead of Error"); + } + return td::Status::OK(); +} + +S_TEST(MessageEncryption, simple) { + std::string secret = "secret"; + { + std::string data = "some private data"; + std::string wrong_secret = "wrong secret"; + auto encrypted_data = MessageEncryption::encrypt_data(data, secret); + LOG(ERROR) << encrypted_data.size(); + TEST_TRY_RESULT(decrypted_data, MessageEncryption::decrypt_data(encrypted_data, secret)); + TEST_ASSERT_EQ(data, decrypted_data, "decryption"); + TEST_TRY_STATUS(expect_error(MessageEncryption::decrypt_data(encrypted_data, wrong_secret))); + TEST_TRY_STATUS(expect_error(MessageEncryption::decrypt_data("", secret))); + TEST_TRY_STATUS(expect_error(MessageEncryption::decrypt_data(std::string(32, 'a'), secret))); + TEST_TRY_STATUS(expect_error(MessageEncryption::decrypt_data(std::string(33, 'a'), secret))); + TEST_TRY_STATUS(expect_error(MessageEncryption::decrypt_data(std::string(64, 'a'), secret))); + TEST_TRY_STATUS(expect_error(MessageEncryption::decrypt_data(std::string(128, 'a'), secret))); + } + + td::Random::Xorshift128plus rnd(123); + for (size_t i = 0; i < 255; i++) { + std::string data; + for (size_t j = 0; j < i; j++) { + data += static_cast<char>(rnd.fast('a', 'z')); + } + auto encrypted_data = MessageEncryption::encrypt_data(data, secret); + TEST_TRY_RESULT(decrypted_data, MessageEncryption::decrypt_data(encrypted_data, secret)); + TEST_ASSERT_EQ(data, decrypted_data, "decryption"); + } + return td::Status::OK(); +} + +struct E2eHandshakeTest { + td::Ed25519::PrivateKey alice; + td::Ed25519::PublicKey alice_public; + td::Ed25519::PrivateKey bob; + td::Ed25519::PublicKey bob_public; + + td::SecureString shared_secret; +}; + +static E2eHandshakeTest gen_test() { + auto alice = td::Ed25519::generate_private_key().move_as_ok(); + auto alice_public = alice.get_public_key().move_as_ok(); + auto bob = td::Ed25519::generate_private_key().move_as_ok(); + auto bob_public = bob.get_public_key().move_as_ok(); + auto shared_secret = td::Ed25519::compute_shared_secret(alice.get_public_key().move_as_ok(), bob).move_as_ok(); + return E2eHandshakeTest{std::move(alice), std::move(alice_public), std::move(bob), std::move(bob_public), + std::move(shared_secret)}; +} + +static void run_test(const E2eHandshakeTest &test) { + auto alice_secret = + td::Ed25519::compute_shared_secret(test.bob.get_public_key().move_as_ok(), test.alice).move_as_ok(); + auto bob_secret = td::Ed25519::compute_shared_secret(test.alice.get_public_key().move_as_ok(), test.bob).move_as_ok(); + CHECK(test.alice.get_public_key().move_as_ok().as_octet_string() == test.alice_public.as_octet_string()); + CHECK(test.bob.get_public_key().move_as_ok().as_octet_string() == test.bob_public.as_octet_string()); + CHECK(alice_secret == bob_secret); + CHECK(alice_secret == test.shared_secret); +} + +static E2eHandshakeTest pregenerated_test() { + auto alice_public_key_str = td::base64url_decode_secure("RvG0CT5i8D-CYnfhp2akVC1tPRBIw-4X6ZqNBjH-mZI").move_as_ok(); + auto alice_private_key_str = td::base64url_decode_secure("8NZGWKfRCJfiks74RG9_xHmYydarLiRsoq8VcJGPglg").move_as_ok(); + auto bob_public_key_str = td::base64url_decode_secure("1V3BGwmbo-Mwsw7QlWKN4OZFPBP9z9VhFlZKRdzTrGw").move_as_ok(); + auto bob_private_key_str = td::base64url_decode_secure("YMGoowtnZ99roUM2y5JRwiQrwGaNJ-ZRE5boy-l4aHg").move_as_ok(); + + auto alice_public_key = td::Ed25519::PublicKey(alice_public_key_str.copy()); + auto alice_private_key = td::Ed25519::PrivateKey(alice_private_key_str.copy()); + auto bob_public_key = td::Ed25519::PublicKey(bob_public_key_str.copy()); + auto bob_private_key = td::Ed25519::PrivateKey(bob_private_key_str.copy()); + auto shared_secret = td::base64url_decode_secure("CU6NsPBw59neM9crFvxKELbtKgAkI7G8tDHsb4CmyVA").move_as_ok(); + + return E2eHandshakeTest{std::move(alice_private_key), std::move(alice_public_key), std::move(bob_private_key), + std::move(bob_public_key), std::move(shared_secret)}; +} + +TEST(Handshake, InvalidKeys) { + auto private_key = td::Ed25519::generate_private_key().move_as_ok(); + auto zero_key = td::Ed25519::PublicKey(td::SecureString(32, 0)); + td::Ed25519::compute_shared_secret(zero_key, private_key).ensure_error(); +} + +TEST(Handshake, Random) { + auto test = gen_test(); + run_test(test); +} + +TEST(Handshake, Pregenerated) { + auto test = pregenerated_test(); + run_test(test); +} + +TEST(QRHandshake, Basic) { + td::int64 alice_user_id = 123; + td::int64 bob_user_id = 321; + auto alice_private_key = PrivateKey::generate().move_as_ok(); + auto bob_private_key = PrivateKey::generate().move_as_ok(); + + auto bob = QRHandshakeBob::create(bob_user_id, bob_private_key); + auto start = bob.generate_start(); // should be passed via QR + auto alice = + QRHandshakeAlice::create(alice_user_id, alice_private_key, bob_user_id, bob_private_key.to_public_key(), start) + .move_as_ok(); + auto accept = alice.generate_accept(); + auto finish = bob.receive_accept(alice_user_id, alice_private_key.to_public_key(), accept).move_as_ok(); + alice.receive_finish(finish).ensure(); +} + +TEST(CheckSharedSecret, Basic) { + auto alice = CheckSharedSecret::create(); + auto bob = CheckSharedSecret::create(); + + alice.recive_commit_nonce(bob.commit_nonce()).ensure(); + bob.recive_commit_nonce(alice.commit_nonce()).ensure(); + + alice.receive_reveal_nonce(bob.reveal_nonce().move_as_ok()).ensure(); + bob.receive_reveal_nonce(alice.reveal_nonce().move_as_ok()).ensure(); + + CHECK(alice.finalize_hash("abc").move_as_ok() == bob.finalize_hash("abc").move_as_ok()); +} + +// node_type enum + +TEST(MiniBlockchain, Basic) { + auto private_key = PrivateKey::generate().move_as_ok(); + Blockchain remote_blockchain = Blockchain::create_empty(); + Blockchain local_blockchain = Blockchain::create_empty(); + + auto block = local_blockchain.set_value(std::string(32, 'a'), "b", private_key); + remote_blockchain.try_apply_block(block, {}).ensure(); + local_blockchain.try_apply_block(block, {}).ensure(); + block = local_blockchain.set_value(std::string(32, 'b'), "c", private_key); + remote_blockchain.try_apply_block(block, {}).ensure(); + local_blockchain.try_apply_block(block, {}).ensure(); +} + +// Example usage +TEST(Tree, BitString) { + LOG(ERROR) << "BitString count: " << BitString::get_counter_value(); + td::UInt256 hash; + sha256("hello world", hash.as_mutable_slice()); + BitString s(hash.as_slice()); + for (auto l = 0; l <= 256; l++) { + for (auto r = l; r <= 256; r++) { + if (l > r) { + return; + } + auto a = s.substr(l, r - l); + BitString b; + auto str = td::serialize(a); + CHECK(str.size() % 4 == 0); + + b = BitString::fetch_from_network(str).move_as_ok(); + ASSERT_EQ(a, b); + } + } + LOG(ERROR) << "BitString count: " << BitString::get_counter_value(); +} + +TEST(Tree, SerializeStress) { + std::string value(32, 'a'); + td::Random::Xorshift128plus rnd(123); + for (size_t i = 0; i < 10000; i++) { + size_t n = rnd.fast(0, 20); + TrieRef root = TrieNode::empty_node(); + for (size_t j = 0; j < n; j++) { + td::UInt256 hash; + rnd.bytes(hash.as_mutable_slice()); + root = set(root, hash.as_slice(), td::to_string(j)).move_as_ok(); + } + + auto old_hash = root->hash; + auto s = TrieNode::serialize_for_network(root).move_as_ok(); + root = TrieNode::fetch_from_network(s).move_as_ok(); + auto new_hash = root->hash; + CHECK(old_hash == new_hash); + + auto snapshot = TrieNode::serialize_for_snapshot(root, "").move_as_ok(); + auto snapshot_root = TrieNode::fetch_from_snapshot(snapshot).move_as_ok(); + auto snapshot2 = TrieNode::serialize_for_snapshot(snapshot_root, snapshot).move_as_ok(); + CHECK(snapshot == snapshot2); + } +} + +TEST(Tree, BitStringCounter) { + CHECK(BitString::get_counter_value() == 0); + { + BitString bs(1); + size_t l = 1; + size_t r = 2; + + auto a = bs.substr(l, r - l); + auto s = td::serialize(a); + CHECK(s.size() % 4 == 0); + + BitString b = BitString::fetch_from_network(s).move_as_ok(); + ASSERT_EQ(a, b); + } + CHECK(BitString::get_counter_value() == 0); +} + +TEST(MerkleTree, Basic) { + // Build the tree + TrieRef root = TrieNode::empty_node(); + root = set(root, "apple", "fruit").move_as_ok(); + print_tree(root); + root = set(root, "application", "software").move_as_ok(); + print_tree(root); + root = set(root, "banana", "fruit").move_as_ok(); + print_tree(root); + + ASSERT_EQ("fruit", get(root, "apple").move_as_ok()); + ASSERT_EQ("software", get(root, "application").move_as_ok()); + ASSERT_EQ("fruit", get(root, "banana").move_as_ok()); + + std::vector<td::Slice> keys = {"apple", "banana"}; + TrieRef pruned_tree = generate_pruned_tree(root, keys).move_as_ok(); + print_tree(pruned_tree); + + ASSERT_EQ("fruit", get(pruned_tree, "apple").move_as_ok()); + ASSERT_EQ("fruit", get(pruned_tree, "banana").move_as_ok()); + get(pruned_tree, "application").ensure_error(); + + auto serialized = TrieNode::serialize_for_network(pruned_tree).move_as_ok(); + TrieRef pruned_tree2 = TrieNode::fetch_from_network(serialized).move_as_ok(); + print_tree(pruned_tree2); + + ASSERT_EQ("fruit", get(pruned_tree2, "apple").move_as_ok()); + ASSERT_EQ("fruit", get(pruned_tree2, "banana").move_as_ok()); + get(pruned_tree2, "application").ensure_error(); +} + +static TrieRef root; +static const int N = 1'000'000; + +TEST(Tree, BenchA) { + LOG(ERROR) << "BitString count: " << BitString::get_counter_value(); + root = TrieNode::empty_node(); + std::string value(32, 'a'); + for (int i = 0; i < N; i++) { + auto key = value + std::to_string(i); + td::UInt256 hash; + sha256(key, hash.as_mutable_slice()); + root = set(std::move(root), hash.as_slice().str(), value).move_as_ok(); + } + LOG(ERROR) << "BitString count: " << BitString::get_counter_value(); +} + +static std::string serialized_root; + +TEST(Tree, Serialize) { + serialized_root = TrieNode::serialize_for_network(root).move_as_ok(); +} + +TEST(Tree, Clear) { + root = {}; + LOG(ERROR) << "BitString count: " << BitString::get_counter_value(); +} + +TEST(Tree, Deserialize) { + root = TrieNode::fetch_from_network(serialized_root).move_as_ok(); + LOG(ERROR) << "BitString count: " << BitString::get_counter_value(); +} + +TEST(Tree, BenchAPruned) { + std::string value(32, 'a'); + size_t step = 1; + std::vector<td::UInt256> keys_str(step); + std::vector<td::Slice> keys(step); + for (size_t i = 0; i < 1000000; i += step) { + for (size_t j = 0; j < step; j++) { + auto key = value + std::to_string((i + j) % N); + sha256(key, keys_str[j].as_mutable_slice()); + keys[j] = keys_str[j].as_slice(); + } + //auto node = generate_pruned_tree(TrieNode::fetch_from_snapshot(serialized_root).move_as_ok(), keys, serialized_root).move_as_ok(); + auto node = generate_pruned_tree(root, keys, serialized_root).move_as_ok(); + auto x = TrieNode::serialize_for_network(node).move_as_ok(); + LOG_IF(ERROR, i == 0) << x.size() << " bytes serialized"; + } +} + +TEST(Tree, BenchAA) { + std::string value(32, 'a'); + for (int i = 0; i < N; i++) { + auto key = value + std::to_string(i); + td::UInt256 hash; + sha256(key, hash.as_mutable_slice()); + CHECK(value == get(root, hash.as_slice().str()).move_as_ok()); + } +} +TEST(Tree, BenchAAA) { + std::string value(32, 'a'); + for (int i = 0; i < N; i++) { + auto key = value + std::to_string(i); + td::UInt256 hash; + sha256(key, hash.as_mutable_slice()); + root = set(std::move(root), hash.as_slice().str(), value).move_as_ok(); + } +} + +static td::FlatHashMap<std::string, std::string> map; + +TEST(Tree, BenchB) { + std::string value(32, 'a'); + for (int i = 0; i < N; i++) { + auto key = value + std::to_string(i); + td::UInt256 hash; + sha256(key, hash.as_mutable_slice()); + map.emplace(hash.as_slice().str(), value); + } +} + +TEST(Tree, BenchBB) { + std::string value(32, 'a'); + for (int i = 0; i < N; i++) { + auto key = value + std::to_string(i); + td::UInt256 hash; + sha256(key, hash.as_mutable_slice()); + CHECK(value == map.find(hash.as_slice().str())->second); + } +} + +TEST(Tree, BenchBBB) { + std::string value(32, 'a'); + for (int i = 0; i < N; i++) { + auto key = value + std::to_string(i); + td::UInt256 hash; + sha256(key, hash.as_mutable_slice()); + map.emplace(hash.as_slice().str(), value); + } +} + +// TODO: to we need both secret and local password?.. +// user_password is known only to user and never stored +// secret used to decrypt the key and should be stored somewhere safe +// encrypted_data could be stored anywhere +// +// Is it too complicated? Is it necessary to encrypt +static td::Result<EncryptedKey> create_new_encrypted_key(td::Slice user_password) { + TRY_RESULT(mnemonic, Mnemonic::create_new()); + auto private_key = mnemonic.to_private_key(); + auto decrypted_key = DecryptedKey(mnemonic.get_words(), std::move(private_key)); + return decrypted_key.encrypt(user_password); +} + +static td::Result<EncryptedKey> change_user_password(const EncryptedKey &encrypted_key, td::Slice user_password, + td::Slice new_user_password) { + TRY_RESULT(decrypted_key, encrypted_key.decrypt(user_password, false)); + return decrypted_key.encrypt(new_user_password); +} + +static td::Result<td::SecureString> export_mnemonic(const EncryptedKey &encrypted_key, td::Slice user_password) { + TRY_RESULT(decrypted_key, encrypted_key.decrypt(user_password, false)); + CHECK(decrypted_key.mnemonic_words.size() == 24); + size_t length = decrypted_key.mnemonic_words.size() - 1; + for (auto &word : decrypted_key.mnemonic_words) { + length += word.size(); + } + td::SecureString res(length); + auto dest = res.as_mutable_slice(); + bool is_first = true; + for (auto &word : decrypted_key.mnemonic_words) { + if (!is_first) { + dest[0] = ' '; + dest.remove_prefix(1); + } else { + is_first = false; + } + dest.copy_from(word); + dest.remove_prefix(word.size()); + } + return res; +} + +static td::Result<EncryptedKey> import_mnemonic(td::Slice mnemonic_words, td::Slice user_password) { + TRY_RESULT(mnemonic, Mnemonic::create(td::SecureString(mnemonic_words), td::SecureString())); + auto decrypted_key = DecryptedKey(mnemonic); + return decrypted_key.encrypt(user_password); +} + +TEST(E2E, GenerateKeys) { + // generate key + auto encrypted_key = create_new_encrypted_key("user_password").move_as_ok(); + change_user_password(encrypted_key, "bad_user_password", "user_password").ensure_error(); + auto new_encrypted_key = change_user_password(encrypted_key, "user_password", "new_password").move_as_ok(); + export_mnemonic(new_encrypted_key, "user_password").ensure_error(); + auto mnemonic = export_mnemonic(new_encrypted_key, "new_password").move_as_ok(); + auto other_encrypted_key = import_mnemonic(mnemonic, "new_password").move_as_ok(); + CHECK(encrypted_key.o_public_key == new_encrypted_key.o_public_key); +} + +TEST(E2E_API, Key) { + using namespace tde2e_api; + auto alice_pk = key_generate_private_key().value(); + auto bob_pk = key_generate_private_key().value(); + auto carol_pk = key_generate_private_key().value(); + + auto secret = key_from_bytes("secret").value(); + auto bad_secret = key_from_bytes("bad_secret").value(); + + auto encrypted_alice_pk = key_to_encrypted_private_key(alice_pk, secret).value(); + key_from_encrypted_private_key(encrypted_alice_pk, bad_secret).error(); + auto alice_pk_copy = key_from_encrypted_private_key(encrypted_alice_pk, secret).value(); + + ASSERT_EQ(key_to_public_key(alice_pk).value(), key_to_public_key(alice_pk_copy).value()); + auto alice_PK = key_from_public_key(key_to_public_key(alice_pk).value()).value(); + auto bob_PK = key_from_public_key(key_to_public_key(bob_pk).value()).value(); + auto carol_PK = key_from_public_key(key_to_public_key(carol_pk).value()).value(); + + key_destroy(alice_pk_copy).value(); + key_to_public_key(alice_pk_copy).error(); + + auto words = key_to_words(alice_pk).value(); + ASSERT_EQ(alice_pk, key_from_words(std::move(words)).value()); + + auto shared_key_ab = key_from_ecdh(alice_pk, bob_PK).value(); + auto shared_key_ba = key_from_ecdh(bob_pk, alice_PK).value(); + auto shared_key_ac = key_from_ecdh(alice_pk, carol_PK).value(); + + auto encrypted = encrypt_message_for_many({shared_key_ab, shared_key_ac}, "very secret message").value(); + ASSERT_EQ( + "very secret message", + decrypt_message_for_many(shared_key_ba, encrypted.encrypted_headers[0], encrypted.encrypted_message).value()); + decrypt_message_for_many(shared_key_ac, encrypted.encrypted_headers[0], encrypted.encrypted_message).error(); + + auto encrypted2 = encrypt_message_for_one(shared_key_ab, "very secret message").value(); + ASSERT_EQ("very secret message", decrypt_message_for_one(shared_key_ba, encrypted2).value()); + decrypt_message_for_one(shared_key_ac, encrypted2).error(); + key_destroy_all(); +} + +TEST(E2E_API, HandshakeVerify) { + using namespace tde2e_api; + auto bob_id = 123; + auto alice_id = 321; + auto alice_pk = key_generate_private_key().value(); + auto bob_pk = key_generate_private_key().value(); + + // Bob creates handshake + auto bob_handshake_id = handshake_create_for_bob(bob_id, bob_pk).value(); + // Start is shown on QR + auto start = handshake_bob_send_start(bob_handshake_id).value(); + + // Alice received qr, received information about qr from server and create handshake + auto alice_handshake_id = + handshake_create_for_alice(alice_id, alice_pk, bob_id, key_to_public_key(bob_pk).value(), start).value(); + auto accept = handshake_alice_send_accept(alice_handshake_id).value(); + // Alice knows shared key. She knows that it is known to author of QR code, but not necessary to bob_pk owner. + auto shared_a = handshake_get_shared_key_id(alice_handshake_id).value(); + + // Bob receives accept and generates finish + auto finish = + handshake_bob_receive_accept_send_finish(bob_handshake_id, alice_id, key_to_public_key(alice_pk).value(), accept) + .value(); + // At this point Bob "verified" Alice + // Bob knows shared key + auto shared_b = handshake_get_shared_key_id(bob_handshake_id).value(); + + // Alice receives and verifies finish + handshake_alice_receive_finish(alice_handshake_id, finish).value(); + // At this point Alice "verified" Bob + + ASSERT_EQ(shared_a, shared_b); + handshake_destroy_all(); +} + +TEST(E2E_API, HandshakeLogin) { + using namespace tde2e_api; + + auto alice_id = 321; + auto alice_pk = key_generate_private_key().value(); + + auto bob_login_id = login_create_for_bob().value(); + auto start = login_bob_send_start(bob_login_id).value(); + auto alice_data = login_create_for_alice(alice_id, alice_pk, start).value(); + auto received_alice_pk = + login_finish_for_bob(bob_login_id, alice_id, key_to_public_key(alice_pk).value(), alice_data).value(); + ASSERT_EQ(key_to_public_key(alice_pk).value(), key_to_public_key(received_alice_pk).value()); + login_destroy_all(); +} + +TEST(Container, Basic) { + using namespace tde2e_api; + Container<TypeInfo<int, false, false>, TypeInfo<std::string, false, true>, TypeInfo<std::vector<int>, true, false>, + TypeInfo<std::vector<std::string>, true, true>> + container; + + auto id_int = container.emplace<int>(1); + td::UInt256 hash; + hash.as_mutable_slice().fill(7); + auto id_string = container.try_emplace<std::string>(hash, "hello"); + auto id_string_2 = container + .try_build<std::string>(hash, + []() -> td::Result<std::string> { + UNREACHABLE(); + return "..."; + }) + .move_as_ok(); + ASSERT_EQ(id_string, id_string_2); + auto id_vec_int = container + .try_build<std::vector<int>>({}, + []() -> td::Result<std::vector<int>> { + return std::vector<int>{1, 2, 3, 4}; + }) + .move_as_ok(); + auto id_vec_string = container.emplace<std::vector<std::string>>(std::vector<std::string>{"a", "b", "c"}); + + container.get_shared<int>(id_int).ensure(); + container.get_shared<std::string>(id_string).ensure(); + container.get_unique<std::vector<int>>(id_vec_int).ensure(); + container.get_unique<std::vector<std::string>>(id_vec_string).ensure(); + + container.get_shared<int>(id_string).ensure_error(); + container.get_shared<std::string>(id_int).ensure_error(); + container.get_unique<std::vector<int>>(id_vec_string).ensure_error(); + container.get_unique<std::vector<std::string>>(id_vec_int).ensure_error(); +} + +struct BaselineBlockchainState { + std::map<std::string, std::string> key_value_state; + GroupStateRef group_state; + GroupSharedKeyRef shared_key; + td::int32 height{-1}; + std::string get_value(const std::string &key) const { + auto it = key_value_state.find(key); + if (it == key_value_state.end()) { + return ""; + } + return it->second; + } + void apply_changes(const std::vector<Change> &changes) { + for (const auto &change_v : changes) { + std::visit(td::overloaded([&](const ChangeNoop &) {}, + [&](const ChangeSetValue &change) { key_value_state[change.key] = change.value; }, + [&](const ChangeSetGroupState &change) { group_state = change.group_state; }, + [&](const ChangeSetSharedKey &change) { shared_key = change.shared_key; }), + change_v.value); + } + height++; + } +}; + +S_TEST(E2E_Blockchain, Base) { + TEST_TRY_RESULT(pk, PrivateKey::generate()); + TEST_TRY_RESULT(pk2, PrivateKey::generate()); + + BlockchainTester tester; + + auto to_hash = [](td::Slice key) { + std::string res(32, 0); + td::sha256(key, res); + return res; + }; + + auto a = to_hash("a"); + auto b = to_hash("b"); + + TEST_ASSERT_EQ("", tester.get_value(a), "empty blockchain"); + TEST_ASSERT_EQ("", tester.get_value(b), "empty blockchain"); + using BB = BlockBuilder; + + TEST_TRY_STATUS( + tester.expect_ok({BB::make_set_value(a, "hello a"), + BB::make_group_change({{2, GroupParticipantFlags::AllPermissions, pk2.to_public_key()}})}, + pk2)); + TEST_ASSERT_EQ("hello a", tester.get_value(a), "hello a"); + TEST_TRY_STATUS(tester.expect_error(E::Any, {BB::make_set_value(a, "hello b")}, pk)); + TEST_TRY_STATUS(tester.expect_ok({BB::make_set_value(a, "hello b")}, pk2)); + TEST_ASSERT_EQ("hello b", tester.get_value(a), "..."); + tester.reindex(); + TEST_ASSERT_EQ("hello b", tester.get_value(a), "..."); + return td::Status::OK(); +} + +S_TEST(E2E_Blockchain, Stress) { + BlockchainTester tester; + TEST_TRY_RESULT(pk, PrivateKey::generate()); + + td::Random::Xorshift128plus rnd(123); + auto gen_string = [&](auto from, auto to, auto size) { + std::string s(size, 0); + for (auto &c : s) { + c = static_cast<char>(rnd.fast(from, to)); + } + return s; + }; + + auto to_hash = [](td::Slice key) { + std::string res(32, 0); + td::sha256(key, res); + return res; + }; + auto gen_key = [&] { + auto len = rnd.fast(1, 15); + return to_hash(gen_string('a', 'b', len)); + }; + + auto gen_value = [&] { + std::string res(rnd.fast(1, 64), 0); + rnd.bytes(res); + return res; + }; + + auto gen_query = [&] { + std::vector<std::string> keys(rnd.fast(1, 1)); + for (auto &key : keys) { + key = gen_key(); + }; + return keys; + }; + + auto gen_changes = [&] { + auto n = rnd.fast(1, 2); + std::vector<Change> changes; + changes.reserve(n); + for (int i = 0; i < n; i++) { + changes.push_back(BlockBuilder::make_set_value(gen_key(), gen_value())); + } + return changes; + }; + + auto run_get = [&] { + auto keys = gen_query(); + TEST_TRY_STATUS(tester.get_values(keys)); + return td::Status::OK(); + }; + + auto run_set = [&] { + auto changes = gen_changes(); + TEST_TRY_STATUS(tester.apply(changes, pk)); + return td::Status::OK(); + }; + + auto reindex = [&] { + tester.reindex(); + return td::Status::OK(); + }; + + td::RandomSteps steps{{{run_set, 10}, {run_get, 100}, {reindex, 1}}}; + for (size_t i = 0; i < 10000; i++) { + steps.step(rnd); + } + return td::Status::OK(); +} + +using namespace tde2e_api; +S_TEST(E2E_Blockchain, Call) { + SET_VERBOSITY_LEVEL(3); + CallTester ct; + TEST_TRY_STATUS(ct.start_call({0, 1, 2})); + TEST_TRY_STATUS(ct.check_shared_key()); + TEST_TRY_STATUS(ct.check_emoji_hash()); + TEST_TRY_STATUS(ct.update_call(0, {0, 3, 4, 5})); + TEST_TRY_STATUS(ct.check_shared_key()); + TEST_TRY_STATUS(ct.check_emoji_hash()); + return td::Status::OK(); +} + +TEST(Call, Basic_API) { + using namespace tde2e_api; + auto F = [](Result<std::string> block) -> Result<std::string> { + if (block.is_ok()) { + return Blockchain::from_local_to_server(block.value()); + } + return block; + }; + + auto key0 = key_generate_temporary_private_key().value(); + auto pkey0 = key_from_public_key(key_to_public_key(key0).value()).value(); + auto key1 = key_generate_temporary_private_key().value(); + auto pkey1 = key_from_public_key(key_to_public_key(key1).value()).value(); + auto key2 = key_generate_temporary_private_key().value(); + auto pkey2 = key_from_public_key(key_to_public_key(key2).value()).value(); + auto key3 = key_generate_temporary_private_key().value(); + auto pkey3 = key_from_public_key(key_to_public_key(key3).value()).value(); + + auto zero_block = F(call_create_zero_block(key0, CallState{0, {CallParticipant{-1, pkey0, 3}}})).value(); + + auto call1 = call_create(-1, key0, zero_block).value(); + auto block0 = F(call_create_self_add_block(key1, zero_block, CallParticipant{1, pkey1, 3})).value(); + call1 = call_create(1, key1, block0).value(); + + auto block1 = F(call_create_self_add_block(key2, block0, CallParticipant{2, pkey2, 3})).value(); + call_apply_block(call1, block1).value(); + auto call2 = call_create(2, key2, block1).value(); + ASSERT_EQ(call_get_verification_words(call2).value().words, call_get_verification_words(call1).value().words); + + auto block2 = F(call_create_change_state_block( + call2, CallState{0, {CallParticipant{2, pkey2, 3}, CallParticipant{3, pkey3, 3}}})) + .value(); + call_describe_block(block2).value(); + auto call3 = call_create(3, key3, block2).value(); + + call_apply_block(call2, block2).value(); + CHECK(!call_apply_block(call1, block2).is_ok()); + + // call2 and call3 verification + ASSERT_EQ(call_get_verification_words(call2).value().words, call_get_verification_words(call3).value().words); + + auto block31 = F(call_create_change_state_block( + call2, CallState{0, {CallParticipant{2, pkey2, 3}, CallParticipant{3, pkey3, 3}}})) + .value(); + + call_apply_block(call2, block31).value(); + auto commit2 = F(call_pull_outbound_messages(call2).value().at(0)).value(); + + call_describe_message(commit2).value(); + + call_receive_inbound_message(call2, commit2).value(); + call_receive_inbound_message(call3, commit2).value(); + + call_apply_block(call3, block31).value(); + auto commit3 = F(call_pull_outbound_messages(call3).value().at(0)).value(); + + CHECK(commit2 != commit3); + call_receive_inbound_message(call2, commit3).value(); + call_receive_inbound_message(call3, commit3).value(); + + auto reveal2 = F(call_pull_outbound_messages(call2).value().at(0)).value(); + auto reveal3 = F(call_pull_outbound_messages(call3).value().at(0)).value(); + call_receive_inbound_message(call2, reveal2).value(); + call_receive_inbound_message(call2, reveal3).value(); + call_receive_inbound_message(call3, reveal2).value(); + call_receive_inbound_message(call3, reveal3).value(); + + ASSERT_EQ(call_get_verification_state(call2).value().emoji_hash.value(), + call_get_verification_state(call3).value().emoji_hash.value()); + + auto e = call_encrypt(call2, 1, "hello").value(); + auto e2 = call_encrypt(call2, 1, "hello").value(); + CHECK(e != "hello"); + LOG(ERROR) << e.size(); + ASSERT_TRUE(!call_decrypt(call2, 2, 1, e).is_ok()); + // ASSERT_TRUE(!call_decrypt(call3, 2, 2, e).is_ok()); Uncomment if we will validate channel_id + ASSERT_TRUE(!call_decrypt(call3, 1, 1, e).is_ok()); + ASSERT_EQ("hello", call_decrypt(call3, 2, 1, e).value()); + ASSERT_TRUE(!call_decrypt(call3, 2, 1, e).is_ok()); + + auto block3 = F(call_create_change_state_block( + call2, CallState{0, {CallParticipant{2, pkey2, 3}, CallParticipant{3, pkey3, 3}}})) + .value(); + call_apply_block(call3, block3).value(); + ASSERT_TRUE(!call_decrypt(call3, 2, 1, e).is_ok()); + ASSERT_EQ("hello", call_decrypt(call3, 2, 1, e2).value()); + ASSERT_TRUE(call_decrypt(call2, 3, 1, call_encrypt(call3, 1, "bye").value()).is_ok()); + + LOG(ERROR) << call_describe(call1).value(); + LOG(ERROR) << call_describe(call2).value(); + + key_destroy_all(); + call_destroy_all(); +} + +TEST(State, Basic) { + SET_VERBOSITY_LEVEL(3); + using namespace tde2e_api; + + ServerBlockchain kv_server; + auto pk = key_generate_private_key().value(); + auto storage = storage_create(pk, {}).value(); + + auto contact_pk = key_generate_private_key().value(); + auto contact_public_key = key_from_public_key(key_to_public_key(contact_pk).value()).value(); + + storage_get_contact(storage, contact_public_key); + storage_get_contact(storage, contact_public_key).error(); + Entry<Name> entry_name; + entry_name.value = Name{"A", "B"}; + auto signed_entry_name = storage_sign_entry(contact_pk, entry_name).value(); + auto update_id = storage_update_contact(storage, contact_public_key, signed_entry_name).value(); + (void)update_id; + + auto load_proofs = [&] { + auto keys = storage_get_blockchain_state(storage).value().required_proofs; + auto proof = kv_server.get_proof(storage_blockchain_height(storage).value(), keys).move_as_ok(); + storage_blockchain_add_proof(storage, proof, keys).value(); + }; + auto update_blockchain = [&] { + auto block = storage_get_blockchain_state(storage).value().next_suggested_block; + if (block.empty()) { + return; + } + kv_server.try_apply_block(block).ensure(); + storage_blockchain_apply_block(storage, block).value(); + }; + + load_proofs(); + + Value value; + value.o_name = entry_name; + + ASSERT_EQ(std::optional<Value>(), storage_get_contact(storage, contact_public_key).value()); + ASSERT_EQ(value, storage_get_contact_optimistic(storage, contact_public_key).value()); + + update_blockchain(); + ASSERT_EQ(value, storage_get_contact(storage, contact_public_key).value()); +} + +class CallEncryptionBench final : public td::Benchmark { + public: + explicit CallEncryptionBench(size_t msg_size) : msg_size_(msg_size) { + msg_ = std::string(msg_size_, '\1'); + } + std::string get_description() const final { + return PSTRING() << "Call encrypt/decrypt msg_size=" << td::format::as_size(msg_size_); + } + + void start_up() final { + auto pk1 = PrivateKey::generate().move_as_ok(); + auto pk2 = PrivateKey::generate().move_as_ok(); + auto group = std::make_shared<GroupState>( + GroupState{{GroupParticipant{1, 0, pk1.to_public_key()}, GroupParticipant{2, 0, pk2.to_public_key()}}}); + td::SecureString shared_key(32, '\0'); + e1_ = std::make_unique<CallEncryption>(1, pk1); + e1_->add_shared_key(1, td::UInt256{}, shared_key.copy(), group); + e2_ = std::make_unique<CallEncryption>(2, pk2); + e2_->add_shared_key(1, td::UInt256{}, shared_key.copy(), group); + } + + void run(int n) final { + for (int i = 0; i < n; i++) { + auto encrypted = e1_->encrypt(1, msg_).move_as_ok(); + CHECK(msg_ == e2_->decrypt(1, 1, encrypted).move_as_ok()); + } + } + + private: + size_t msg_size_{}; + std::string msg_; + std::unique_ptr<CallEncryption> e1_; + std::unique_ptr<CallEncryption> e2_; +}; + +TEST(Call, Bench) { + td::bench(CallEncryptionBench(16)); + td::bench(CallEncryptionBench(1024)); + td::bench(CallEncryptionBench(16 * 1024)); + td::bench(CallEncryptionBench(64 * 1024)); +} + +TEST(Keys, Sanity) { + auto pk = PrivateKey::generate().move_as_ok(); + auto hello_sign = pk.sign("hello").move_as_ok(); + pk.to_public_key().verify("hello", hello_sign).ensure(); + auto bad_sign = hello_sign.to_u512(); + bad_sign.raw[0]++; + pk.to_public_key().verify("hello", Signature::from_u512(bad_sign)).ensure_error(); +} + +#if TG_ENGINE +int main() { + td::TestsRunner::get_default().run_all(); + _Exit(0); +} +#endif |