From 45c05ffd3009e8416b71a1dd0f88b5fe3ac81f95 Mon Sep 17 00:00:00 2001
From: toddouska <todd@wolfssl.com>
Date: Thu, 23 Jan 2014 12:34:27 -0800
Subject: [PATCH] add non block size AesCtr support

---
 ctaocrypt/src/aes.c    | 34 +++++++++++++++++++++++++++++++---
 cyassl/ctaocrypt/aes.h |  3 +++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/ctaocrypt/src/aes.c b/ctaocrypt/src/aes.c
index 1aea29f8c..4cb723d31 100644
--- a/ctaocrypt/src/aes.c
+++ b/ctaocrypt/src/aes.c
@@ -1476,6 +1476,10 @@ static int AesSetKeyLocal(Aes* aes, const byte* userKey, word32 keylen,
     #ifdef CYASSL_AESNI
         aes->use_aesni = 0;
     #endif /* CYASSL_AESNI */
+    #ifdef CYASSL_AES_COUNTER
+        aes->left = 0;
+    #endif /* CYASSL_AES_COUNTER */
+
     aes->rounds = keylen/4 + 6;
 
     XMEMCPY(rk, userKey, keylen);
@@ -2129,15 +2133,39 @@ static INLINE void IncrementAesCounter(byte* inOutCtr)
 
 void AesCtrEncrypt(Aes* aes, byte* out, const byte* in, word32 sz)
 {
-    word32 blocks = sz / AES_BLOCK_SIZE;
+    byte* tmp = (byte*)aes->tmp + AES_BLOCK_SIZE - aes->left;
 
-    while (blocks--) {
+    /* consume any unused bytes left in aes->tmp */
+    while (aes->left && sz) {
+       *(out++) = *(in++) ^ *(tmp++);
+       aes->left--;
+       sz--;
+    }
+
+    /* do as many block size ops as possible */
+    while (sz >= AES_BLOCK_SIZE) {
         AesEncrypt(aes, (byte*)aes->reg, out);
         IncrementAesCounter((byte*)aes->reg);
         xorbuf(out, in, AES_BLOCK_SIZE);
 
         out += AES_BLOCK_SIZE;
-        in  += AES_BLOCK_SIZE; 
+        in  += AES_BLOCK_SIZE;
+        sz  -= AES_BLOCK_SIZE;
+        aes->left = 0;
+    }
+
+    /* handle non block size remaining and sotre unused byte count in left */
+    if (sz) {
+        AesEncrypt(aes, (byte*)aes->reg, (byte*)aes->tmp);
+        IncrementAesCounter((byte*)aes->reg);
+
+        aes->left = AES_BLOCK_SIZE;
+        tmp = (byte*)aes->tmp;
+
+        while (sz--) {
+            *(out++) = *(in++) ^ *(tmp++);
+            aes->left--;
+        }
     }
 }
 
diff --git a/cyassl/ctaocrypt/aes.h b/cyassl/ctaocrypt/aes.h
index 37861903e..e280e006b 100644
--- a/cyassl/ctaocrypt/aes.h
+++ b/cyassl/ctaocrypt/aes.h
@@ -92,6 +92,9 @@ typedef struct Aes {
     word32  magic;           /* using cavium magic */
     word64  contextHandle;   /* nitrox context memory handle */
 #endif
+#ifdef CYASSL_AES_COUNTER
+    word32  left;            /* unsued bytes left from last call */
+#endif 
 } Aes;