Skip to content

Commit c4ed5ce

Browse files
arseny30arseny30
committed
tdutils: aes experiments
GitOrigin-RevId: 1dd2812041778679fb1ef0462d9e8b1899181e0b
1 parent 541f15b commit c4ed5ce

File tree

6 files changed

+136
-1
lines changed

6 files changed

+136
-1
lines changed

‎benchmark/bench_crypto.cpp

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,36 @@ class SHA1Bench : public td::Benchmark {
4848
}
4949
};
5050

51+
class AesBench : public td::Benchmark {
52+
public:
53+
alignas(64) unsigned char data[DATA_SIZE];
54+
td::UInt256 key;
55+
td::UInt256 iv;
56+
57+
std::string get_description() const override {
58+
return PSTRING() << "AES ECB OpenSSL [" << (DATA_SIZE >> 10) << "KB]";
59+
}
60+
61+
void start_up() override {
62+
for (int i = 0; i < DATA_SIZE; i++) {
63+
data[i] = 123;
64+
}
65+
td::Random::secure_bytes(key.raw, sizeof(key));
66+
td::Random::secure_bytes(iv.raw, sizeof(iv));
67+
}
68+
69+
void run(int n) override {
70+
td::AesState state;
71+
state.init(td::as_slice(key), true);
72+
td::MutableSlice data_slice(data, DATA_SIZE);
73+
for (int i = 0; i <= n; i++) {
74+
for (size_t offset = 0; offset + 16 <= data_slice.size(); offset += 16) {
75+
state.encrypt(data_slice.ubegin() + offset, data_slice.ubegin() + offset);
76+
}
77+
}
78+
}
79+
};
80+
5181
class AESBench : public td::Benchmark {
5282
public:
5383
alignas(64) unsigned char data[DATA_SIZE];
@@ -197,6 +227,8 @@ class Crc64Bench : public td::Benchmark {
197227

198228
int main() {
199229
td::init_openssl_threads();
230+
td::bench(AesBench());
231+
td::bench(AESBench());
200232

201233
td::bench(Pbkdf2Bench());
202234
td::bench(RandBench());
@@ -208,7 +240,6 @@ int main() {
208240
#endif
209241
td::bench(SslRandBufBench());
210242
td::bench(SHA1Bench());
211-
td::bench(AESBench());
212243
td::bench(Crc32Bench());
213244
td::bench(Crc64Bench());
214245
return 0;

‎tdutils/td/utils/Random.cpp

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,4 +164,18 @@ int Random::Xorshift128plus::fast(int min, int max) {
164164
return static_cast<int>((*this)() % (max - min + 1) + min);
165165
}
166166

167+
void Random::Xorshift128plus::bytes(MutableSlice dest) {
168+
int cnt = 0;
169+
td::uint64 buf = 0;
170+
for (auto &c : dest) {
171+
if (cnt == 0) {
172+
buf = operator()();
173+
cnt = 8;
174+
}
175+
cnt--;
176+
c = buf & 255;
177+
buf >>= 8;
178+
}
179+
}
180+
167181
} // namespace td

‎tdutils/td/utils/Random.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ class Random {
3737
Xorshift128plus(uint64 seed_a, uint64 seed_b);
3838
uint64 operator()();
3939
int fast(int min, int max);
40+
void bytes(MutableSlice dest);
4041

4142
private:
4243
uint64 seed_[2];

‎tdutils/td/utils/crypto.cpp

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -249,6 +249,56 @@ int pq_factorize(Slice pq_str, string *p_str, string *q_str) {
249249
return 0;
250250
}
251251

252+
class AesState::Impl {
253+
public:
254+
EVP_CIPHER_CTX *ctx{nullptr};
255+
AES_KEY key;
256+
bool encrypt;
257+
258+
~Impl() {
259+
if (ctx != nullptr) {
260+
EVP_CIPHER_CTX_free(ctx);
261+
}
262+
}
263+
};
264+
265+
AesState::AesState() = default;
266+
AesState::~AesState() = default;
267+
268+
void AesState::init(Slice key, bool encrypt) {
269+
CHECK(key.size() == 32);
270+
impl_ = make_unique<Impl>();
271+
impl_->ctx = EVP_CIPHER_CTX_new();
272+
CHECK(impl_->ctx);
273+
274+
if (encrypt) {
275+
CHECK(1 == EVP_EncryptInit_ex(impl_->ctx, EVP_aes_256_ecb(), nullptr, key.ubegin(), nullptr));
276+
AES_set_encrypt_key(key.ubegin(), 256, &impl_->key);
277+
} else {
278+
CHECK(1 == EVP_DecryptInit_ex(impl_->ctx, EVP_aes_256_ecb(), nullptr, key.ubegin(), nullptr));
279+
AES_set_decrypt_key(key.ubegin(), 256, &impl_->key);
280+
}
281+
EVP_CIPHER_CTX_set_padding(impl_->ctx, 0);
282+
impl_->encrypt = encrypt;
283+
}
284+
285+
void AesState::encrypt(const uint8 *src, uint8 *dst) {
286+
CHECK(impl_->encrypt);
287+
CHECK(impl_->ctx);
288+
int len;
289+
CHECK(1 == EVP_EncryptUpdate(impl_->ctx, dst, &len, src, 16));
290+
CHECK(len == 16);
291+
//AES_encrypt(src, dst, &impl_->key);
292+
}
293+
void AesState::decrypt(const uint8 *src, uint8 *dst) {
294+
CHECK(!impl_->encrypt);
295+
CHECK(impl_->ctx);
296+
int len;
297+
CHECK(1 == EVP_DecryptUpdate(impl_->ctx, dst, &len, src, 16));
298+
LOG_CHECK(len == 16) << len;
299+
//AES_decrypt(src, dst, &impl_->key);
300+
}
301+
252302
static void aes_ige_xcrypt(Slice aes_key, MutableSlice aes_iv, Slice from, MutableSlice to, bool encrypt_flag) {
253303
CHECK(aes_key.size() == 32);
254304
CHECK(aes_iv.size() == 32);

‎tdutils/td/utils/crypto.h

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,23 @@ uint64 pq_factorize(uint64 pq);
1919
#if TD_HAVE_OPENSSL
2020
void init_crypto();
2121

22+
struct AesState {
23+
public:
24+
AesState();
25+
AesState(const AesState &from) = delete;
26+
AesState &operator=(const AesState &from) = delete;
27+
AesState(AesState &&from);
28+
AesState &operator=(AesState &&from);
29+
~AesState();
30+
void init(Slice key, bool encrypt);
31+
void encrypt(const uint8 *src, uint8 *dst);
32+
void decrypt(const uint8 *src, uint8 *dst);
33+
34+
private:
35+
class Impl;
36+
unique_ptr<Impl> impl_;
37+
};
38+
2239
int pq_factorize(Slice pq_str, string *p_str, string *q_str);
2340

2441
void aes_ige_encrypt(Slice aes_key, MutableSlice aes_iv, Slice from, MutableSlice to);

‎tdutils/test/crypto.cpp

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,28 @@
1818
static td::vector<td::string> strings{"", "1", "short test string", td::string(1000000, 'a')};
1919

2020
#if TD_HAVE_OPENSSL
21+
TEST(Crypto, Aes) {
22+
td::Random::Xorshift128plus rnd(123);
23+
td::UInt256 key;
24+
rnd.bytes(as_slice(key));
25+
std::string plaintext(16, 0);
26+
std::string encrypted(16, 0);
27+
std::string decrypted(16, 0);
28+
rnd.bytes(plaintext);
29+
30+
td::AesState encryptor;
31+
encryptor.init(as_slice(key), true);
32+
td::AesState decryptor;
33+
decryptor.init(as_slice(key), false);
34+
35+
encryptor.encrypt(td::as_slice(plaintext).ubegin(), td::as_slice(encrypted).ubegin());
36+
decryptor.decrypt(td::as_slice(encrypted).ubegin(), td::as_slice(decrypted).ubegin());
37+
38+
CHECK(decrypted == plaintext);
39+
CHECK(decrypted != encrypted);
40+
CHECK(td::crc32(encrypted) == 178892237);
41+
}
42+
2143
TEST(Crypto, AesCtrState) {
2244
td::vector<td::uint32> answers1{0u, 1141589763u, 596296607u, 3673001485u, 2302125528u,
2345
330967191u, 2047392231u, 3537459563u, 307747798u, 2149598133u};

0 commit comments

Comments
 (0)