fix GCM test

Signed-off-by: Sameeh Jubran <[email protected]>

Author: sameehj

tests/api/test_aes.c

@@ -5283,6 +5283,7 @@ static int test_CryptoCb_Aes_Cb(int devId, wc_CryptoInfo* info, void* ctx)
 
         /* Store handle in aes->devCtx - this is what wolfSSL will use */
         aes->devCtx = cryptoCbAesMockHandle;
+        aes->devFlags |= WC_AES_FLAG_GCM_OFFLOAD;  /* We handle GCM ops */
 
         cryptoCbAesSetKeyCalled++;
 
@@ -5627,6 +5628,7 @@ static int test_CryptoCb_AesGcm_Offload_Cb(int devId, wc_CryptoInfo* info, void*
 
         /* Store handle in aes->devCtx - this is what wolfSSL will use */
         aes->devCtx = cryptoCbAesGcmMockHandle;
+        aes->devFlags |= WC_AES_FLAG_GCM_OFFLOAD;  /* We handle GCM ops */
 
         cryptoCbAesGcmSetKeyCalled++;
 

wolfcrypt/src/aes.c

@@ -7511,7 +7511,8 @@ int wc_AesGcmSetKey(Aes* aes, const byte* key, word32 len)
 #ifdef WOLF_CRYPTO_CB_AES_SETKEY
         /* In CryptoCB key import mode, skip H table generation.
          * The secure element handles GCM internally. */
-        if (aes->devId != INVALID_DEVID && aes->devCtx != NULL) {
+        if (aes->devId != INVALID_DEVID && aes->devCtx != NULL &&
+            (aes->devFlags & WC_AES_FLAG_GCM_OFFLOAD)) {
             /* H table not needed - SE does GCM */
         }
         else
@@ -7529,7 +7530,8 @@ int wc_AesGcmSetKey(Aes* aes, const byte* key, word32 len)
 #ifdef WOLF_CRYPTO_CB_AES_SETKEY
         /* In CryptoCB key import mode, skip M0 table generation.
          * The secure element handles GCM internally. */
-        if (aes->devId != INVALID_DEVID && aes->devCtx != NULL) {
+        if (aes->devId != INVALID_DEVID && aes->devCtx != NULL &&
+            (aes->devFlags & WC_AES_FLAG_GCM_OFFLOAD)) {
             /* M0 table not needed - SE does GCM */
         }
         else
@@ -13473,14 +13475,17 @@ void wc_AesFree(Aes* aes)
     {
         int ret = wc_CryptoCb_Free(aes->devId, WC_ALGO_TYPE_CIPHER,
                                    WC_CIPHER_AES, aes);
+        if (ret == WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE)) {
+            /* Callback didn't handle it - fall through to software cleanup */
+        }
+        else {
+            /* Callback handled cleanup */
     #ifdef WOLF_CRYPTO_CB_AES_SETKEY
-        aes->devCtx = NULL;  /* Clear device context handle */
+            aes->devCtx = NULL;
+            aes->devFlags = 0;
     #endif
-        /* If callback wants standard free, it can set devId to INVALID_DEVID.
-         * Otherwise assume the callback handled cleanup. */
-        if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
             return;
-        /* fall-through when unavailable */
+        }
     }
 #endif /* WOLF_CRYPTO_CB && WOLF_CRYPTO_CB_FREE */
 

wolfssl/wolfcrypt/aes.h

@@ -237,6 +237,10 @@ enum {
     WOLF_ENUM_DUMMY_LAST_ELEMENT(AES)
 };
 
+/* CryptoCB device capability flags */
+#define WC_AES_FLAG_NONE        0x00
+#define WC_AES_FLAG_GCM_OFFLOAD 0x01  /* Device handles AES-GCM encrypt/decrypt */
+
 #ifdef WC_AES_BITSLICED
     #ifdef WC_AES_BS_WORD_SIZE
         #define BS_WORD_SIZE        WC_AES_BS_WORD_SIZE
@@ -335,6 +339,9 @@ struct Aes {
 #if defined(WOLF_CRYPTO_CB) || defined(WOLFSSL_STM32U5_DHUK)
     int    devId;
     void*  devCtx;
+#ifdef WOLF_CRYPTO_CB
+    byte   devFlags;  /* CryptoCB capability flags */
+#endif
 #endif
 #ifdef WOLF_PRIVATE_KEY_ID
     byte id[AES_MAX_ID_LEN];