From 41c19f6aa0284bd5afa8c1d2be747e067f538165 Mon Sep 17 00:00:00 2001
From: jordan <jordan@wolfssl.com>
Date: Sun, 9 Aug 2026 22:53:40 -0500
Subject: [PATCH] bsdkm driver: wolfkdriv_aes_ctx_clear cleanup.

---
 bsdkm/wolfkmod.c | 42 ++++++++++++++++++++++++++++++++++--------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/bsdkm/wolfkmod.c b/bsdkm/wolfkmod.c
index faa570e87b..32df739eb4 100644
--- a/bsdkm/wolfkmod.c
+++ b/bsdkm/wolfkmod.c
@@ -383,16 +383,19 @@ static void km_AesFree(Aes * aes) {
     #endif
 }
 
-static void wolfkdriv_aes_ctx_clear(km_aes_ctx * ctx)
+/* clean up allocated km_aes_ctx struct.
+ *   - cbc allocates both encrypt and decrypt, and frees both.
+ *   - gcm uses only aes_encrypt.
+ * */
+static void wolfkdriv_aes_ctx_clear(km_aes_ctx * ctx, int free_decrypt)
 {
     if (ctx != NULL) {
         km_AesFree(&ctx->aes_encrypt);
-        km_AesFree(&ctx->aes_decrypt);
-    }
 
-    #ifdef WOLFKM_DEBUG_AES
-    printf("info: exiting km_AesExitCommon\n");
-    #endif /* WOLFKM_DEBUG_AES */
+        if (free_decrypt) {
+            km_AesFree(&ctx->aes_decrypt);
+        }
+    }
 }
 
 static void wolfkdriv_identify(driver_t * driver, device_t parent)
@@ -671,7 +674,16 @@ static int wolfkdriv_newsession_aes(device_t dev,
 
 newsession_cipher_out:
     if (error != 0) {
-        wolfkdriv_aes_ctx_clear(&session->aes_ctx);
+        switch (csp->csp_cipher_alg) {
+        case CRYPTO_AES_NIST_GCM_16:
+            wolfkdriv_aes_ctx_clear(&session->aes_ctx, 0);
+            break;
+        case CRYPTO_AES_CBC:
+            wolfkdriv_aes_ctx_clear(&session->aes_ctx, 1);
+        default:
+            break;
+        }
+
         return (EINVAL);
     }
 
@@ -713,13 +725,27 @@ static void
 wolfkdriv_freesession(device_t dev, crypto_session_t cses)
 {
     wolfkdriv_session_t * session = NULL;
+    const struct crypto_session_params * csp = NULL;
     (void)dev;
 
     /* get the wolfkdriv_session_t context */
     session = crypto_get_driver_session(cses);
+    csp = crypto_get_params(cses);
 
     /* clean it up */
-    wolfkdriv_aes_ctx_clear(&session->aes_ctx);
+    switch (csp->csp_mode) {
+    case CSP_MODE_CIPHER:
+        wolfkdriv_aes_ctx_clear(&session->aes_ctx, 1);
+        break;
+    case CSP_MODE_DIGEST:
+    case CSP_MODE_ETA:
+        break;
+    case CSP_MODE_AEAD:
+        wolfkdriv_aes_ctx_clear(&session->aes_ctx, 0);
+        break;
+    default:
+        __assert_unreachable();
+    }
 
     #if defined(WOLFSSL_BSDKM_VERBOSE_DEBUG)
     device_printf(dev, "info: exiting freesession\n");