From be655c774837ed28a91d3d435c31198f7291b524 Mon Sep 17 00:00:00 2001
From: Gleb Panchishen <gleb.panchishen@nordigy.ru>
Date: Fri, 22 Mar 2019 02:20:13 +0300
Subject: [PATCH 1/4] evp support enabled for both cbc and ctr ciphers

---
 crypto/aesni256.c            | 111 +++++++++--------------------------
 crypto/aesni256.h            |  12 ++--
 net/net-crypto-aes.c         |  24 +++++---
 net/net-crypto-aes.h         |   3 -
 net/net-msg.c                |  16 ++---
 net/net-msg.h                |   2 +-
 net/net-tcp-connections.c    |   8 +--
 net/net-tcp-rpc-ext-server.c |   2 +-
 8 files changed, 60 insertions(+), 118 deletions(-)

diff --git a/crypto/aesni256.c b/crypto/aesni256.c
index f3bdc1a..08efa94 100644
--- a/crypto/aesni256.c
+++ b/crypto/aesni256.c
@@ -29,108 +29,51 @@
 
 #include <openssl/opensslv.h>
 
-#if OPENSSL_VERSION_NUMBER >= 0x10100000L
-#include <openssl/modes.h>
-void AES_ctr128_encrypt(
-		const unsigned char *in, 
-		unsigned char *out,
-		size_t length, 
-		const AES_KEY *key,
-		unsigned char ivec[AES_BLOCK_SIZE],
-		unsigned char ecount_buf[AES_BLOCK_SIZE],
-		unsigned int *num) {
-	CRYPTO_ctr128_encrypt(in, out, length, key, ivec, ecount_buf, num, (block128_f)AES_encrypt);
-}
-#endif
-
-void tg_ssl_aes_ctr_crypt (tg_aes_ctx_t *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[16], unsigned long long offset) {
-  unsigned char iv_copy[16];
-  memcpy (iv_copy, iv, 16);
-  unsigned long long *p = (unsigned long long *) (iv_copy + 8);
-  (*p) += offset >> 4;
-  union {
-    unsigned char c[16];
-    unsigned long long d[2];
-  } u;
-  int i = offset & 15, l;
-  if (i) {
-    AES_encrypt (iv_copy, u.c, &ctx->u.key);
-    (*p)++;
-    l = i + size;
-    if (l > 16) {
-      l = 16;
-    }
-    size -= l - i;
-    do {
-      *out++ = (*in++) ^ u.c[i++];
-    } while (i < l);
-  }
-  const unsigned long long *I = (const unsigned long long *) in;
-  unsigned long long *O = (unsigned long long *) out;
-  int n = size >> 4;
-  while (--n >= 0) {
-    AES_encrypt (iv_copy, (unsigned char *) u.d, &ctx->u.key);
-    (*p)++;
-    *O++ = (*I++) ^ u.d[0];
-    *O++ = (*I++) ^ u.d[1];
-  }
-  l = size & 15;
-  if (l) {
-    AES_encrypt (iv_copy, u.c, &ctx->u.key);
-    in = (const unsigned char *) I;
-    out = (unsigned char *) O;
-    i = 0;
-    do {
-      *out++ = (*in++) ^ u.c[i++];
-    } while (i < l);
-  }
-}
-
-
-static void tg_ssl_aes_cbc_encrypt (tg_aes_ctx_t *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[16]) {
-  AES_cbc_encrypt (in, out, size, &ctx->u.key, iv, AES_ENCRYPT);
-}
-
-static void tg_ssl_aes_cbc_decrypt (tg_aes_ctx_t *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[16]) {
-  AES_cbc_encrypt (in, out, size, &ctx->u.key, iv, AES_DECRYPT);
-}
+void tg_ssl_evp_encrypt(tg_aes_ctx_t *ctx, const unsigned char *in, unsigned char *out, int size) {
+  int len;
 
-static void tg_ssl_aes_ige_encrypt (tg_aes_ctx_t *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[32]) {
-  AES_ige_encrypt (in, out, size, &ctx->u.key, iv, AES_ENCRYPT);
+  assert(EVP_EncryptUpdate(ctx->evp_ctx, out, &len, in, size) == 1);
+  assert(EVP_EncryptFinal_ex(ctx->evp_ctx, out + len, &len) == 1);
 }
 
-static void tg_ssl_aes_ige_decrypt (tg_aes_ctx_t *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[32]) {
-  AES_ige_encrypt (in, out, size, &ctx->u.key, iv, AES_DECRYPT);
-}
+void tg_ssl_evp_decrypt(tg_aes_ctx_t *ctx, const unsigned char *in, unsigned char *out, int size) {
+  int len;
 
-void tg_ssl_aes_ctr128_crypt (struct tg_aes_ctx *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[16], unsigned char ecount_buf[16], unsigned int *num) {
-  AES_ctr128_encrypt (in, out, size, &ctx->u.key, iv, ecount_buf, num);
+  assert(EVP_DecryptUpdate(ctx->evp_ctx, out, &len, in, size) == 1);
+  assert(EVP_DecryptFinal_ex(ctx->evp_ctx, out + len, &len) == 1);
 }
 
 static const struct tg_aes_methods ssl_aes_encrypt_methods = {
-  .cbc_crypt = tg_ssl_aes_cbc_encrypt,
-  .ige_crypt = tg_ssl_aes_ige_encrypt,
-  .ctr_crypt = tg_ssl_aes_ctr_crypt,
-  .ctr128_crypt = tg_ssl_aes_ctr128_crypt
+  .cbc_crypt = tg_ssl_evp_encrypt,
+  .ctr_crypt = tg_ssl_evp_encrypt
 };
 
-void tg_aes_set_encrypt_key (tg_aes_ctx_t *ctx, unsigned char *key, int bits) {
-  AES_set_encrypt_key (key, bits, &ctx->u.key);
+void tg_aes_encrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char iv[16], const EVP_CIPHER *cipher) {
+  ctx->evp_ctx = EVP_CIPHER_CTX_new();
+  assert(ctx->evp_ctx);
+
+  assert(EVP_EncryptInit(ctx->evp_ctx, cipher, key, iv) == 1);
+  assert(EVP_CIPHER_CTX_set_padding(ctx->evp_ctx, 0));
+
   ctx->type = &ssl_aes_encrypt_methods;
 }
 
 static const struct tg_aes_methods ssl_aes_decrypt_methods = {
-  .cbc_crypt = tg_ssl_aes_cbc_decrypt,
-  .ige_crypt = tg_ssl_aes_ige_decrypt,
-  .ctr_crypt = NULL,
-  .ctr128_crypt = NULL
+  .cbc_crypt = tg_ssl_evp_decrypt,
+  .ctr_crypt = NULL
 };
 
-void tg_aes_set_decrypt_key (tg_aes_ctx_t *ctx, unsigned char *key, int bits) {
-  AES_set_decrypt_key (key, bits, &ctx->u.key);
+void tg_aes_decrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char iv[16], const EVP_CIPHER *cipher) {
+  ctx->evp_ctx = EVP_CIPHER_CTX_new();
+  assert(ctx->evp_ctx);
+
+  assert(EVP_DecryptInit(ctx->evp_ctx, cipher, key, iv) == 1);
+  assert(EVP_CIPHER_CTX_set_padding(ctx->evp_ctx, 0));
+
   ctx->type = &ssl_aes_decrypt_methods;
 }
 
 void tg_aes_ctx_cleanup (tg_aes_ctx_t *ctx) {
+  EVP_CIPHER_CTX_free(ctx->evp_ctx);
   memset (ctx, 0, sizeof (tg_aes_ctx_t));
 }
diff --git a/crypto/aesni256.h b/crypto/aesni256.h
index f2d1d17..2353372 100644
--- a/crypto/aesni256.h
+++ b/crypto/aesni256.h
@@ -24,6 +24,7 @@
 #pragma once
 
 #include <openssl/aes.h>
+#include <openssl/evp.h>
 
 struct aesni256_ctx {
   unsigned char a[256];
@@ -33,10 +34,8 @@ struct aesni256_ctx {
 struct tg_aes_ctx;
 
 struct tg_aes_methods {
-  void (*cbc_crypt) (struct tg_aes_ctx *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[16]);
-  void (*ige_crypt) (struct tg_aes_ctx *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[32]);
-  void (*ctr_crypt) (struct tg_aes_ctx *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[16], unsigned long long offset);
-  void (*ctr128_crypt) (struct tg_aes_ctx *ctx, const unsigned char *in, unsigned char *out, int size, unsigned char iv[16], unsigned char ecount_buf[16], unsigned int *num);
+  void (*cbc_crypt) (struct tg_aes_ctx *ctx, const unsigned char *in, unsigned char *out, int size);
+  void (*ctr_crypt) (struct tg_aes_ctx *ctx, const unsigned char *in, unsigned char *out, int size);
 };
 
 typedef struct tg_aes_ctx {
@@ -44,9 +43,10 @@ typedef struct tg_aes_ctx {
     AES_KEY key;
     struct aesni256_ctx ctx;
   } u;
+  EVP_CIPHER_CTX *evp_ctx;
   const struct tg_aes_methods *type;
 } tg_aes_ctx_t;
 
-void tg_aes_set_encrypt_key (tg_aes_ctx_t *ctx, unsigned char *key, int bits);
-void tg_aes_set_decrypt_key (tg_aes_ctx_t *ctx, unsigned char *key, int bits);
+void tg_aes_encrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char iv[16], const EVP_CIPHER *cipher);
+void tg_aes_decrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char iv[16], const EVP_CIPHER *cipher);
 void tg_aes_ctx_cleanup (tg_aes_ctx_t *ctx);
diff --git a/net/net-crypto-aes.c b/net/net-crypto-aes.c
index b6bc155..4461f60 100644
--- a/net/net-crypto-aes.c
+++ b/net/net-crypto-aes.c
@@ -84,14 +84,15 @@ int aes_crypto_init (connection_job_t c, void *key_data, int key_data_len) {
   struct aes_key_data *D = key_data;
   assert (T);
 
+  memset(&T->read_aeskey, 0, sizeof(T->read_aeskey));
+  memset(&T->write_aeskey, 0, sizeof(T->write_aeskey));
+
   MODULE_STAT->allocated_aes_crypto ++;
   
-  tg_aes_set_decrypt_key (&T->read_aeskey, D->read_key, 256);
+  tg_aes_decrypt_init (&T->read_aeskey, D->read_key, D->read_iv, EVP_aes_256_cbc());
   memcpy (T->read_iv, D->read_iv, 16);
-  tg_aes_set_encrypt_key (&T->write_aeskey, D->write_key, 256);
+  tg_aes_encrypt_init (&T->write_aeskey, D->write_key, D->write_iv, EVP_aes_256_cbc());
   memcpy (T->write_iv, D->write_iv, 16);
-  // T->read_pos = T->write_pos = 0;
-  T->read_num = T->write_num = 0;
   CONN_INFO(c)->crypto = T;
   return 0;
 }
@@ -103,20 +104,25 @@ int aes_crypto_ctr128_init (connection_job_t c, void *key_data, int key_data_len
   struct aes_key_data *D = key_data;
   assert (T);
 
+  memset(&T->read_aeskey, 0, sizeof(T->read_aeskey));
+  memset(&T->write_aeskey, 0, sizeof(T->write_aeskey));
+
   MODULE_STAT->allocated_aes_crypto ++;
   
-  tg_aes_set_encrypt_key (&T->read_aeskey, D->read_key, 256); // NB: *_encrypt_key here!
+  tg_aes_encrypt_init (&T->read_aeskey, D->read_key, D->read_iv, EVP_aes_256_ctr()); // NB: *_encrypt_init here!
   memcpy (T->read_iv, D->read_iv, 16);
-  tg_aes_set_encrypt_key (&T->write_aeskey, D->write_key, 256);
+  tg_aes_encrypt_init (&T->write_aeskey, D->write_key, D->write_iv, EVP_aes_256_ctr());
   memcpy (T->write_iv, D->write_iv, 16);
-  // T->read_pos = T->write_pos = 0;
-  T->read_num = T->write_num = 0;
   CONN_INFO(c)->crypto = T;
   return 0;
 }
 
 int aes_crypto_free (connection_job_t c) {
-  if (CONN_INFO(c)->crypto) {
+  struct aes_crypto *crypto = CONN_INFO(c)->crypto;
+  if (crypto) {
+    tg_aes_ctx_cleanup(&crypto->read_aeskey);
+    tg_aes_ctx_cleanup(&crypto->write_aeskey);
+
     free (CONN_INFO(c)->crypto);
     CONN_INFO(c)->crypto = 0;
     MODULE_STAT->allocated_aes_crypto --;
diff --git a/net/net-crypto-aes.h b/net/net-crypto-aes.h
index cdc123b..cf9afe4 100644
--- a/net/net-crypto-aes.h
+++ b/net/net-crypto-aes.h
@@ -71,11 +71,8 @@ struct aes_key_data {
 /* for c->crypto */
 struct aes_crypto {
   unsigned char read_iv[16], write_iv[16];
-  unsigned char read_ebuf[16], write_ebuf[16]; /* for AES-CTR modes */ 
   tg_aes_ctx_t read_aeskey __attribute__ ((aligned (16)));
   tg_aes_ctx_t write_aeskey __attribute__ ((aligned (16)));
-  unsigned int read_num, write_num; /* for AES-CTR modes */
-  // long long read_pos, write_pos; /* for AES-CTR modes */
 };
 
 extern int aes_initialized;
diff --git a/net/net-msg.c b/net/net-msg.c
index 71d2571..13e2031 100644
--- a/net/net-msg.c
+++ b/net/net-msg.c
@@ -1230,10 +1230,8 @@ struct rwm_encrypt_decrypt_tmp  {
   int block_size;
   struct raw_message *raw;
   struct tg_aes_ctx *ctx;
-  void (*crypt)(struct tg_aes_ctx *, const void *, void *, int, unsigned char *, void *, void *);
+  void (*crypt)(struct tg_aes_ctx *, const void *, void *, int, unsigned char *);
   unsigned char *iv;
-  void *extra;
-  void *extra2;
   char buf[16] __attribute__((aligned(16)));
 };
 
@@ -1261,12 +1259,12 @@ int rwm_process_encrypt_decrypt (struct rwm_encrypt_decrypt_tmp *x, const void *
       data += to_fill;
       x->bp = 0;     
       if (x->buf_left >= bsize) {
-        x->crypt (x->ctx, x->buf, res->last->part->data + res->last_offset, bsize, x->iv, x->extra, x->extra2);
+        x->crypt (x->ctx, x->buf, res->last->part->data + res->last_offset, bsize, x->iv);
         res->last->data_end += bsize;
         res->last_offset += bsize;
         x->buf_left -= bsize;
       } else {
-        x->crypt (x->ctx, x->buf, x->buf, bsize, x->iv, x->extra, x->extra2);
+        x->crypt (x->ctx, x->buf, x->buf, bsize, x->iv);
         memcpy (res->last->part->data + res->last_offset, x->buf, x->buf_left);
         int t = x->buf_left;
         res->last->data_end += t;
@@ -1316,7 +1314,7 @@ int rwm_process_encrypt_decrypt (struct rwm_encrypt_decrypt_tmp *x, const void *
     assert (x->buf_left + res->last_offset <= res->last->part->chunk->buffer_size);
     if (len <= x->buf_left) {
       assert (!(len & (bsize - 1)));
-      x->crypt (x->ctx, data, (res->last->part->data + res->last_offset), len, x->iv, x->extra, x->extra2);
+      x->crypt (x->ctx, data, (res->last->part->data + res->last_offset), len, x->iv);
       res->last->data_end += len;
       res->last_offset += len;
       res->total_bytes += len;
@@ -1324,7 +1322,7 @@ int rwm_process_encrypt_decrypt (struct rwm_encrypt_decrypt_tmp *x, const void *
       return 0;
     } else {
       int t = x->buf_left & -bsize;
-      x->crypt (x->ctx, data, res->last->part->data + res->last_offset, t, x->iv, x->extra, x->extra2);
+      x->crypt (x->ctx, data, res->last->part->data + res->last_offset, t, x->iv);
       res->last->data_end += t;
       res->last_offset += t;
       res->total_bytes += t;
@@ -1336,7 +1334,7 @@ int rwm_process_encrypt_decrypt (struct rwm_encrypt_decrypt_tmp *x, const void *
 }
 
 
-int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, int bytes, struct tg_aes_ctx *ctx, void (*crypt)(struct tg_aes_ctx *ctx, const void *src, void *dst, int l, unsigned char *iv, void *extra, void *extra2), unsigned char *iv, int block_size, void *extra, void *extra2) {
+int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, int bytes, struct tg_aes_ctx *ctx, void (*crypt)(struct tg_aes_ctx *ctx, const void *src, void *dst, int l, unsigned char *iv), unsigned char *iv, int block_size) {
   assert (bytes >= 0);
   assert (block_size && !(block_size & (block_size - 1)));
   if (bytes > raw->total_bytes) {
@@ -1375,8 +1373,6 @@ int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, in
   t.ctx = ctx;
   t.iv = iv;
   t.left = bytes;
-  t.extra = extra;
-  t.extra2 = extra2;
   t.block_size = block_size;
   int r = rwm_process_and_advance (raw, bytes, (void *)rwm_process_encrypt_decrypt, &t);
   if (locked) {
diff --git a/net/net-msg.h b/net/net-msg.h
index fb2c908..e41cc2d 100644
--- a/net/net-msg.h
+++ b/net/net-msg.h
@@ -147,7 +147,7 @@ int rwm_process_and_advance (struct raw_message *raw, int bytes, int (*process_b
 int rwm_sha1 (struct raw_message *raw, int bytes, unsigned char output[20]);
 // int rwm_encrypt_decrypt (struct raw_message *raw, int bytes, tg_aes_ctx_t *ctx, unsigned char iv[32]);
 // int rwm_encrypt_decrypt_cbc (struct raw_message *raw, int bytes, tg_aes_ctx_t *ctx, unsigned char iv[16]);
-int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, int bytes, tg_aes_ctx_t *ctx, void (*crypt)(tg_aes_ctx_t *ctx, const void *src, void *dst, int l, unsigned char *iv, void *extra, void *extra2), unsigned char *iv, int block_size, void *extra, void *extra2);
+int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, int bytes, tg_aes_ctx_t *ctx, void (*crypt)(tg_aes_ctx_t *ctx, const void *src, void *dst, int l, unsigned char *iv), unsigned char *iv, int block_size);
 
 void *rwm_get_block_ptr (struct raw_message *raw);
 int rwm_get_block_ptr_bytes (struct raw_message *raw);
diff --git a/net/net-tcp-connections.c b/net/net-tcp-connections.c
index 9b2f462..541d9ec 100644
--- a/net/net-tcp-connections.c
+++ b/net/net-tcp-connections.c
@@ -201,7 +201,7 @@ int cpu_tcp_aes_crypto_encrypt_output (connection_job_t C) /* {{{ */ {
   int l = out->total_bytes;
   l &= ~15;
   if (l) {
-    assert (rwm_encrypt_decrypt_to (&c->out, &c->out_p, l, &T->write_aeskey, (void *)T->write_aeskey.type->cbc_crypt, T->write_iv, 16, 0, 0) == l);
+    assert (rwm_encrypt_decrypt_to (&c->out, &c->out_p, l, &T->write_aeskey, (void *)T->write_aeskey.type->cbc_crypt, T->write_iv, 16) == l);
   }
 
   return (-out->total_bytes) & 15;
@@ -218,7 +218,7 @@ int cpu_tcp_aes_crypto_decrypt_input (connection_job_t C) /* {{{ */ {
   int l = in->total_bytes;
   l &= ~15;
   if (l) {
-    assert (rwm_encrypt_decrypt_to (&c->in_u, &c->in, l, &T->read_aeskey, (void *)T->read_aeskey.type->cbc_crypt, T->read_iv, 16, 0, 0) == l);
+    assert (rwm_encrypt_decrypt_to (&c->in_u, &c->in, l, &T->read_aeskey, (void *)T->read_aeskey.type->cbc_crypt, T->read_iv, 16) == l);
   }
 
   return (-in->total_bytes) & 15;
@@ -242,7 +242,7 @@ int cpu_tcp_aes_crypto_ctr128_encrypt_output (connection_job_t C) /* {{{ */ {
 
   int l = out->total_bytes;
   if (l) {
-    assert (rwm_encrypt_decrypt_to (&c->out, &c->out_p, l, &T->write_aeskey, (void *)T->write_aeskey.type->ctr128_crypt, T->write_iv, 1, T->write_ebuf, &T->write_num) == l);
+    assert (rwm_encrypt_decrypt_to (&c->out, &c->out_p, l, &T->write_aeskey, (void *)T->write_aeskey.type->ctr_crypt, T->write_iv, 1) == l);
   }
 
   return 0;
@@ -258,7 +258,7 @@ int cpu_tcp_aes_crypto_ctr128_decrypt_input (connection_job_t C) /* {{{ */ {
 
   int l = in->total_bytes;
   if (l) {
-    assert (rwm_encrypt_decrypt_to (&c->in_u, &c->in, l, &T->read_aeskey, (void *)T->read_aeskey.type->ctr128_crypt, T->read_iv, 1, T->read_ebuf, &T->read_num) == l);
+    assert (rwm_encrypt_decrypt_to (&c->in_u, &c->in, l, &T->read_aeskey, (void *)T->read_aeskey.type->ctr_crypt, T->read_iv, 1) == l);
   }
 
   return 0;
diff --git a/net/net-tcp-rpc-ext-server.c b/net/net-tcp-rpc-ext-server.c
index efc29cf..ec36edf 100644
--- a/net/net-tcp-rpc-ext-server.c
+++ b/net/net-tcp-rpc-ext-server.c
@@ -220,7 +220,7 @@ int tcp_rpcs_compact_parse_execute (connection_job_t C) {
         assert (c->crypto);
         struct aes_crypto *T = c->crypto;
 
-        T->read_aeskey.type->ctr128_crypt (&T->read_aeskey, random_header, random_header, 64, T->read_iv, T->read_ebuf, &T->read_num);
+        T->read_aeskey.type->ctr_crypt (&T->read_aeskey, random_header, random_header, 64);
         unsigned tag = *(unsigned *)(random_header + 56);
 
         if (tag == 0xdddddddd || tag == 0xeeeeeeee || tag == 0xefefefef) {

From e38fc7381ffb97ccf7bc1408fe73140568ba8f38 Mon Sep 17 00:00:00 2001
From: Gleb Panchishen <gleb.panchishen@nordigy.ru>
Date: Wed, 27 Mar 2019 01:52:12 +0300
Subject: [PATCH 2/4] crypto module drastically simplified

---
 crypto/aesni256.c            | 38 +++++++++++-------------------------
 crypto/aesni256.h            |  7 +------
 net/net-msg.c                | 15 ++++++--------
 net/net-msg.h                |  2 +-
 net/net-tcp-connections.c    |  8 ++++----
 net/net-tcp-rpc-ext-server.c |  2 +-
 6 files changed, 24 insertions(+), 48 deletions(-)

diff --git a/crypto/aesni256.c b/crypto/aesni256.c
index 08efa94..c431a58 100644
--- a/crypto/aesni256.c
+++ b/crypto/aesni256.c
@@ -29,48 +29,32 @@
 
 #include <openssl/opensslv.h>
 
-void tg_ssl_evp_encrypt(tg_aes_ctx_t *ctx, const unsigned char *in, unsigned char *out, int size) {
-  int len;
-
-  assert(EVP_EncryptUpdate(ctx->evp_ctx, out, &len, in, size) == 1);
-  assert(EVP_EncryptFinal_ex(ctx->evp_ctx, out + len, &len) == 1);
-}
-
-void tg_ssl_evp_decrypt(tg_aes_ctx_t *ctx, const unsigned char *in, unsigned char *out, int size) {
-  int len;
-
-  assert(EVP_DecryptUpdate(ctx->evp_ctx, out, &len, in, size) == 1);
-  assert(EVP_DecryptFinal_ex(ctx->evp_ctx, out + len, &len) == 1);
-}
-
-static const struct tg_aes_methods ssl_aes_encrypt_methods = {
-  .cbc_crypt = tg_ssl_evp_encrypt,
-  .ctr_crypt = tg_ssl_evp_encrypt
-};
-
 void tg_aes_encrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char iv[16], const EVP_CIPHER *cipher) {
   ctx->evp_ctx = EVP_CIPHER_CTX_new();
   assert(ctx->evp_ctx);
 
   assert(EVP_EncryptInit(ctx->evp_ctx, cipher, key, iv) == 1);
   assert(EVP_CIPHER_CTX_set_padding(ctx->evp_ctx, 0));
-
-  ctx->type = &ssl_aes_encrypt_methods;
 }
 
-static const struct tg_aes_methods ssl_aes_decrypt_methods = {
-  .cbc_crypt = tg_ssl_evp_decrypt,
-  .ctr_crypt = NULL
-};
-
 void tg_aes_decrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char iv[16], const EVP_CIPHER *cipher) {
   ctx->evp_ctx = EVP_CIPHER_CTX_new();
   assert(ctx->evp_ctx);
 
   assert(EVP_DecryptInit(ctx->evp_ctx, cipher, key, iv) == 1);
   assert(EVP_CIPHER_CTX_set_padding(ctx->evp_ctx, 0));
+}
+
+void tg_aes_crypt(tg_aes_ctx_t *ctx, const void *in, void *out, int size) {
+  int len;
 
-  ctx->type = &ssl_aes_decrypt_methods;
+  if (EVP_CIPHER_CTX_encrypting(ctx->evp_ctx)) {
+    assert(EVP_EncryptUpdate(ctx->evp_ctx, out, &len, in, size) == 1);
+    assert(EVP_EncryptFinal_ex(ctx->evp_ctx, out + len, &len) == 1);
+  } else {
+    assert(EVP_DecryptUpdate(ctx->evp_ctx, out, &len, in, size) == 1);
+    assert(EVP_DecryptFinal_ex(ctx->evp_ctx, out + len, &len) == 1);
+  }
 }
 
 void tg_aes_ctx_cleanup (tg_aes_ctx_t *ctx) {
diff --git a/crypto/aesni256.h b/crypto/aesni256.h
index 2353372..8e1be24 100644
--- a/crypto/aesni256.h
+++ b/crypto/aesni256.h
@@ -33,20 +33,15 @@ struct aesni256_ctx {
 //TODO: move cbc_crypt, ige_crypt, ctr_crypt to the virtual method table
 struct tg_aes_ctx;
 
-struct tg_aes_methods {
-  void (*cbc_crypt) (struct tg_aes_ctx *ctx, const unsigned char *in, unsigned char *out, int size);
-  void (*ctr_crypt) (struct tg_aes_ctx *ctx, const unsigned char *in, unsigned char *out, int size);
-};
-
 typedef struct tg_aes_ctx {
   union {
     AES_KEY key;
     struct aesni256_ctx ctx;
   } u;
   EVP_CIPHER_CTX *evp_ctx;
-  const struct tg_aes_methods *type;
 } tg_aes_ctx_t;
 
 void tg_aes_encrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char iv[16], const EVP_CIPHER *cipher);
 void tg_aes_decrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char iv[16], const EVP_CIPHER *cipher);
+void tg_aes_crypt (tg_aes_ctx_t *ctx, const void *in, void *out, int size); //bidirectional
 void tg_aes_ctx_cleanup (tg_aes_ctx_t *ctx);
diff --git a/net/net-msg.c b/net/net-msg.c
index 13e2031..13fb0d5 100644
--- a/net/net-msg.c
+++ b/net/net-msg.c
@@ -1230,8 +1230,6 @@ struct rwm_encrypt_decrypt_tmp  {
   int block_size;
   struct raw_message *raw;
   struct tg_aes_ctx *ctx;
-  void (*crypt)(struct tg_aes_ctx *, const void *, void *, int, unsigned char *);
-  unsigned char *iv;
   char buf[16] __attribute__((aligned(16)));
 };
 
@@ -1259,12 +1257,12 @@ int rwm_process_encrypt_decrypt (struct rwm_encrypt_decrypt_tmp *x, const void *
       data += to_fill;
       x->bp = 0;     
       if (x->buf_left >= bsize) {
-        x->crypt (x->ctx, x->buf, res->last->part->data + res->last_offset, bsize, x->iv);
+        tg_aes_crypt (x->ctx, x->buf, res->last->part->data + res->last_offset, bsize);
         res->last->data_end += bsize;
         res->last_offset += bsize;
         x->buf_left -= bsize;
       } else {
-        x->crypt (x->ctx, x->buf, x->buf, bsize, x->iv);
+        tg_aes_crypt (x->ctx, x->buf, x->buf, bsize);
         memcpy (res->last->part->data + res->last_offset, x->buf, x->buf_left);
         int t = x->buf_left;
         res->last->data_end += t;
@@ -1314,7 +1312,7 @@ int rwm_process_encrypt_decrypt (struct rwm_encrypt_decrypt_tmp *x, const void *
     assert (x->buf_left + res->last_offset <= res->last->part->chunk->buffer_size);
     if (len <= x->buf_left) {
       assert (!(len & (bsize - 1)));
-      x->crypt (x->ctx, data, (res->last->part->data + res->last_offset), len, x->iv);
+      tg_aes_crypt (x->ctx, data, (res->last->part->data + res->last_offset), len);
       res->last->data_end += len;
       res->last_offset += len;
       res->total_bytes += len;
@@ -1322,7 +1320,7 @@ int rwm_process_encrypt_decrypt (struct rwm_encrypt_decrypt_tmp *x, const void *
       return 0;
     } else {
       int t = x->buf_left & -bsize;
-      x->crypt (x->ctx, data, res->last->part->data + res->last_offset, t, x->iv);
+      tg_aes_crypt (x->ctx, data, res->last->part->data + res->last_offset, t);
       res->last->data_end += t;
       res->last_offset += t;
       res->total_bytes += t;
@@ -1334,9 +1332,10 @@ int rwm_process_encrypt_decrypt (struct rwm_encrypt_decrypt_tmp *x, const void *
 }
 
 
-int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, int bytes, struct tg_aes_ctx *ctx, void (*crypt)(struct tg_aes_ctx *ctx, const void *src, void *dst, int l, unsigned char *iv), unsigned char *iv, int block_size) {
+int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, int bytes, struct tg_aes_ctx *ctx, int block_size) {
   assert (bytes >= 0);
   assert (block_size && !(block_size & (block_size - 1)));
+
   if (bytes > raw->total_bytes) {
     bytes = raw->total_bytes;
   }
@@ -1363,7 +1362,6 @@ int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, in
   }
   struct rwm_encrypt_decrypt_tmp t;
   t.bp = 0;
-  t.crypt = crypt;
   if (res->last->part->refcnt == 1) {
     t.buf_left = res->last->part->chunk->buffer_size - res->last_offset;
   } else {
@@ -1371,7 +1369,6 @@ int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, in
   }
   t.raw = res;
   t.ctx = ctx;
-  t.iv = iv;
   t.left = bytes;
   t.block_size = block_size;
   int r = rwm_process_and_advance (raw, bytes, (void *)rwm_process_encrypt_decrypt, &t);
diff --git a/net/net-msg.h b/net/net-msg.h
index e41cc2d..a850cbd 100644
--- a/net/net-msg.h
+++ b/net/net-msg.h
@@ -147,7 +147,7 @@ int rwm_process_and_advance (struct raw_message *raw, int bytes, int (*process_b
 int rwm_sha1 (struct raw_message *raw, int bytes, unsigned char output[20]);
 // int rwm_encrypt_decrypt (struct raw_message *raw, int bytes, tg_aes_ctx_t *ctx, unsigned char iv[32]);
 // int rwm_encrypt_decrypt_cbc (struct raw_message *raw, int bytes, tg_aes_ctx_t *ctx, unsigned char iv[16]);
-int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, int bytes, tg_aes_ctx_t *ctx, void (*crypt)(tg_aes_ctx_t *ctx, const void *src, void *dst, int l, unsigned char *iv), unsigned char *iv, int block_size);
+int rwm_encrypt_decrypt_to (struct raw_message *raw, struct raw_message *res, int bytes, tg_aes_ctx_t *ctx, int block_size);
 
 void *rwm_get_block_ptr (struct raw_message *raw);
 int rwm_get_block_ptr_bytes (struct raw_message *raw);
diff --git a/net/net-tcp-connections.c b/net/net-tcp-connections.c
index 541d9ec..d8e3e55 100644
--- a/net/net-tcp-connections.c
+++ b/net/net-tcp-connections.c
@@ -201,7 +201,7 @@ int cpu_tcp_aes_crypto_encrypt_output (connection_job_t C) /* {{{ */ {
   int l = out->total_bytes;
   l &= ~15;
   if (l) {
-    assert (rwm_encrypt_decrypt_to (&c->out, &c->out_p, l, &T->write_aeskey, (void *)T->write_aeskey.type->cbc_crypt, T->write_iv, 16) == l);
+    assert (rwm_encrypt_decrypt_to (&c->out, &c->out_p, l, &T->write_aeskey, 16) == l);
   }
 
   return (-out->total_bytes) & 15;
@@ -218,7 +218,7 @@ int cpu_tcp_aes_crypto_decrypt_input (connection_job_t C) /* {{{ */ {
   int l = in->total_bytes;
   l &= ~15;
   if (l) {
-    assert (rwm_encrypt_decrypt_to (&c->in_u, &c->in, l, &T->read_aeskey, (void *)T->read_aeskey.type->cbc_crypt, T->read_iv, 16) == l);
+    assert (rwm_encrypt_decrypt_to (&c->in_u, &c->in, l, &T->read_aeskey, 16) == l);
   }
 
   return (-in->total_bytes) & 15;
@@ -242,7 +242,7 @@ int cpu_tcp_aes_crypto_ctr128_encrypt_output (connection_job_t C) /* {{{ */ {
 
   int l = out->total_bytes;
   if (l) {
-    assert (rwm_encrypt_decrypt_to (&c->out, &c->out_p, l, &T->write_aeskey, (void *)T->write_aeskey.type->ctr_crypt, T->write_iv, 1) == l);
+    assert (rwm_encrypt_decrypt_to (&c->out, &c->out_p, l, &T->write_aeskey, 1) == l);
   }
 
   return 0;
@@ -258,7 +258,7 @@ int cpu_tcp_aes_crypto_ctr128_decrypt_input (connection_job_t C) /* {{{ */ {
 
   int l = in->total_bytes;
   if (l) {
-    assert (rwm_encrypt_decrypt_to (&c->in_u, &c->in, l, &T->read_aeskey, (void *)T->read_aeskey.type->ctr_crypt, T->read_iv, 1) == l);
+    assert (rwm_encrypt_decrypt_to (&c->in_u, &c->in, l, &T->read_aeskey, 1) == l);
   }
 
   return 0;
diff --git a/net/net-tcp-rpc-ext-server.c b/net/net-tcp-rpc-ext-server.c
index ec36edf..ff0d96c 100644
--- a/net/net-tcp-rpc-ext-server.c
+++ b/net/net-tcp-rpc-ext-server.c
@@ -220,7 +220,7 @@ int tcp_rpcs_compact_parse_execute (connection_job_t C) {
         assert (c->crypto);
         struct aes_crypto *T = c->crypto;
 
-        T->read_aeskey.type->ctr_crypt (&T->read_aeskey, random_header, random_header, 64);
+        tg_aes_crypt(&T->read_aeskey, random_header, random_header, 64);
         unsigned tag = *(unsigned *)(random_header + 56);
 
         if (tag == 0xdddddddd || tag == 0xeeeeeeee || tag == 0xefefefef) {

From afba07c330d46f14a8307a320b18e140dbf2404c Mon Sep 17 00:00:00 2001
From: Gleb Panchishen <gleb.panchishen@nordigy.ru>
Date: Wed, 27 Mar 2019 02:40:26 +0300
Subject: [PATCH 3/4] added nullptr check and fixed assert conditions

---
 crypto/aesni256.c | 8 +++++---
 crypto/aesni256.h | 1 -
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/crypto/aesni256.c b/crypto/aesni256.c
index c431a58..c96449c 100644
--- a/crypto/aesni256.c
+++ b/crypto/aesni256.c
@@ -34,7 +34,7 @@ void tg_aes_encrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char i
   assert(ctx->evp_ctx);
 
   assert(EVP_EncryptInit(ctx->evp_ctx, cipher, key, iv) == 1);
-  assert(EVP_CIPHER_CTX_set_padding(ctx->evp_ctx, 0));
+  assert(EVP_CIPHER_CTX_set_padding(ctx->evp_ctx, 0) == 1);
 }
 
 void tg_aes_decrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char iv[16], const EVP_CIPHER *cipher) {
@@ -42,7 +42,7 @@ void tg_aes_decrypt_init (tg_aes_ctx_t *ctx, unsigned char *key, unsigned char i
   assert(ctx->evp_ctx);
 
   assert(EVP_DecryptInit(ctx->evp_ctx, cipher, key, iv) == 1);
-  assert(EVP_CIPHER_CTX_set_padding(ctx->evp_ctx, 0));
+  assert(EVP_CIPHER_CTX_set_padding(ctx->evp_ctx, 0) == 1);
 }
 
 void tg_aes_crypt(tg_aes_ctx_t *ctx, const void *in, void *out, int size) {
@@ -58,6 +58,8 @@ void tg_aes_crypt(tg_aes_ctx_t *ctx, const void *in, void *out, int size) {
 }
 
 void tg_aes_ctx_cleanup (tg_aes_ctx_t *ctx) {
-  EVP_CIPHER_CTX_free(ctx->evp_ctx);
+  if (ctx->evp_ctx) {
+    EVP_CIPHER_CTX_free(ctx->evp_ctx);
+  }
   memset (ctx, 0, sizeof (tg_aes_ctx_t));
 }
diff --git a/crypto/aesni256.h b/crypto/aesni256.h
index 8e1be24..8bf82e3 100644
--- a/crypto/aesni256.h
+++ b/crypto/aesni256.h
@@ -30,7 +30,6 @@ struct aesni256_ctx {
   unsigned char a[256];
 };
 
-//TODO: move cbc_crypt, ige_crypt, ctr_crypt to the virtual method table
 struct tg_aes_ctx;
 
 typedef struct tg_aes_ctx {

From 54ea5ee93e68b8dc14730e68b6f453a82cf7e18b Mon Sep 17 00:00:00 2001
From: Gleb Panchishen <gleb.panchishen@nordigy.ru>
Date: Fri, 29 Mar 2019 01:58:26 +0300
Subject: [PATCH 4/4] removed excess code

---
 crypto/aesni256.c    | 4 +---
 net/net-crypto-aes.c | 3 ---
 2 files changed, 1 insertion(+), 6 deletions(-)

diff --git a/crypto/aesni256.c b/crypto/aesni256.c
index c96449c..7eb9e6a 100644
--- a/crypto/aesni256.c
+++ b/crypto/aesni256.c
@@ -58,8 +58,6 @@ void tg_aes_crypt(tg_aes_ctx_t *ctx, const void *in, void *out, int size) {
 }
 
 void tg_aes_ctx_cleanup (tg_aes_ctx_t *ctx) {
-  if (ctx->evp_ctx) {
-    EVP_CIPHER_CTX_free(ctx->evp_ctx);
-  }
+  EVP_CIPHER_CTX_free(ctx->evp_ctx);
   memset (ctx, 0, sizeof (tg_aes_ctx_t));
 }
diff --git a/net/net-crypto-aes.c b/net/net-crypto-aes.c
index 4461f60..946fe71 100644
--- a/net/net-crypto-aes.c
+++ b/net/net-crypto-aes.c
@@ -84,9 +84,6 @@ int aes_crypto_init (connection_job_t c, void *key_data, int key_data_len) {
   struct aes_key_data *D = key_data;
   assert (T);
 
-  memset(&T->read_aeskey, 0, sizeof(T->read_aeskey));
-  memset(&T->write_aeskey, 0, sizeof(T->write_aeskey));
-
   MODULE_STAT->allocated_aes_crypto ++;
   
   tg_aes_decrypt_init (&T->read_aeskey, D->read_key, D->read_iv, EVP_aes_256_cbc());