From 360c7497032d1e6394ffa8277a06ce0a5fbb1062 Mon Sep 17 00:00:00 2001
From: Daniel Pouzzner <douzzer@wolfssl.com>
Date: Fri, 28 Aug 2020 17:16:44 -0500
Subject: [PATCH] add {SAVE,RESTORE}_VECTOR_REGISTERS() macros for
 kernel_fpu_{begin,end} when WOLFSSL_LINUXKM, to allow safe use of AESNI and
 SIMD instructions in the kernel.

---
 wolfcrypt/src/aes.c         | 12 ++++++++++++
 wolfssl/wolfcrypt/wc_port.h | 12 +++++++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/wolfcrypt/src/aes.c b/wolfcrypt/src/aes.c
index 5b0338f60..5ed473f26 100644
--- a/wolfcrypt/src/aes.c
+++ b/wolfcrypt/src/aes.c
@@ -6257,23 +6257,29 @@ int wc_AesGcmEncrypt(Aes* aes, byte* out, const byte* in, word32 sz,
 #ifdef WOLFSSL_AESNI
     #ifdef HAVE_INTEL_AVX2
     if (IS_INTEL_AVX2(intel_flags)) {
+        SAVE_VECTOR_REGISTERS();
         AES_GCM_encrypt_avx2(in, out, authIn, iv, authTag, sz, authInSz, ivSz,
                                  authTagSz, (const byte*)aes->key, aes->rounds);
+        RESTORE_VECTOR_REGISTERS();
         return 0;
     }
     else
     #endif
     #ifdef HAVE_INTEL_AVX1
     if (IS_INTEL_AVX1(intel_flags)) {
+        SAVE_VECTOR_REGISTERS();
         AES_GCM_encrypt_avx1(in, out, authIn, iv, authTag, sz, authInSz, ivSz,
                                  authTagSz, (const byte*)aes->key, aes->rounds);
+        RESTORE_VECTOR_REGISTERS();
         return 0;
     }
     else
     #endif
     if (haveAESNI) {
+        SAVE_VECTOR_REGISTERS();
         AES_GCM_encrypt(in, out, authIn, iv, authTag, sz, authInSz, ivSz,
                                  authTagSz, (const byte*)aes->key, aes->rounds);
+        RESTORE_VECTOR_REGISTERS();
         return 0;
     }
     else
@@ -6727,8 +6733,10 @@ int wc_AesGcmDecrypt(Aes* aes, byte* out, const byte* in, word32 sz,
 #ifdef WOLFSSL_AESNI
     #ifdef HAVE_INTEL_AVX2
     if (IS_INTEL_AVX2(intel_flags)) {
+        SAVE_VECTOR_REGISTERS();
         AES_GCM_decrypt_avx2(in, out, authIn, iv, authTag, sz, authInSz, ivSz,
                                  authTagSz, (byte*)aes->key, aes->rounds, &res);
+        RESTORE_VECTOR_REGISTERS();
         if (res == 0)
             return AES_GCM_AUTH_E;
         return 0;
@@ -6737,8 +6745,10 @@ int wc_AesGcmDecrypt(Aes* aes, byte* out, const byte* in, word32 sz,
     #endif
     #ifdef HAVE_INTEL_AVX1
     if (IS_INTEL_AVX1(intel_flags)) {
+        SAVE_VECTOR_REGISTERS();
         AES_GCM_decrypt_avx1(in, out, authIn, iv, authTag, sz, authInSz, ivSz,
                                  authTagSz, (byte*)aes->key, aes->rounds, &res);
+        RESTORE_VECTOR_REGISTERS();
         if (res == 0)
             return AES_GCM_AUTH_E;
         return 0;
@@ -6746,8 +6756,10 @@ int wc_AesGcmDecrypt(Aes* aes, byte* out, const byte* in, word32 sz,
     else
     #endif
     if (haveAESNI) {
+        SAVE_VECTOR_REGISTERS();
         AES_GCM_decrypt(in, out, authIn, iv, authTag, sz, authInSz, ivSz,
                                  authTagSz, (byte*)aes->key, aes->rounds, &res);
+        RESTORE_VECTOR_REGISTERS();
         if (res == 0)
             return AES_GCM_AUTH_E;
         return 0;
diff --git a/wolfssl/wolfcrypt/wc_port.h b/wolfssl/wolfcrypt/wc_port.h
index ef6c0fd57..ae68674f1 100644
--- a/wolfssl/wolfcrypt/wc_port.h
+++ b/wolfssl/wolfcrypt/wc_port.h
@@ -88,6 +88,7 @@
     #endif
     #include <linux/net.h>
     #include <linux/slab.h>
+    #include <asm/simd.h>
     _Pragma("GCC diagnostic pop");
 
     /* remove this multifariously conflicting macro, picked up from
@@ -114,7 +115,16 @@
     #define XSNPRINTF snprintf /* needed to suppress inclusion of stdio.h in wolfssl/wolfcrypt/types.h */
     /* the rigmarole around kstrtol() here is to accommodate its warn-unused-result attribute. */
     #define XATOI(s) ({ long _xatoi_res = 0; int _xatoi_ret = kstrtol(s, 10, &_xatoi_res); if (_xatoi_ret != 0) { _xatoi_res = 0; } (int)_xatoi_res; })
-#endif
+
+    #define SAVE_VECTOR_REGISTERS() kernel_fpu_begin()
+    #define RESTORE_VECTOR_REGISTERS() kernel_fpu_end()
+
+#else /* ! WOLFSSL_LINUXKM */
+
+    #define SAVE_VECTOR_REGISTERS()
+    #define RESTORE_VECTOR_REGISTERS()
+
+#endif /* WOLFSSL_LINUXKM */
 
 /* THREADING/MUTEX SECTION */
 #ifdef USE_WINDOWS_API