From 132caf5c8f518272f0f2402715c89c205941cd73 Mon Sep 17 00:00:00 2001
From: Arseny Smirnov <arseny30@gmail.com>
Date: Mon, 15 Jun 2020 19:23:40 +0300
Subject: [PATCH] tdutils: use new aes ige for long plaintext

GitOrigin-RevId: 4bd8ddd20508e235c0fb8b40ac42b9dcabfed30c
---
 benchmark/bench_crypto.cpp  | 31 +++++++++++++++++++++++++++++++
 tdutils/td/utils/crypto.cpp | 24 +++++++++++++++++++-----
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/benchmark/bench_crypto.cpp b/benchmark/bench_crypto.cpp
index 38ef179fe..781372d0a 100644
--- a/benchmark/bench_crypto.cpp
+++ b/benchmark/bench_crypto.cpp
@@ -161,6 +161,36 @@ class AesCbcBench : public td::Benchmark {
   }
 };
 
+class AesIgeShortBench : public td::Benchmark {
+ public:
+  static constexpr int DATA_SIZE = 16;
+  alignas(64) unsigned char data[DATA_SIZE];
+  td::UInt256 key;
+  td::UInt256 iv;
+
+  std::string get_description() const override {
+    return PSTRING() << "AES IGE OpenSSL [" << (DATA_SIZE) << "B]";
+  }
+
+  void start_up() override {
+    for (int i = 0; i < DATA_SIZE; i++) {
+      data[i] = 123;
+    }
+    td::Random::secure_bytes(as_slice(key));
+    td::Random::secure_bytes(as_slice(iv));
+  }
+
+  void run(int n) override {
+    td::MutableSlice data_slice(data, DATA_SIZE);
+    td::AesIgeState ige;
+    for (int i = 0; i < n; i++) {
+      ige.init(as_slice(key), as_slice(iv), true);
+      ige.encrypt(data_slice, data_slice);
+      //td::aes_ige_encrypt(as_slice(key), as_slice(iv), data_slice, data_slice);
+    }
+  }
+};
+
 BENCH(Rand, "std_rand") {
   int res = 0;
   for (int i = 0; i < n; i++) {
@@ -285,6 +315,7 @@ class Crc64Bench : public td::Benchmark {
 int main() {
   td::init_openssl_threads();
 
+  td::bench(AesIgeShortBench());
   td::bench(AesCtrBench());
   td::bench(AesEcbBench());
   td::bench(AesIgeBench());
diff --git a/tdutils/td/utils/crypto.cpp b/tdutils/td/utils/crypto.cpp
index d9edb37bd..70c2b175a 100644
--- a/tdutils/td/utils/crypto.cpp
+++ b/tdutils/td/utils/crypto.cpp
@@ -393,8 +393,10 @@ AesState::~AesState() = default;
 
 void AesState::init(Slice key, bool encrypt) {
   CHECK(key.size() == 32);
-  impl_ = make_unique<Impl>();
-  impl_->ctx = EVP_CIPHER_CTX_new();
+  if (!impl_) {
+    impl_ = make_unique<Impl>();
+    impl_->ctx = EVP_CIPHER_CTX_new();
+  }
   CHECK(impl_->ctx);
 
   if (encrypt) {
@@ -446,11 +448,21 @@ static void aes_ige_xcrypt(Slice aes_key, MutableSlice aes_iv, Slice from, Mutab
 }
 
 void aes_ige_encrypt(Slice aes_key, MutableSlice aes_iv, Slice from, MutableSlice to) {
-  aes_ige_xcrypt(aes_key, aes_iv, from, to, true);
+  if (from.size() <= 128) {
+    return aes_ige_xcrypt(aes_key, aes_iv, from, to, true);
+  }
+  AesIgeState state;
+  state.init(aes_key, aes_iv, true);
+  state.encrypt(from, to);
 }
 
 void aes_ige_decrypt(Slice aes_key, MutableSlice aes_iv, Slice from, MutableSlice to) {
-  aes_ige_xcrypt(aes_key, aes_iv, from, to, false);
+  if (from.size() <= 128) {
+    return aes_ige_xcrypt(aes_key, aes_iv, from, to, false);
+  }
+  AesIgeState state;
+  state.init(aes_key, aes_iv, false);
+  state.decrypt(from, to);
 }
 
 class AesIgeState::Impl {
@@ -512,7 +524,9 @@ AesIgeState::~AesIgeState() = default;
 void AesIgeState::init(Slice key, Slice iv, bool encrypt) {
   CHECK(key.size() == 32);
   CHECK(iv.size() == 32);
-  impl_ = make_unique<Impl>();
+  if (!impl_) {
+    impl_ = make_unique<Impl>();
+  }
   impl_->state.init(key, encrypt);
   impl_->iv.load(iv.ubegin());
   impl_->iv2.load(iv.ubegin() + AES_BLOCK_SIZE);