From db013c6bb857b606f1e494962f08cb6ae625a7e4 Mon Sep 17 00:00:00 2001
From: kaleb-himes <kaleb@wolfssl.com>
Date: Fri, 17 May 2024 17:30:00 -0400
Subject: [PATCH] Add SP800-38E limitation

---
 wolfcrypt/src/aes.c     | 17 ++++++++++-------
 wolfssl/wolfcrypt/aes.h |  6 ++++++
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/wolfcrypt/src/aes.c b/wolfcrypt/src/aes.c
index 0416cd2947..73e8167473 100644
--- a/wolfcrypt/src/aes.c
+++ b/wolfcrypt/src/aes.c
@@ -12785,13 +12785,16 @@ int wc_AesXtsEncrypt(XtsAes* xaes, byte* out, const byte* in, word32 sz,
 
     aes = &xaes->aes;
 
-/* FIPS TODO: SP800-38E - Restrict data unit to 2^20 blocks per key. A block is
- * AES_BLOCK_SIZE or 16-bytes (128-bits). So each key may only be used to
- * protect up to 1,048,576 blocks of AES_BLOCK_SIZE (16,777,216 bytes or
- * 134,217,728-bits) Add helpful printout and message along with BAD_FUNC_ARG
- * return whenever sz / AES_BLOCK_SIZE > 1,048,576 or equal to that and sz is
- * not a sequence of complete blocks.
- */
+    #if FIPS_VERSION3_GE(6,0,0)
+    /* SP800-38E - Restrict data unit to 2^20 blocks per key. A block is
+     * AES_BLOCK_SIZE or 16-bytes (128-bits). So each key may only be used to
+     * protect up to 1,048,576 blocks of AES_BLOCK_SIZE (16,777,216 bytes)
+     */
+    if (sz > FIPS_XTS_LIMIT) {
+        WOLFSSL_MSG("Request exceeds allowed bytes per SP800-38E");
+        return BAD_FUNC_ARG;
+    }
+    #endif
 
     if (aes->keylen == 0) {
         WOLFSSL_MSG("wc_AesXtsEncrypt called with unset encryption key.");
diff --git a/wolfssl/wolfcrypt/aes.h b/wolfssl/wolfcrypt/aes.h
index 38bc8c32ce..31a39991ab 100644
--- a/wolfssl/wolfcrypt/aes.h
+++ b/wolfssl/wolfcrypt/aes.h
@@ -58,6 +58,12 @@ typedef struct Gcm {
 #if FIPS_VERSION3_GE(6,0,0)
     extern const unsigned int wolfCrypt_FIPS_aes_ro_sanity[2];
     WOLFSSL_LOCAL int wolfCrypt_FIPS_AES_sanity(void);
+
+    /* SP800-38E - Restrict data unit to 2^20 blocks per key. A block is
+     * AES_BLOCK_SIZE or 16-bytes (128-bits). So each key may only be used to
+     * protect up to 1,048,576 blocks of AES_BLOCK_SIZE (16,777,216 bytes)
+     */
+    #define FIPS_XTS_LIMIT 16777216
 #endif
 
 WOLFSSL_LOCAL void GenerateM0(Gcm* gcm);