Fix issue #198: AESNI breaks with messages shorter than 16 bytes

lib/Crypto/SelfTest/Cipher/test_AES.py

@@ -1265,16 +1265,39 @@ def runTest(self):
             self.assertEqual(SHA256.new(ct).hexdigest(), expected)
 
 
+class TestIncompleteBlocks(unittest.TestCase):
+
+    def __init__(self, use_aesni):
+        unittest.TestCase.__init__(self)
+        self.use_aesni = use_aesni
+
+    def runTest(self):
+        # Encrypt data with length not multiple of 16 bytes
+
+        cipher = AES.new(b'4'*16, AES.MODE_ECB, use_aesni=self.use_aesni)
+
+        for msg_len in range(1, 16):
+            self.assertRaises(ValueError, cipher.encrypt, b'1' * msg_len)
+            self.assertRaises(ValueError, cipher.encrypt, b'1' * (msg_len+16))
+            self.assertRaises(ValueError, cipher.decrypt, b'1' * msg_len)
+            self.assertRaises(ValueError, cipher.decrypt, b'1' * (msg_len+16))
+
+        self.assertEqual(cipher.encrypt(b''), b'')
+        self.assertEqual(cipher.decrypt(b''), b'')
+
+
 def get_tests(config={}):
     from Crypto.Util import _cpu_features
     from common import make_block_tests
 
     tests = make_block_tests(AES, "AES", test_data, {'use_aesni': False})
     tests += [ TestMultipleBlocks(False) ]
+    tests += [ TestIncompleteBlocks(False) ]
     if _cpu_features.have_aes_ni():
         # Run tests with AES-NI instructions if they are available.
         tests += make_block_tests(AES, "AESNI", test_data, {'use_aesni': True})
         tests += [ TestMultipleBlocks(True) ]
+        tests += [ TestIncompleteBlocks(True) ]
     else:
         print "Skipping AESNI tests"
     return tests

src/AESNI.c

@@ -222,7 +222,7 @@ static int AESNI_encrypt(const BlockBase *bb, const uint8_t *in, uint8_t *out, s
     }
 
     /** There are 7 blocks or fewer left **/
-    for (;data_len>0; data_len-=16, in+=16, out+=16) {
+    for (;data_len>=BLOCK_SIZE; data_len-=BLOCK_SIZE, in+=BLOCK_SIZE, out+=BLOCK_SIZE) {
         __m128i pt, data;
         unsigned i;
 
@@ -331,7 +331,7 @@ static int AESNI_decrypt(const BlockBase *bb, const uint8_t *in, uint8_t *out, s
     }
 
     /** There are 7 blocks or fewer left **/
-    for (;data_len>0; data_len-=16, in+=16, out+=16) {
+    for (;data_len>=BLOCK_SIZE; data_len-=BLOCK_SIZE, in+=BLOCK_SIZE, out+=BLOCK_SIZE) {
         __m128i ct, data;
         unsigned i;