1 files changed, 14 insertions, 3 deletions
diff --git a/aes-encrypt.c b/aes-encrypt.c
index f962924a..d1496e54 100644
--- a/aes-encrypt.c
+++ b/aes-encrypt.c
@@ -36,6 +36,7 @@
 #endif
 
 #include <assert.h>
+#include <stdlib.h>
 
 #include "aes-internal.h"
 
@@ -47,9 +48,19 @@ aes_encrypt(const struct aes_ctx *ctx,
 	    size_t length, uint8_t *dst,
 	    const uint8_t *src)
 {
-  assert(!(length % AES_BLOCK_SIZE) );
-  _aes_encrypt(ctx->rounds, ctx->keys, &_aes_encrypt_table,
-	       length, dst, src);
+  switch (ctx->key_size)
+    {
+    default: abort();
+    case AES128_KEY_SIZE:
+      aes128_encrypt(&ctx->u.ctx128, length, dst, src);
+      break;
+    case AES192_KEY_SIZE:
+      aes192_encrypt(&ctx->u.ctx192, length, dst, src);
+      break;
+    case AES256_KEY_SIZE:
+      aes256_encrypt(&ctx->u.ctx256, length, dst, src);
+      break;
+    }
 }
 
 void