Merge pull request #62 (Fixes AESNI alignment bug)

9 files changed, 108 insertions, 17 deletions

diff --git a/configure.ac b/configure.ac
index b62db5e..5f22d00 100644
--- a/configure.ac
+++ b/configure.ac
@@ -103,5 +103,6 @@ AC_TYPE_UINT8_T
 # Checks for library functions.
 AC_FUNC_MALLOC
 AC_CHECK_FUNCS([memmove memset])
+AC_CHECK_FUNCS([posix_memalign aligned_alloc _aligned_malloc])
 
 AC_OUTPUT
diff --git a/setup.py b/setup.py
index 998122e..5269e9d 100644
--- a/setup.py
+++ b/setup.py
@@ -166,6 +166,10 @@ class PCTBuildExt (build_ext):
         if not ac.get("HAVE_WMMINTRIN_H"):
             # AES-NI instrincs not available
             self.__remove_extensions(["Crypto.Cipher._AESNI"])
+        elif not (ac.get("HAVE_POSIX_MEMALIGN") or ac.get("HAVE_ALIGNED_ALLOC")
+                  or ac.get("HAVE__ALIGNED_MALLOC")):
+            # no function to allocate aligned memory is available
+            self.__remove_extensions(["Crypto.Cipher._AESNI"])
         elif ac.get("HAVE_MAES"):
             # -maes has to be passed to the compiler to use the AES-NI instrincs
             self.__add_extension_compile_option(["Crypto.Cipher._AESNI"],
diff --git a/src/AES.c b/src/AES.c
index 5da189f..1e705da 100644
--- a/src/AES.c
+++ b/src/AES.c
@@ -1446,6 +1446,10 @@ static void block_init(block_state *state, unsigned char *key,
 	rijndaelKeySetupDec(state->dk, key, keylen*8);
 }
 
+static void block_finalize(block_state* self)
+{
+}
+
 static void block_encrypt(block_state *self, u8 *in, u8 *out)
 {
 	rijndaelEncrypt(self->ek, self->rounds, in, out);
diff --git a/src/AESNI.c b/src/AESNI.c
index 750f719..cbd1448 100644
--- a/src/AESNI.c
+++ b/src/AESNI.c
@@ -22,25 +22,64 @@
  * ===================================================================
  */
 
-#include "Python.h"
+#include "pycrypto_common.h"
 #include <wmmintrin.h>
+#include <stdlib.h>
+#if defined(HAVE__ALIGNED_MALLOC)
+#include <malloc.h>
+#endif
 
 #define MODULE_NAME _AESNI
 #define BLOCK_SIZE 16
 #define KEY_SIZE 0
 
-#define MAXKC	(256/32)
-#define MAXKB	(256/8)
-#define MAXNR	14
+#define MAXKC (256/32)
+#define MAXKB (256/8)
+#define MAXNR 14
 
 typedef unsigned char u8;
 
 typedef struct {
-	__m128i ek[MAXNR + 1];
-	__m128i dk[MAXNR + 1];
-	int rounds;
+    __m128i* ek;
+    __m128i* dk;
+    int rounds;
 } block_state;
 
+/* Wrapper function for malloc and free with memory alignment */
+
+static void* memalign_wrapper(size_t alignment, size_t size)
+{
+#if defined(HAVE_ALIGNED_ALLOC)
+    /* aligned_alloc is defined by C11 */
+    return aligned_alloc(alignment, size);
+#elif defined(HAVE_POSIX_MEMALIGN)
+    /* posix_memalign is defined by POSIX */
+    void* tmp = NULL;
+    int result = posix_memalign(&tmp, alignment, size);
+    if (result != 0)
+        return NULL;
+    return tmp;
+#elif defined(HAVE__ALIGNED_MALLOC)
+    /* _aligned_malloc is available on Windows */
+    return _aligned_malloc(size, alignment);
+#else
+#error "No function to allocate/free aligned memory is available."
+#endif
+}
+
+static void free_wrapper(void* ptr)
+{
+#if defined(HAVE_POSIX_MEMALIGN) || defined(HAVE_ALIGNED_ALLOC)
+    /* free is fine for aligned_alloc and posix_memalign */
+    free(ptr);
+#elif defined(HAVE__ALIGNED_MALLOC)
+    /* _aligned_malloc requires _aligned_free */
+    _aligned_free(ptr);
+#else
+#error "No function to allocate/free aligned memory is available."
+#endif
+}
+
 /* Helper functions to expand keys */
 
 static __m128i aes128_keyexpand(__m128i key, __m128i keygened, int shuf)
@@ -151,21 +190,46 @@ static void aes_key_setup_dec(__m128i dk[], const __m128i ek[], int rounds)
 
 static void block_init(block_state* self, unsigned char* key, int keylen)
 {
-	int nr = 0;
-	switch (keylen) {
-	    case 16: nr = 10; break;
-	    case 24: nr = 12; break;
-	    case 32: nr = 14; break;
+    int nr = 0;
+    switch (keylen) {
+        case 16: nr = 10; break;
+        case 24: nr = 12; break;
+        case 32: nr = 14; break;
         default:
             PyErr_SetString(PyExc_ValueError,
                 "AES key must be either 16, 24, or 32 bytes long");
-		    return;
-	}
-	self->rounds = nr;
-	aes_key_setup_enc(self->ek, key, keylen);
+            return;
+    }
+
+    /* ensure that self->ek and self->dk are aligned to 16 byte boundaries */
+    void* tek = memalign_wrapper(16, (nr + 1) * sizeof(__m128i));
+    void* tdk = memalign_wrapper(16, (nr + 1) * sizeof(__m128i));
+    if (!tek || !tdk) {
+        free_wrapper(tek);
+        free_wrapper(tdk);
+        PyErr_SetString(PyExc_MemoryError,
+                "failed to allocate memory for keys");
+        return;
+    }
+
+    self->ek = tek;
+    self->dk = tdk;
+
+    self->rounds = nr;
+    aes_key_setup_enc(self->ek, key, keylen);
     aes_key_setup_dec(self->dk, self->ek, nr);
 }
 
+static void block_finalize(block_state* self)
+{
+    /* overwrite contents of ek and dk */
+    memset(self->ek, 0, (self->rounds + 1) * sizeof(__m128i));
+    memset(self->dk, 0, (self->rounds + 1) * sizeof(__m128i));
+
+    free_wrapper(self->ek);
+    free_wrapper(self->dk);
+}
+
 static void block_encrypt(block_state* self, const u8* in, u8* out)
 {
     __m128i m = _mm_loadu_si128((const __m128i*) in);
@@ -196,7 +260,7 @@ static void block_encrypt(block_state* self, const u8* in, u8* out)
 
 static void block_decrypt(block_state* self, const u8* in, u8* out)
 {
-	__m128i m = _mm_loadu_si128((const __m128i*) in);
+    __m128i m = _mm_loadu_si128((const __m128i*) in);
     /* first 9 rounds */
     m = _mm_xor_si128(m, self->dk[0]);
     m = _mm_aesdec_si128(m, self->dk[1]);
diff --git a/src/ARC2.c b/src/ARC2.c
index 255e77e..71858f6 100644
--- a/src/ARC2.c
+++ b/src/ARC2.c
@@ -216,5 +216,9 @@ block_init(block_state *self, U8 *key, int keylength)
 	} while (i--);
 }
 
+static void
+block_finalize(block_state* self)
+{
+}
 
 #include "block_template.c"
diff --git a/src/Blowfish.c b/src/Blowfish.c
index 96f8628..f1ab55a 100644
--- a/src/Blowfish.c
+++ b/src/Blowfish.c
@@ -231,6 +231,10 @@ static void Blowfish_init(Blowfish_state *self, const unsigned char *key, int ke
 #define block_encrypt Blowfish_encrypt
 #define block_decrypt Blowfish_decrypt
 
+static void block_finalize(block_state *self)
+{
+}
+
 #include "block_template.c"
 
 /* vim:set ts=4 sw=4 sts=4 expandtab: */
diff --git a/src/CAST.c b/src/CAST.c
index ca5d0f0..d7b00f3 100644
--- a/src/CAST.c
+++ b/src/CAST.c
@@ -435,6 +435,11 @@ block_init(block_state *self, unsigned char *key, int keylength)
 }
 
 static void
+block_finalize(block_state* self)
+{
+}
+
+static void
 block_encrypt(block_state *self, unsigned char *in,
 	      unsigned char *out)
 {
diff --git a/src/DES.c b/src/DES.c
index 41e4153..65171fb 100644
--- a/src/DES.c
+++ b/src/DES.c
@@ -80,6 +80,10 @@ static void block_init(block_state *self, unsigned char *key, int keylen)
     }
 }
 
+static void block_finalize(block_state *self)
+{
+}
+
 static void block_encrypt(block_state *self, unsigned char *in, unsigned char *out)
 {
     int rc;
diff --git a/src/block_template.c b/src/block_template.c
index d555ceb..99aee43 100644
--- a/src/block_template.c
+++ b/src/block_template.c
@@ -84,6 +84,7 @@ static void
 ALGdealloc(PyObject *ptr)
 {		
 	ALGobject *self = (ALGobject *)ptr;
+	block_finalize(&self->st);
 
 	/* Overwrite the contents of the object */
 	Py_XDECREF(self->counter);