lavu/aes: add runtime dispatch for crypt function

2 files changed, 33 insertions, 17 deletions

diff --git a/libavutil/aes.c b/libavutil/aes.c
index 61e9dd1b32..dde15be19c 100644
--- a/libavutil/aes.c
+++ b/libavutil/aes.c
@@ -126,31 +126,44 @@ static inline void aes_crypt(AVAES *a, int s, const uint8_t *sbox,
     subshift(&a->state[0], s, sbox);
 }
 
-void av_aes_crypt(AVAES *a, uint8_t *dst, const uint8_t *src,
-                  int count, uint8_t *iv, int decrypt)
+static void aes_encrypt(AVAES *a, uint8_t *dst, const uint8_t *src,
+                        int count, uint8_t *iv, int rounds)
 {
     while (count--) {
-        addkey_s(&a->state[1], src, &a->round_key[a->rounds]);
-        if (decrypt) {
-            aes_crypt(a, 0, inv_sbox, dec_multbl);
-            if (iv) {
-                addkey_s(&a->state[0], iv, &a->state[0]);
-                memcpy(iv, src, 16);
-            }
-            addkey_d(dst, &a->state[0], &a->round_key[0]);
-        } else {
-            if (iv)
-                addkey_s(&a->state[1], iv, &a->state[1]);
-            aes_crypt(a, 2, sbox, enc_multbl);
-            addkey_d(dst, &a->state[0], &a->round_key[0]);
-            if (iv)
-                memcpy(iv, dst, 16);
+        addkey_s(&a->state[1], src, &a->round_key[rounds]);
+        if (iv)
+            addkey_s(&a->state[1], iv, &a->state[1]);
+        aes_crypt(a, 2, sbox, enc_multbl);
+        addkey_d(dst, &a->state[0], &a->round_key[0]);
+        if (iv)
+            memcpy(iv, dst, 16);
+        src += 16;
+        dst += 16;
+    }
+}
+
+static void aes_decrypt(AVAES *a, uint8_t *dst, const uint8_t *src,
+                        int count, uint8_t *iv, int rounds)
+{
+    while (count--) {
+        addkey_s(&a->state[1], src, &a->round_key[rounds]);
+        aes_crypt(a, 0, inv_sbox, dec_multbl);
+        if (iv) {
+            addkey_s(&a->state[0], iv, &a->state[0]);
+            memcpy(iv, src, 16);
         }
+        addkey_d(dst, &a->state[0], &a->round_key[0]);
         src += 16;
         dst += 16;
     }
 }
 
+void av_aes_crypt(AVAES *a, uint8_t *dst, const uint8_t *src,
+                  int count, uint8_t *iv, int decrypt)
+{
+    a->crypt(a, dst, src, count, iv, a->rounds);
+}
+
 static void init_multbl2(uint32_t tbl[][256], const int c[4],
                          const uint8_t *log8, const uint8_t *alog8,
                          const uint8_t *sbox)
@@ -186,6 +199,8 @@ int av_aes_init(AVAES *a, const uint8_t *key, int key_bits, int decrypt)
     uint8_t log8[256];
     uint8_t alog8[512];
 
+    a->crypt = decrypt ? aes_decrypt : aes_encrypt;
+
     if (!enc_multbl[FF_ARRAY_ELEMS(enc_multbl)-1][FF_ARRAY_ELEMS(enc_multbl[0])-1]) {
         j = 1;
         for (i = 0; i < 255; i++) {
diff --git a/libavutil/aes_internal.h b/libavutil/aes_internal.h
index e2841efc1a..e5bf4bde07 100644
--- a/libavutil/aes_internal.h
+++ b/libavutil/aes_internal.h
@@ -36,6 +36,7 @@ typedef struct AVAES {
     av_aes_block round_key[15];
     av_aes_block state[2];
     int rounds;
+    void (*crypt)(struct AVAES *a, uint8_t *dst, const uint8_t *src, int count, uint8_t *iv, int rounds);
 } AVAES;
 
 #endif /* AVUTIL_AES_INTERNAL_H */