tdutils: aes experiments

GitOrigin-RevId: 1dd2812041778679fb1ef0462d9e8b1899181e0b

benchmark/bench_crypto.cpp

@@ -48,6 +48,36 @@ class SHA1Bench : public td::Benchmark {
   }
 };
 
+class AesBench : public td::Benchmark {
+ public:
+  alignas(64) unsigned char data[DATA_SIZE];
+  td::UInt256 key;
+  td::UInt256 iv;
+
+  std::string get_description() const override {
+    return PSTRING() << "AES ECB OpenSSL [" << (DATA_SIZE >> 10) << "KB]";
+  }
+
+  void start_up() override {
+    for (int i = 0; i < DATA_SIZE; i++) {
+      data[i] = 123;
+    }
+    td::Random::secure_bytes(key.raw, sizeof(key));
+    td::Random::secure_bytes(iv.raw, sizeof(iv));
+  }
+
+  void run(int n) override {
+    td::AesState state;
+    state.init(td::as_slice(key), true);
+    td::MutableSlice data_slice(data, DATA_SIZE);
+    for (int i = 0; i <= n; i++) {
+      for (size_t offset = 0; offset + 16 <= data_slice.size(); offset += 16) {
+        state.encrypt(data_slice.ubegin() + offset, data_slice.ubegin() + offset);
+      }
+    }
+  }
+};
+
 class AESBench : public td::Benchmark {
  public:
   alignas(64) unsigned char data[DATA_SIZE];
@@ -197,6 +227,8 @@ class Crc64Bench : public td::Benchmark {
 
 int main() {
   td::init_openssl_threads();
+  td::bench(AesBench());
+  td::bench(AESBench());
 
   td::bench(Pbkdf2Bench());
   td::bench(RandBench());
@@ -208,7 +240,6 @@ int main() {
 #endif
   td::bench(SslRandBufBench());
   td::bench(SHA1Bench());
-  td::bench(AESBench());
   td::bench(Crc32Bench());
   td::bench(Crc64Bench());
   return 0;

tdutils/td/utils/Random.cpp

@@ -164,4 +164,18 @@ int Random::Xorshift128plus::fast(int min, int max) {
   return static_cast<int>((*this)() % (max - min + 1) + min);
 }
 
+void Random::Xorshift128plus::bytes(MutableSlice dest) {
+  int cnt = 0;
+  td::uint64 buf = 0;
+  for (auto &c : dest) {
+    if (cnt == 0) {
+      buf = operator()();
+      cnt = 8;
+    }
+    cnt--;
+    c = buf & 255;
+    buf >>= 8;
+  }
+}
+
 }  // namespace td

tdutils/td/utils/Random.h

@@ -37,6 +37,7 @@ class Random {
     Xorshift128plus(uint64 seed_a, uint64 seed_b);
     uint64 operator()();
     int fast(int min, int max);
+    void bytes(MutableSlice dest);
 
    private:
     uint64 seed_[2];

tdutils/td/utils/crypto.cpp

@@ -249,6 +249,56 @@ int pq_factorize(Slice pq_str, string *p_str, string *q_str) {
   return 0;
 }
 
+class AesState::Impl {
+ public:
+  EVP_CIPHER_CTX *ctx{nullptr};
+  AES_KEY key;
+  bool encrypt;
+
+  ~Impl() {
+    if (ctx != nullptr) {
+      EVP_CIPHER_CTX_free(ctx);
+    }
+  }
+};
+
+AesState::AesState() = default;
+AesState::~AesState() = default;
+
+void AesState::init(Slice key, bool encrypt) {
+  CHECK(key.size() == 32);
+  impl_ = make_unique<Impl>();
+  impl_->ctx = EVP_CIPHER_CTX_new();
+  CHECK(impl_->ctx);
+
+  if (encrypt) {
+    CHECK(1 == EVP_EncryptInit_ex(impl_->ctx, EVP_aes_256_ecb(), nullptr, key.ubegin(), nullptr));
+    AES_set_encrypt_key(key.ubegin(), 256, &impl_->key);
+  } else {
+    CHECK(1 == EVP_DecryptInit_ex(impl_->ctx, EVP_aes_256_ecb(), nullptr, key.ubegin(), nullptr));
+    AES_set_decrypt_key(key.ubegin(), 256, &impl_->key);
+  }
+  EVP_CIPHER_CTX_set_padding(impl_->ctx, 0);
+  impl_->encrypt = encrypt;
+}
+
+void AesState::encrypt(const uint8 *src, uint8 *dst) {
+  CHECK(impl_->encrypt);
+  CHECK(impl_->ctx);
+  int len;
+  CHECK(1 == EVP_EncryptUpdate(impl_->ctx, dst, &len, src, 16));
+  CHECK(len == 16);
+  //AES_encrypt(src, dst, &impl_->key);
+}
+void AesState::decrypt(const uint8 *src, uint8 *dst) {
+  CHECK(!impl_->encrypt);
+  CHECK(impl_->ctx);
+  int len;
+  CHECK(1 == EVP_DecryptUpdate(impl_->ctx, dst, &len, src, 16));
+  LOG_CHECK(len == 16) << len;
+  //AES_decrypt(src, dst, &impl_->key);
+}
+
 static void aes_ige_xcrypt(Slice aes_key, MutableSlice aes_iv, Slice from, MutableSlice to, bool encrypt_flag) {
   CHECK(aes_key.size() == 32);
   CHECK(aes_iv.size() == 32);

tdutils/td/utils/crypto.h

@@ -19,6 +19,23 @@ uint64 pq_factorize(uint64 pq);
 #if TD_HAVE_OPENSSL
 void init_crypto();
 
+struct AesState {
+ public:
+  AesState();
+  AesState(const AesState &from) = delete;
+  AesState &operator=(const AesState &from) = delete;
+  AesState(AesState &&from);
+  AesState &operator=(AesState &&from);
+  ~AesState();
+  void init(Slice key, bool encrypt);
+  void encrypt(const uint8 *src, uint8 *dst);
+  void decrypt(const uint8 *src, uint8 *dst);
+
+ private:
+  class Impl;
+  unique_ptr<Impl> impl_;
+};
+
 int pq_factorize(Slice pq_str, string *p_str, string *q_str);
 
 void aes_ige_encrypt(Slice aes_key, MutableSlice aes_iv, Slice from, MutableSlice to);

tdutils/test/crypto.cpp

@@ -18,6 +18,28 @@
 static td::vector<td::string> strings{"", "1", "short test string", td::string(1000000, 'a')};
 
 #if TD_HAVE_OPENSSL
+TEST(Crypto, Aes) {
+  td::Random::Xorshift128plus rnd(123);
+  td::UInt256 key;
+  rnd.bytes(as_slice(key));
+  std::string plaintext(16, 0);
+  std::string encrypted(16, 0);
+  std::string decrypted(16, 0);
+  rnd.bytes(plaintext);
+
+  td::AesState encryptor;
+  encryptor.init(as_slice(key), true);
+  td::AesState decryptor;
+  decryptor.init(as_slice(key), false);
+
+  encryptor.encrypt(td::as_slice(plaintext).ubegin(), td::as_slice(encrypted).ubegin());
+  decryptor.decrypt(td::as_slice(encrypted).ubegin(), td::as_slice(decrypted).ubegin());
+
+  CHECK(decrypted == plaintext);
+  CHECK(decrypted != encrypted);
+  CHECK(td::crc32(encrypted) == 178892237);
+}
+
 TEST(Crypto, AesCtrState) {
   td::vector<td::uint32> answers1{0u,         1141589763u, 596296607u,  3673001485u, 2302125528u,
                                   330967191u, 2047392231u, 3537459563u, 307747798u,  2149598133u};