1 files changed, 208 insertions, 0 deletions
diff --git a/tde2e/td/e2e/Keys.cpp b/tde2e/td/e2e/Keys.cpp
new file mode 100644
index 000000000..e7aafd289
--- /dev/null
+++ b/tde2e/td/e2e/Keys.cpp
@@ -0,0 +1,208 @@
+//
+// 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/Keys.h"
+
+#include "td/e2e/MessageEncryption.h"
+
+#include "td/utils/Ed25519.h"
+#include "td/utils/misc.h"
+#include "td/utils/SliceBuilder.h"
+
+namespace tde2e_core {
+
+struct PublicKeyRaw {
+  td::Ed25519::PublicKey public_key;
+};
+
+struct PrivateKeyRaw {
+  PublicKeyRaw public_key;
+  td::Ed25519::PrivateKey private_key;
+  std::shared_ptr<const td::Ed25519::PreparedPrivateKey> prepared_private_key;
+};
+
+struct PrivateKeyWithMnemonicRaw {
+  std::vector<td::SecureString> mnemonic;
+  PrivateKeyRaw key_pair;
+};
+
+Signature Signature::from_u512(const td::UInt512 &signature) {
+  return Signature{signature};
+}
+
+td::UInt512 Signature::to_u512() const {
+  return signature_;
+}
+
+td::Result<Signature> Signature::from_slice(const td::Slice &slice) {
+  td::UInt512 signature;
+  if (slice.size() != 64) {
+    return td::Status::Error(PSLICE() << "Invalid signature length: " << slice.size());
+  }
+  signature.as_mutable_slice().copy_from(slice);
+  return Signature{signature};
+}
+
+td::Slice Signature::to_slice() const {
+  return signature_.as_slice();
+}
+
+auto empty_public_key() {
+  static auto pk = PublicKey::from_u256({});
+  return pk;
+}
+
+PublicKey::PublicKey() : raw_(empty_public_key().raw_) {
+}
+
+PublicKey::PublicKey(std::shared_ptr<const PublicKeyRaw> public_key) : raw_(std::move(public_key)) {
+  CHECK(raw_);
+}
+
+td::Result<PublicKey> PublicKey::from_slice(td::Slice slice) {
+  if (slice.size() != td::Ed25519::PublicKey::LENGTH) {
+    return td::Status::Error("Invalid length of public key");
+  }
+  PublicKeyRaw public_key_raw{td::Ed25519::PublicKey(td::SecureString(slice))};
+  return PublicKey(std::make_shared<PublicKeyRaw>(std::move(public_key_raw)));
+}
+
+PublicKey PublicKey::from_u256(const td::UInt256 &public_key) {
+  PublicKeyRaw public_key_raw{td::Ed25519::PublicKey(td::SecureString(public_key.as_slice()))};
+  return PublicKey(std::make_shared<PublicKeyRaw>(std::move(public_key_raw)));
+}
+
+td::UInt256 PublicKey::to_u256() const {
+  CHECK(raw_);
+  td::UInt256 result;
+  result.as_mutable_slice().copy_from(raw_->public_key.as_octet_string());
+  return result;
+}
+
+td::Status PublicKey::verify(td::Slice data, const Signature &signature) const {
+  CHECK(raw_);
+  return raw_->public_key.verify_signature(data, signature.to_slice());
+}
+
+td::SecureString PublicKey::to_secure_string() const {
+  return raw_->public_key.as_octet_string();
+}
+
+bool PublicKey::operator==(const PublicKey &other) const {
+  return to_u256() == other.to_u256();
+}
+
+bool PublicKey::operator!=(const PublicKey &other) const {
+  return !(*this == other);
+}
+
+bool PublicKey::operator<(const PublicKey &other) const {
+  return to_u256() < other.to_u256();
+}
+
+auto empty_private_key() {
+  static auto pk = PrivateKey::from_slice(std::string(32, 1)).move_as_ok();
+  return pk;
+}
+
+PrivateKey::PrivateKey() : raw_(empty_private_key().raw_) {
+}
+
+PrivateKey::PrivateKey(std::shared_ptr<const PrivateKeyRaw> key_pair) : raw_(std::move(key_pair)) {
+  CHECK(raw_);
+}
+
+td::Result<PrivateKey> PrivateKey::generate() {
+  TRY_RESULT(private_key, td::Ed25519::generate_private_key());
+  TRY_RESULT(public_key, private_key.get_public_key());
+  TRY_RESULT(prepared_private_key, private_key.prepare());
+  return std::make_shared<PrivateKeyRaw>(
+      PrivateKeyRaw{{std::move(public_key)}, std::move(private_key), std::move(prepared_private_key)});
+}
+
+td::Result<PrivateKey> PrivateKey::from_slice(const td::Slice &slice) {
+  if (slice.size() != td::Ed25519::PublicKey::LENGTH) {
+    return td::Status::Error("Invalid private key length");
+  }
+  auto private_key = td::Ed25519::PrivateKey(td::SecureString(slice));
+  TRY_RESULT(public_key, private_key.get_public_key());
+  TRY_RESULT(prepared_private_key, private_key.prepare());
+  return std::make_shared<PrivateKeyRaw>(
+      PrivateKeyRaw{{std::move(public_key)}, std::move(private_key), std::move(prepared_private_key)});
+}
+
+td::Result<td::SecureString> PrivateKey::compute_shared_secret(const PublicKey &public_key) const {
+  TRY_RESULT(x25519_shared_secret, td::Ed25519::compute_shared_secret(public_key.raw().public_key, raw_->private_key));
+  return td::SecureString(
+      MessageEncryption::hmac_sha512("tde2e_shared_secret", x25519_shared_secret).as_slice().substr(0, 32));
+}
+
+td::Result<Signature> PrivateKey::sign(const td::Slice &data) const {
+  CHECK(raw_);
+  TRY_RESULT(signature, td::Ed25519::PrivateKey::sign(*raw_->prepared_private_key, data));
+  //TRY_RESULT(signature, raw_->private_key.sign(data));
+  return Signature::from_slice(signature);
+}
+
+PublicKey PrivateKey::to_public_key() const {
+  CHECK(raw_);
+  return PublicKey(std::shared_ptr<const PublicKeyRaw>(raw_, &raw_->public_key));
+}
+
+td::SecureString PrivateKey::to_secure_string() const {
+  return raw_->private_key.as_octet_string();
+}
+
+PrivateKeyWithMnemonic::PrivateKeyWithMnemonic(std::shared_ptr<const PrivateKeyWithMnemonicRaw> raw)
+    : raw_(std::move(raw)) {
+  CHECK(raw_);
+}
+
+PrivateKeyWithMnemonic PrivateKeyWithMnemonic::from_private_key(const PrivateKey &private_key,
+                                                                std::vector<td::SecureString> words) {
+  return PrivateKeyWithMnemonic(std::make_shared<PrivateKeyWithMnemonicRaw>(PrivateKeyWithMnemonicRaw{
+      std::move(words),
+      PrivateKeyRaw{{td::Ed25519::PublicKey(td::SecureString(private_key.to_public_key().to_u256().as_slice()))},
+                    td::Ed25519::PrivateKey(private_key.to_secure_string()),
+                    private_key.raw().prepared_private_key}}));
+}
+
+PrivateKey PrivateKeyWithMnemonic::to_private_key() const {
+  return PrivateKey(std::shared_ptr<const PrivateKeyRaw>(raw_, &raw_->key_pair));
+}
+
+td::Span<td::SecureString> PrivateKeyWithMnemonic::words() const {
+  CHECK(raw_);
+  return raw_->mnemonic;
+}
+
+PublicKey PrivateKeyWithMnemonic::to_public_key() const {
+  return to_private_key().to_public_key();
+}
+
+td::Result<Signature> PrivateKeyWithMnemonic::sign(const td::Slice &data) const {
+  return to_private_key().sign(data);
+}
+
+td::StringBuilder &operator<<(td::StringBuilder &sb, const PrivateKeyWithMnemonic &key_pair_with_mnemonic) {
+  return sb << "EdPrivateKey(pub="
+            << td::hex_encode(key_pair_with_mnemonic.to_public_key().to_u256().as_slice().substr(0, 8)) << "...)";
+}
+
+td::StringBuilder &operator<<(td::StringBuilder &sb, const PrivateKey &key_pair) {
+  return sb << "EdPrivateKey(pub=" << td::hex_encode(key_pair.to_public_key().to_u256().as_slice().substr(0, 8))
+            << "...)";
+}
+
+td::StringBuilder &operator<<(td::StringBuilder &sb, const PublicKey &public_key) {
+  return sb << "EdPublicKey(" << td::hex_encode(public_key.to_u256().as_slice().substr(0, 8)) << "...)";
+}
+
+td::StringBuilder &operator<<(td::StringBuilder &sb, const Signature &signature) {
+  return sb << "Signature(" << td::hex_encode(signature.signature_.as_slice().substr(0, 8)) << "...)";
+}
+
+}  // namespace tde2e_core