From 2c25481e7d52ab54d2a886bd209f291facfcb6bf Mon Sep 17 00:00:00 2001
From: toddouska <todd@yassl.com>
Date: Fri, 28 Sep 2012 15:01:07 -0700
Subject: [PATCH] add AesSetKeyDirect for Ctr and Direct when also using aesni

---
 ctaocrypt/src/aes.c    | 26 ++++++++++++++++++++++++++
 cyassl/ctaocrypt/aes.h |  3 ++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/ctaocrypt/src/aes.c b/ctaocrypt/src/aes.c
index f21e7a1a2..eabd0f4ee 100644
--- a/ctaocrypt/src/aes.c
+++ b/ctaocrypt/src/aes.c
@@ -1027,6 +1027,13 @@ static void AesEncrypt(Aes* aes, const byte* inBlock, byte* outBlock)
         CYASSL_MSG("AesEncrypt encountered improper key, set it up");
         return;  /* stop instead of segfaulting, set up your keys! */
     }
+#ifdef CYASSL_AESNI
+    if (haveAESNI) {
+        CYASSL_MSG("AesEncrypt encountered aesni keysetup, don't use direct");
+        return;  /* just stop now */
+    } 
+#endif
+
     /*
      * map byte array block to cipher state
      * and add initial round key:
@@ -1165,6 +1172,13 @@ static void AesDecrypt(Aes* aes, const byte* inBlock, byte* outBlock)
         CYASSL_MSG("AesDecrypt encountered improper key, set it up");
         return;  /* stop instead of segfaulting, set up your keys! */
     }
+#ifdef CYASSL_AESNI
+    if (haveAESNI) {
+        CYASSL_MSG("AesEncrypt encountered aesni keysetup, don't use direct");
+        return;  /* just stop now */
+    } 
+#endif
+
     /*
      * map byte array block to cipher state
      * and add initial round key:
@@ -1381,6 +1395,18 @@ void AesDecryptDirect(Aes* aes, byte* out, const byte* in)
 #endif /* CYASSL_AES_DIRECT */
 
 
+#if defined(CYASSL_AES_DIRECT) || defined(CYASSL_AES_COUNTER)
+
+/* AES-CTR and AES-DIRECT need to use this for key setup, no aesni yet */
+int AesSetKeyDirect(Aes* aes, const byte* userKey, word32 keylen,
+                    const byte* iv, int dir)
+{
+    return AesSetKeyLocal(aes, userKey, keylen, iv, dir);
+}
+
+#endif /* CYASSL_AES_DIRECT || CYASSL_AES_COUNTER */
+
+
 #ifdef CYASSL_AES_COUNTER
 
 /* Increment AES counter */
diff --git a/cyassl/ctaocrypt/aes.h b/cyassl/ctaocrypt/aes.h
index 9ab625dfc..e8dc53312 100644
--- a/cyassl/ctaocrypt/aes.h
+++ b/cyassl/ctaocrypt/aes.h
@@ -87,7 +87,8 @@ CYASSL_API void AesCbcDecrypt(Aes* aes, byte* out, const byte* in, word32 sz);
 CYASSL_API void AesCtrEncrypt(Aes* aes, byte* out, const byte* in, word32 sz);
 CYASSL_API void AesEncryptDirect(Aes* aes, byte* out, const byte* in);
 CYASSL_API void AesDecryptDirect(Aes* aes, byte* out, const byte* in);
-
+CYASSL_API int  AesSetKeyDirect(Aes* aes, const byte* key, word32 len,
+                                const byte* iv, int dir);
 #ifdef HAVE_AESGCM
 CYASSL_API void AesGcmSetKey(Aes* aes, const byte* key, word32 len,
                               const byte* implicitIV);