From afbc65a6c37331ab70333a7f132458f3cc360faa Mon Sep 17 00:00:00 2001
From: night1rider <backman23@live.com>
Date: Mon, 22 Dec 2025 10:48:49 -0700
Subject: [PATCH] Aes Free callback support

---
 tests/api.c           | 17 +++++++++++++++++
 wolfcrypt/src/aes.c   | 24 ++++++++++++++++++++++++
 wolfcrypt/test/test.c | 17 +++++++++++++++++
 3 files changed, 58 insertions(+)

diff --git a/tests/api.c b/tests/api.c
index 1a3162b24..bdfa09f26 100644
--- a/tests/api.c
+++ b/tests/api.c
@@ -35747,6 +35747,23 @@ static int test_CryptoCb_Func(int thisDevId, wc_CryptoInfo* info, void* ctx)
                     ret = 0;
                     break;
                 }
+    #endif
+                default:
+                    ret = WC_NO_ERR_TRACE(NOT_COMPILED_IN);
+                    break;
+            }
+        }
+        else if (info->free.algo == WC_ALGO_TYPE_CIPHER) {
+            switch (info->free.type) {
+    #ifndef NO_AES
+                case WC_CIPHER_AES:
+                {
+                    Aes* aes = (Aes*)info->free.obj;
+                    aes->devId = INVALID_DEVID;
+                    wc_AesFree(aes);
+                    ret = 0;
+                    break;
+                }
     #endif
                 default:
                     ret = WC_NO_ERR_TRACE(NOT_COMPILED_IN);
diff --git a/wolfcrypt/src/aes.c b/wolfcrypt/src/aes.c
index cdc32c739..ddb7e31db 100644
--- a/wolfcrypt/src/aes.c
+++ b/wolfcrypt/src/aes.c
@@ -13402,10 +13402,34 @@ int wc_AesInit_Label(Aes* aes, const char* label, void* heap, int devId)
 /* Free Aes resources */
 void wc_AesFree(Aes* aes)
 {
+#if defined(WOLF_CRYPTO_CB) && defined(WOLF_CRYPTO_CB_FREE)
+    int ret = 0;
+#endif
+
     if (aes == NULL) {
         return;
     }
 
+#if defined(WOLF_CRYPTO_CB) && defined(WOLF_CRYPTO_CB_FREE)
+    #ifndef WOLF_CRYPTO_CB_FIND
+    if (aes->devId != INVALID_DEVID)
+    #endif
+    {
+        ret = wc_CryptoCb_Free(aes->devId, WC_ALGO_TYPE_CIPHER,
+                         WC_CIPHER_AES, (void*)aes);
+        /* If they want the standard free, they can call it themselves */
+        /* via their callback setting devId to INVALID_DEVID */
+        /* otherwise assume the callback handled it */
+        if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
+            return;
+        /* fall-through when unavailable */
+    }
+
+    /* silence compiler warning */
+    (void)ret;
+
+#endif /* WOLF_CRYPTO_CB && WOLF_CRYPTO_CB_FREE */
+
 #ifdef WC_DEBUG_CIPHER_LIFECYCLE
     (void)wc_debug_CipherLifecycleFree(&aes->CipherLifecycleTag, aes->heap, 1);
 #endif
diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index 971502e29..c87ab9c53 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -62522,6 +62522,23 @@ static int myCryptoDevCb(int devIdArg, wc_CryptoInfo* info, void* ctx)
                     ret = 0;
                     break;
                 }
+#endif
+                default:
+                    ret = WC_NO_ERR_TRACE(NOT_COMPILED_IN);
+                    break;
+            }
+        }
+        else if (info->free.algo == WC_ALGO_TYPE_CIPHER) {
+            switch (info->free.type) {
+#ifndef NO_AES
+                case WC_CIPHER_AES:
+                {
+                    Aes* aes = (Aes*)info->free.obj;
+                    aes->devId = INVALID_DEVID;
+                    wc_AesFree(aes);
+                    ret = 0;
+                    break;
+                }
 #endif
                 default:
                     ret = WC_NO_ERR_TRACE(NOT_COMPILED_IN);