make AES Offload fallback to SW

Signed-off-by: Sameeh Jubran <sameeh@wolfssl.com>

Author: sameehj

wolfcrypt/src/aes.c

@@ -132,6 +132,14 @@ block cipher mechanism that uses n-bit binary string parameter key with 128-bits
     }
 #endif
 
+/* TLS-safe default:
+ * When TLS is enabled, AES-GCM must always retain software fallback.
+ * CryptoCB may accelerate, but must not remove fallback implicitly.
+ */
+#if !defined(NO_TLS)
+    #define WOLFSSL_AES_GCM_TLS_SAFE
+#endif
+
 /* Define AES implementation includes and functions */
 #if defined(STM32_CRYPTO)
      /* STM32F2/F4/F7/L4/L5/H7/WB55 hardware AES support for ECB, CBC, CTR and GCM modes */
@@ -4366,31 +4374,37 @@ static WARN_UNUSED_RESULT int wc_AesDecrypt(
             ret = wc_CryptoCb_AesSetKey(aes, userKey, keylen, &capabilities);
 
             if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE)) {
-                /* Callback handled it (success or error) */
-                if (ret == 0) {
-                    /* Check if callback declared full GCM offload capability.
-                     * If NOT, fall through to software path - we don't support
-                     * key-import-only mode for GCM (need key for table generation). */
-                    if (!(capabilities & WC_CRYPTOCB_AES_GCM)) {
-                        /* Fall through to software path */
-                        goto cryptocb_aes_setkey_sw_fallback;
-                    }
-                    
-                    /* Full GCM offload mode */
+                if (ret != 0) {
+                    /* real error from callback */
+                    return ret;
+                }
+
+                /* Callback handled SetKey */
+                if (capabilities & WC_CRYPTOCB_AES_GCM) {
                     aes->gcmCryptoCbOffload = 1;
                     aes->keylen = (int)keylen;
-                    
+
+                #ifdef WOLFSSL_AES_GCM_TLS_SAFE
+                    /* TLS-safe: DO NOT return
+                     * We must still set up software AES key for fallback.
+                     */
+                    goto cryptocb_aes_setkey_fallback;
+                #else
+                    /* Cryptonly build: full offload allowed */
                     /* Set IV if provided */
                     if (iv != NULL) {
                         XMEMCPY(aes->reg, iv, WC_AES_BLOCK_SIZE);
                     } else {
                         XMEMSET(aes->reg, 0, WC_AES_BLOCK_SIZE);
                     }
+                    return 0;
+                #endif
                 }
-                return ret;
+
+                /* Key-import-only or partial support → fallback */
             }
-            /* CRYPTOCB_UNAVAILABLE: fall through to software path */
-        cryptocb_aes_setkey_sw_fallback: (void)0;
+            /* Either UNAVAILABLE or TLS-safe fallback */
+        cryptocb_aes_setkey_fallback: (void)0;
         #else
             /* Standard CryptoCB path - copy key to devKey */
             if (keylen > sizeof(aes->devKey)) {
@@ -4830,31 +4844,37 @@ static void AesSetKey_C(Aes* aes, const byte* key, word32 keySz, int dir)
             ret = wc_CryptoCb_AesSetKey(aes, userKey, keylen, &capabilities);
 
             if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE)) {
-                /* Callback handled it (success or error) */
-                if (ret == 0) {
-                    /* Check if callback declared full GCM offload capability.
-                     * If NOT, fall through to software path - we don't support
-                     * key-import-only mode for GCM (need key for table generation). */
-                    if (!(capabilities & WC_CRYPTOCB_AES_GCM)) {
-                        /* Fall through to software path */
-                        goto cryptocb_aes_setkey_sw_fallback2;
-                    }
-                    
-                    /* Full GCM offload mode */
+                if (ret != 0) {
+                    /* real error from callback */
+                    return ret;
+                }
+
+                /* Callback handled SetKey */
+                if (capabilities & WC_CRYPTOCB_AES_GCM) {
                     aes->gcmCryptoCbOffload = 1;
                     aes->keylen = (int)keylen;
-                    
+
+                #ifdef WOLFSSL_AES_GCM_TLS_SAFE
+                    /* TLS-safe: DO NOT return
+                     * We must still set up software AES key for fallback.
+                     */
+                    goto cryptocb_aes_setkey_fallback2;
+                #else
+                    /* Cryptonly build: full offload allowed */
                     /* Set IV if provided */
                     if (iv != NULL) {
                         XMEMCPY(aes->reg, iv, WC_AES_BLOCK_SIZE);
                     } else {
                         XMEMSET(aes->reg, 0, WC_AES_BLOCK_SIZE);
                     }
+                    return 0;
+                #endif
                 }
-                return ret;
+
+                /* Key-import-only or partial support → fallback */
             }
-            /* CRYPTOCB_UNAVAILABLE: fall through to software path */
-        cryptocb_aes_setkey_sw_fallback2: (void)0;
+            /* Either UNAVAILABLE or TLS-safe fallback */
+        cryptocb_aes_setkey_fallback2: (void)0;
         #else
             /* Copy key to devKey for standard CryptoCB path */
             XMEMCPY(aes->devKey, userKey, keylen);
@@ -7533,24 +7553,16 @@ int wc_AesGcmSetKey(Aes* aes, const byte* key, word32 len)
 #if !defined(FREESCALE_LTC_AES_GCM) && !defined(WOLFSSL_PSOC6_CRYPTO)
     if (ret == 0) {
 #ifdef WOLF_CRYPTO_CB_AES_SETKEY
-        /* CRITICAL FIX: Only skip H table if callback EXPLICITLY declared
-         * full AES-GCM offload capability via WC_CRYPTOCB_AES_GCM flag.
-         *
-         * Checking devCtx alone is NOT sufficient because:
-         * - Callback may only support key import, not GCM operations
-         * - If we skip tables and GCM ops return UNAVAILABLE, software fallback fails
-         * - This was a critical bug in the original implementation
-         *
-         * The callback must explicitly opt-in to GCM offload by setting the
-         * WC_CRYPTOCB_AES_GCM capability flag during SetKey. Only then do we
-         * skip table generation.
-         */
+    #ifdef WOLFSSL_AES_GCM_TLS_SAFE
+        /* TLS-safe: always generate H for fallback */
+    #else
         if (aes->devId != INVALID_DEVID && 
             aes->devCtx != NULL &&
             aes->gcmCryptoCbOffload) {
-            /* H table not needed - callback confirmed full GCM offload */
+            /* full offload allowed: skip H */
         }
         else
+    #endif
 #endif
         {
             VECTOR_REGISTERS_PUSH;
@@ -7561,14 +7573,16 @@ int wc_AesGcmSetKey(Aes* aes, const byte* key, word32 len)
     }
     if (ret == 0) {
 #ifdef WOLF_CRYPTO_CB_AES_SETKEY
-        /* CRITICAL FIX: Same logic as H table.
-         * Only skip M0 generation if callback explicitly declared full GCM offload. */
+    #ifdef WOLFSSL_AES_GCM_TLS_SAFE
+        /* TLS-safe: always generate M0 for fallback */
+    #else
         if (aes->devId != INVALID_DEVID && 
             aes->devCtx != NULL &&
             aes->gcmCryptoCbOffload) {
-            /* M0 table not needed - callback confirmed full GCM offload */
+            /* full offload allowed: skip M0 */
         }
         else
+    #endif
 #endif
         {
 #if defined(GCM_TABLE) || defined(GCM_TABLE_4BIT)
@@ -7610,14 +7624,16 @@ int wc_AesGcmSetKey(Aes* aes, const byte* key, word32 len)
 #ifdef WOLF_CRYPTO_CB
     if (aes->devId != INVALID_DEVID) {
     #ifdef WOLF_CRYPTO_CB_AES_SETKEY
+    #if !defined(WOLFSSL_AES_GCM_TLS_SAFE)
         /* In CryptoCB key import mode, key is in SE - devKey not used */
         if (aes->devCtx != NULL) {
             /* Skip - key handled by CryptoCB */
         }
         else
+    #endif
     #endif
         {
-            /* Copy key to devKey for standard CryptoCB path */
+            /* Copy key to devKey for standard CryptoCB path or TLS-safe fallback */
             XMEMCPY(aes->devKey, key, len);
         }
     }