tdutils: simplify aes ige

GitOrigin-RevId: bc99b8c159cc866875ef69e44cd1e88f5e5dffee

benchmark/bench_crypto.cpp

@@ -79,34 +79,6 @@ class AesEcbBench : public td::Benchmark {
   }
 };
 
-class AesCtrBench : public td::Benchmark {
- public:
-  alignas(64) unsigned char data[DATA_SIZE];
-  td::UInt256 key;
-  td::UInt128 iv;
-
-  std::string get_description() const override {
-    return PSTRING() << "AES CTR 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::AesCtrState state;
-    state.init(as_slice(key), as_slice(iv));
-    td::MutableSlice data_slice(data, DATA_SIZE);
-    for (int i = 0; i < n; i++) {
-      state.encrypt(data_slice, data_slice);
-    }
-  }
-};
-
 class AesIgeBench : public td::Benchmark {
  public:
   alignas(64) unsigned char data[DATA_SIZE];
@@ -135,6 +107,60 @@ class AesIgeBench : public td::Benchmark {
   }
 };
 
+class AesCtrBench : public td::Benchmark {
+ public:
+  alignas(64) unsigned char data[DATA_SIZE];
+  td::UInt256 key;
+  td::UInt128 iv;
+
+  std::string get_description() const override {
+    return PSTRING() << "AES CTR 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::MutableSlice data_slice(data, DATA_SIZE);
+    td::AesCtrState state;
+    state.init(as_slice(key), as_slice(iv));
+    for (int i = 0; i < n; i++) {
+      state.encrypt(data_slice, data_slice);
+    }
+  }
+};
+
+class AesCbcBench : public td::Benchmark {
+ public:
+  alignas(64) unsigned char data[DATA_SIZE];
+  td::UInt256 key;
+  td::UInt128 iv;
+
+  std::string get_description() const override {
+    return PSTRING() << "AES CBC OpenSSL [" << (DATA_SIZE >> 10) << "KB]";
+  }
+
+  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);
+    for (int i = 0; i < n; i++) {
+      td::aes_cbc_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++) {
@@ -258,6 +284,7 @@ class Crc64Bench : public td::Benchmark {
 
 int main() {
   td::init_openssl_threads();
+
   td::bench(AesEcbBench());
   td::bench(AesIgeBench());
   td::bench(AesCtrBench());

tdutils/td/utils/crypto.cpp

@@ -32,16 +32,6 @@
 #include <openssl/sha.h>
 #endif
 
-#if TD_HAVE_OPENSSL
-#define N_WORDS (AES_BLOCK_SIZE / sizeof(unsigned long))
-typedef struct {
-  unsigned long data[N_WORDS];
-} aes_block_t;
-
-#define load_block(d, s) memcpy((d).data, (s), AES_BLOCK_SIZE)
-#define store_block(d, s) memcpy((d), (s).data, AES_BLOCK_SIZE)
-#endif
-
 #if TD_HAVE_ZLIB
 #include <zlib.h>
 #endif
@@ -58,6 +48,34 @@ typedef struct {
 
 namespace td {
 
+struct alignas(8) AesBlock {
+  uint64 hi;
+  uint64 lo;
+
+  uint8 *raw() {
+    return reinterpret_cast<uint8 *>(this);
+  }
+
+  AesBlock operator^(const AesBlock &b) const {
+    AesBlock res;
+    res.hi = hi ^ b.hi;
+    res.lo = lo ^ b.lo;
+    return res;
+  }
+  AesBlock &operator^=(const AesBlock &b) {
+    hi ^= b.hi;
+    lo ^= b.lo;
+    return *this;
+  }
+
+  void load(const uint8 *from) {
+    *this = as<AesBlock>(from);
+  }
+  void store(uint8 *to) {
+    as<AesBlock>(to) = *this;
+  }
+};
+
 static uint64 gcd(uint64 a, uint64 b) {
   if (a == 0) {
     return b;
@@ -345,28 +363,24 @@ void aes_ige_decrypt(Slice aes_key, MutableSlice aes_iv, Slice from, MutableSlic
 class AesIgeState::Impl {
  public:
   AesState state;
-  aes_block_t iv;
+  AesBlock iv;
-  aes_block_t iv2;
+  AesBlock iv2;
   void encrypt(Slice from, MutableSlice to) {
     CHECK(from.size() % AES_BLOCK_SIZE == 0);
     CHECK(to.size() >= from.size());
     auto len = to.size() / AES_BLOCK_SIZE;
     auto in = from.ubegin();
-    auto out = to.begin();
+    auto out = to.ubegin();
 
-    aes_block_t tmp, tmp2;
+    AesBlock tmp, tmp2;
 
     while (len) {
-      load_block(tmp, in);
+      tmp.load(in);
-      for (size_t n = 0; n < N_WORDS; ++n) {
+      tmp2 = tmp ^ iv;
-        tmp2.data[n] = tmp.data[n] ^ iv.data[n];
+      state.encrypt(tmp2.raw(), tmp2.raw(), AES_BLOCK_SIZE);
-      }
 
-      state.encrypt((unsigned char *)tmp2.data, (unsigned char *)tmp2.data, AES_BLOCK_SIZE);
+      tmp2 ^= iv2;
-      for (size_t n = 0; n < N_WORDS; ++n) {
+      tmp2.store(out);
-        tmp2.data[n] ^= iv2.data[n];
-      }
-      store_block(out, tmp2);
       iv = tmp2;
       iv2 = tmp;
       --len;
@@ -379,21 +393,17 @@ class AesIgeState::Impl {
     CHECK(to.size() >= from.size());
     auto len = to.size() / AES_BLOCK_SIZE;
     auto in = from.ubegin();
-    auto out = to.begin();
+    auto out = to.ubegin();
 
-    aes_block_t tmp, tmp2;
+    AesBlock tmp, tmp2;
 
     while (len) {
-      load_block(tmp, in);
+      tmp.load(in);
       tmp2 = tmp;
-      for (size_t n = 0; n < N_WORDS; ++n) {
+      tmp ^= iv2;
-        tmp.data[n] ^= iv2.data[n];
+      state.decrypt(tmp.raw(), tmp.raw(), AES_BLOCK_SIZE);
-      }
+      tmp ^= iv;
-      state.decrypt((unsigned char *)tmp.data, (unsigned char *)tmp.data, AES_BLOCK_SIZE);
+      tmp.store(out);
-      for (size_t n = 0; n < N_WORDS; ++n) {
-        tmp.data[n] ^= iv.data[n];
-      }
-      store_block(out, tmp);
       iv = tmp2;
       iv2 = tmp;
       --len;
@@ -411,8 +421,8 @@ void AesIgeState::init(Slice key, Slice iv, bool encrypt) {
   CHECK(iv.size() == 32);
   impl_ = make_unique<Impl>();
   impl_->state.init(key, encrypt);
-  load_block(impl_->iv, iv.ubegin());
+  impl_->iv.load(iv.ubegin());
-  load_block(impl_->iv2, iv.ubegin() + AES_BLOCK_SIZE);
+  impl_->iv2.load(iv.ubegin() + AES_BLOCK_SIZE);
 }
 
 void AesIgeState::encrypt(Slice from, MutableSlice to) {