feat(bt): migrate mbedtls to PSA APIs

components/protocomm/src/crypto/srp6a/esp_srp_mpi.c

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#define MBEDTLS_DECLARE_PRIVATE_IDENTIFIERS
 #include "esp_srp_mpi.h"
 
 esp_mpi_t *esp_mpi_new(void)

components/protocomm/src/security/security1.c

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2018-2022 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2018-2025 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Apache-2.0
  */
@@ -32,6 +32,7 @@
 #include <mbedtls/ecdh.h>
 #include <mbedtls/error.h>
 #include <mbedtls/constant_time.h>
+#include "psa/crypto.h"
 
 #include <protocomm_security.h>
 #include <protocomm_security1.h>
@@ -66,7 +67,8 @@ typedef struct session {
     uint8_t rand[SZ_RANDOM];
 
     /* mbedtls context data for AES */
-    mbedtls_aes_context ctx_aes;
+    psa_cipher_operation_t ctx_aes;
+    psa_key_id_t key_id;
     unsigned char stb[16];
     size_t nc_off;
 } session_t;
@@ -94,7 +96,7 @@ static esp_err_t handle_session_command1(session_t *cur_session,
     ESP_LOGD(TAG, "Request to handle setup1_command");
     Sec1Payload *in = (Sec1Payload *) req->sec1;
     uint8_t check_buf[PUBLIC_KEY_LEN];
-    int mbed_err;
+    // int mbed_err;
 
     if (cur_session->state != SESSION_STATE_CMD1) {
         ESP_LOGE(TAG, "Invalid state of session %d (expected %d)", SESSION_STATE_CMD1, cur_session->state);
@@ -102,38 +104,49 @@ static esp_err_t handle_session_command1(session_t *cur_session,
     }
 
     /* Initialize crypto context */
-    mbedtls_aes_init(&cur_session->ctx_aes);
     memset(cur_session->stb, 0, sizeof(cur_session->stb));
     cur_session->nc_off = 0;
 
     hexdump("Client verifier", in->sc1->client_verify_data.data,
             in->sc1->client_verify_data.len);
 
-    mbed_err = mbedtls_aes_setkey_enc(&cur_session->ctx_aes, cur_session->sym_key,
-                                      sizeof(cur_session->sym_key)*8);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failure at mbedtls_aes_setkey_enc with error code : -0x%x", -mbed_err);
-        mbedtls_aes_free(&cur_session->ctx_aes);
+    psa_status_t status;
+    psa_key_id_t key_id = 0;
+    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_algorithm_t alg = PSA_ALG_CTR;
+    psa_set_key_usage_flags(&key_attributes, PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT);
+    psa_set_key_algorithm(&key_attributes, alg);
+    psa_set_key_type(&key_attributes, PSA_KEY_TYPE_AES);
+    psa_set_key_bits(&key_attributes, 128);
+    status = psa_import_key(&key_attributes, cur_session->sym_key, sizeof(cur_session->sym_key), &key_id);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_import_key failed with status=%d", status);
         return ESP_FAIL;
     }
-
-    mbed_err = mbedtls_aes_crypt_ctr(&cur_session->ctx_aes,
-                                     PUBLIC_KEY_LEN, &cur_session->nc_off,
-                                     cur_session->rand, cur_session->stb,
-                                     in->sc1->client_verify_data.data, check_buf);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failure at mbedtls_aes_crypt_ctr with error code : -0x%x", -mbed_err);
-        mbedtls_aes_free(&cur_session->ctx_aes);
+    psa_reset_key_attributes(&key_attributes);
+    status = psa_cipher_encrypt_setup(&cur_session->ctx_aes, key_id, alg);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_cipher_encrypt_setup failed with status=%d", status);
+        psa_destroy_key(key_id);
+        return ESP_FAIL;
+    }
+    size_t output_len = 0;
+    status = psa_cipher_encrypt(key_id, alg, in->sc1->client_verify_data.data,
+                                      in->sc1->client_verify_data.len, check_buf, sizeof(check_buf), &output_len);
+    if (status != PSA_SUCCESS || output_len != sizeof(check_buf)) {
+        ESP_LOGE(TAG, "psa_cipher_encrypt failed with status=%d", status);
+        psa_cipher_abort(&cur_session->ctx_aes);
+        psa_destroy_key(key_id);
         return ESP_FAIL;
     }
-
     hexdump("Dec Client verifier", check_buf, sizeof(check_buf));
 
     /* constant time memcmp */
     if (mbedtls_ct_memcmp(check_buf, cur_session->device_pubkey,
                                  sizeof(cur_session->device_pubkey)) != 0) {
         ESP_LOGE(TAG, "Key mismatch. Close connection");
-        mbedtls_aes_free(&cur_session->ctx_aes);
+        psa_cipher_abort(&cur_session->ctx_aes);
+        psa_destroy_key(key_id);
         if (esp_event_post(PROTOCOMM_SECURITY_SESSION_EVENT, PROTOCOMM_SECURITY_SESSION_CREDENTIALS_MISMATCH, NULL, 0, portMAX_DELAY) != ESP_OK) {
             ESP_LOGE(TAG, "Failed to post credential mismatch event");
         }
@@ -146,7 +159,6 @@ static esp_err_t handle_session_command1(session_t *cur_session,
         ESP_LOGE(TAG, "Error allocating memory for response1");
         free(out);
         free(out_resp);
-        mbedtls_aes_free(&cur_session->ctx_aes);
         return ESP_ERR_NO_MEM;
     }
 
@@ -159,20 +171,18 @@ static esp_err_t handle_session_command1(session_t *cur_session,
         ESP_LOGE(TAG, "Error allocating ciphertext buffer");
         free(out);
         free(out_resp);
-        mbedtls_aes_free(&cur_session->ctx_aes);
         return ESP_ERR_NO_MEM;
     }
 
-    mbed_err = mbedtls_aes_crypt_ctr(&cur_session->ctx_aes,
-                                     PUBLIC_KEY_LEN, &cur_session->nc_off,
-                                     cur_session->rand, cur_session->stb,
-                                     cur_session->client_pubkey, outbuf);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failure at mbedtls_aes_crypt_ctr with error code : -0x%x", -mbed_err);
+    status = psa_cipher_encrypt(key_id, alg, cur_session->client_pubkey,
+                                        PUBLIC_KEY_LEN, outbuf, PUBLIC_KEY_LEN, &output_len);
+    if (status != PSA_SUCCESS || output_len != PUBLIC_KEY_LEN) {
+        ESP_LOGE(TAG, "psa_cipher_encrypt failed with status=%d", status);
         free(outbuf);
         free(out);
         free(out_resp);
-        mbedtls_aes_free(&cur_session->ctx_aes);
+        psa_cipher_abort(&cur_session->ctx_aes);
+        psa_destroy_key(key_id);
         return ESP_FAIL;
     }
 
@@ -206,7 +216,6 @@ static esp_err_t handle_session_command0(session_t *cur_session,
     ESP_LOGD(TAG, "Request to handle setup0_command");
     Sec1Payload *in = (Sec1Payload *) req->sec1;
     esp_err_t ret;
-    int mbed_err;
 
     if (cur_session->state != SESSION_STATE_CMD0) {
         ESP_LOGW(TAG, "Invalid state of session %d (expected %d). Restarting session.",
@@ -233,42 +242,23 @@ static esp_err_t handle_session_command0(session_t *cur_session,
         return ESP_ERR_NO_MEM;
     }
 
-    mbedtls_ecdh_init(ctx_server);
-    mbedtls_ecdh_setup(ctx_server, MBEDTLS_ECP_DP_CURVE25519);
-    mbedtls_ctr_drbg_init(ctr_drbg);
-    mbedtls_entropy_init(entropy);
-
-    mbed_err = mbedtls_ctr_drbg_seed(ctr_drbg, mbedtls_entropy_func,
-                                     entropy, NULL, 0);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_ctr_drbg_seed with error code : -0x%x", -mbed_err);
-        ret = ESP_FAIL;
-        goto exit_cmd0;
+    psa_status_t status;
+    psa_key_id_t key_id = 0;
+    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_set_key_type(&key_attributes, PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_MONTGOMERY));
+    psa_set_key_bits(&key_attributes, 256);
+    psa_set_key_lifetime(&key_attributes, PSA_KEY_LIFETIME_VOLATILE);
+    psa_set_key_usage_flags(&key_attributes, PSA_KEY_USAGE_DERIVE | PSA_KEY_USAGE_EXPORT);
+    psa_set_key_algorithm(&key_attributes, PSA_ALG_ECDH);
+    status = psa_generate_key(&key_attributes, &key_id);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_generate_key failed with status=%d", status);
+        psa_reset_key_attributes(&key_attributes);
+        return ESP_FAIL;
     }
+    psa_reset_key_attributes(&key_attributes);
+    size_t olen = 0;
 
-    mbed_err = mbedtls_ecp_group_load(ACCESS_ECDH(&ctx_server, grp), MBEDTLS_ECP_DP_CURVE25519);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_ecp_group_load with error code : -0x%x", -mbed_err);
-        ret = ESP_FAIL;
-        goto exit_cmd0;
-    }
-
-    mbed_err = mbedtls_ecdh_gen_public(ACCESS_ECDH(&ctx_server, grp), ACCESS_ECDH(&ctx_server, d), ACCESS_ECDH(&ctx_server, Q),
-                                       mbedtls_ctr_drbg_random, ctr_drbg);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_ecdh_gen_public with error code : -0x%x", -mbed_err);
-        ret = ESP_FAIL;
-        goto exit_cmd0;
-    }
-
-    mbed_err = mbedtls_mpi_write_binary(ACCESS_ECDH(&ctx_server, Q).MBEDTLS_PRIVATE(X),
-                                        cur_session->device_pubkey,
-                                        PUBLIC_KEY_LEN);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_mpi_write_binary with error code : -0x%x", -mbed_err);
-        ret = ESP_FAIL;
-        goto exit_cmd0;
-    }
     flip_endian(cur_session->device_pubkey, PUBLIC_KEY_LEN);
 
     memcpy(cur_session->client_pubkey, in->sc0->client_pubkey.data, PUBLIC_KEY_LEN);
@@ -278,45 +268,54 @@ static esp_err_t handle_session_command0(session_t *cur_session,
     hexdump("Device pubkey", dev_pubkey, PUBLIC_KEY_LEN);
     hexdump("Client pubkey", cli_pubkey, PUBLIC_KEY_LEN);
 
-    mbed_err = mbedtls_mpi_lset(ACCESS_ECDH(&ctx_server, Qp).MBEDTLS_PRIVATE(Z), 1);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_mpi_lset with error code : -0x%x", -mbed_err);
+    status = psa_raw_key_agreement(PSA_ALG_ECDH, key_id, cur_session->client_pubkey, PUBLIC_KEY_LEN,
+                                   cur_session->sym_key, sizeof(cur_session->sym_key), &olen);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_raw_key_agreement failed with status=%d", status);
         ret = ESP_FAIL;
         goto exit_cmd0;
     }
+    if (olen != sizeof(cur_session->sym_key)) {
+        ESP_LOGE(TAG, "psa_raw_key_agreement output length mismatch: expected %zu, got %zu",
+                 sizeof(cur_session->sym_key), olen);
+        ret = ESP_FAIL;
+        goto exit_cmd0;
+    }
+    cur_session->key_id = key_id;
 
-    flip_endian(cur_session->client_pubkey, PUBLIC_KEY_LEN);
-    mbed_err = mbedtls_mpi_read_binary(ACCESS_ECDH(&ctx_server, Qp).MBEDTLS_PRIVATE(X), cli_pubkey, PUBLIC_KEY_LEN);
-    flip_endian(cur_session->client_pubkey, PUBLIC_KEY_LEN);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_mpi_read_binary with error code : -0x%x", -mbed_err);
-        ret = ESP_FAIL;
-        goto exit_cmd0;
-    }
-
-    mbed_err = mbedtls_ecdh_compute_shared(ACCESS_ECDH(&ctx_server, grp), ACCESS_ECDH(&ctx_server, z), ACCESS_ECDH(&ctx_server, Qp),
-                                           ACCESS_ECDH(&ctx_server, d), mbedtls_ctr_drbg_random, ctr_drbg);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_ecdh_compute_shared with error code : -0x%x", -mbed_err);
-        ret = ESP_FAIL;
-        goto exit_cmd0;
-    }
-
-    mbed_err = mbedtls_mpi_write_binary(ACCESS_ECDH(&ctx_server, z), cur_session->sym_key, PUBLIC_KEY_LEN);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_mpi_write_binary with error code : -0x%x", -mbed_err);
-        ret = ESP_FAIL;
-        goto exit_cmd0;
-    }
     flip_endian(cur_session->sym_key, PUBLIC_KEY_LEN);
 
     if (pop != NULL && pop->data != NULL && pop->len != 0) {
         ESP_LOGD(TAG, "Adding proof of possession");
         uint8_t sha_out[PUBLIC_KEY_LEN];
 
-        mbed_err = mbedtls_sha256((const unsigned char *) pop->data, pop->len, sha_out, 0);
-        if (mbed_err != 0) {
-            ESP_LOGE(TAG, "Failed at mbedtls_sha256_ret with error code : -0x%x", -mbed_err);
+        // mbed_err = mbedtls_sha256((const unsigned char *) pop->data, pop->len, sha_out, 0);
+        // if (mbed_err != 0) {
+        //     ESP_LOGE(TAG, "Failed at mbedtls_sha256_ret with error code : -0x%x", -mbed_err);
+        //     ret = ESP_FAIL;
+        //     goto exit_cmd0;
+        // }
+
+        psa_mac_operation_t mac_operation = PSA_MAC_OPERATION_INIT;
+        status = psa_mac_sign_setup(&mac_operation, key_id, PSA_ALG_SHA_256);
+        if (status != PSA_SUCCESS) {
+            ESP_LOGE(TAG, "psa_mac_sign_setup failed with status=%d", status);
+            ret = ESP_FAIL;
+            goto exit_cmd0;
+        }
+
+        status = psa_mac_update(&mac_operation, pop->data, pop->len);
+        if (status != PSA_SUCCESS) {
+            ESP_LOGE(TAG, "psa_mac_update failed with status=%d", status);
+            psa_mac_abort(&mac_operation);
+            ret = ESP_FAIL;
+            goto exit_cmd0;
+        }
+
+        status = psa_mac_sign_finish(&mac_operation, sha_out, sizeof(sha_out), &olen);
+        if (status != PSA_SUCCESS || olen != sizeof(sha_out)) {
+            ESP_LOGE(TAG, "psa_mac_sign_finish failed with status=%d", status);
+            psa_mac_abort(&mac_operation);
             ret = ESP_FAIL;
             goto exit_cmd0;
         }
@@ -328,9 +327,9 @@ static esp_err_t handle_session_command0(session_t *cur_session,
 
     hexdump("Shared key", cur_session->sym_key, PUBLIC_KEY_LEN);
 
-    mbed_err = mbedtls_ctr_drbg_random(ctr_drbg, cur_session->rand, SZ_RANDOM);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_ctr_drbg_random with error code : -0x%x", -mbed_err);
+    status = psa_generate_random(cur_session->rand, SZ_RANDOM);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_generate_random failed with status=%d", status);
         ret = ESP_FAIL;
         goto exit_cmd0;
     }
@@ -371,14 +370,6 @@ static esp_err_t handle_session_command0(session_t *cur_session,
     ret = ESP_OK;
 
 exit_cmd0:
-    mbedtls_ecdh_free(ctx_server);
-    free(ctx_server);
-
-    mbedtls_ctr_drbg_free(ctr_drbg);
-    free(ctr_drbg);
-
-    mbedtls_entropy_free(entropy);
-    free(entropy);
 
     return ret;
 }
@@ -460,7 +451,17 @@ static esp_err_t sec1_close_session(protocomm_security_handle_t handle, uint32_t
 
     if (cur_session->state == SESSION_STATE_DONE) {
         /* Free AES context data */
-        mbedtls_aes_free(&cur_session->ctx_aes);
+        // mbedtls_aes_free(&cur_session->ctx_aes);
+        psa_status_t status = psa_destroy_key(cur_session->id);
+        if (status != PSA_SUCCESS) {
+            ESP_LOGE(TAG, "psa_destroy_key failed with status=%d", status);
+            return ESP_FAIL;
+        }
+        status = psa_cipher_abort(&cur_session->ctx_aes);
+        if (status != PSA_SUCCESS) {
+            ESP_LOGE(TAG, "psa_cipher_abort failed with status=%d", status);
+            return ESP_FAIL;
+        }
     }
 
     memset(cur_session, 0, sizeof(session_t));
@@ -537,12 +538,17 @@ static esp_err_t sec1_decrypt(protocomm_security_handle_t handle,
         return ESP_ERR_NO_MEM;
     }
 
-    int ret = mbedtls_aes_crypt_ctr(&cur_session->ctx_aes, inlen, &cur_session->nc_off,
-                                    cur_session->rand, cur_session->stb, inbuf, *outbuf);
-    if (ret != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_aes_crypt_ctr with error code : %d", ret);
+    psa_status_t status;
+    size_t output_len = 0;
+
+    psa_algorithm_t alg = PSA_ALG_CTR;
+    status = psa_cipher_decrypt(cur_session->key_id, alg, inbuf, inlen, *outbuf, *outlen, &output_len);
+    if (status != PSA_SUCCESS || output_len != *outlen) {
+        ESP_LOGE(TAG, "psa_cipher_decrypt failed with status=%d", status);
+        free(*outbuf);
         return ESP_FAIL;
     }
+
     return ESP_OK;
 }
 

components/protocomm/src/security/security2.c

@@ -16,6 +16,7 @@
 #include <mbedtls/error.h>
 #include <mbedtls/entropy.h>
 #include <mbedtls/ctr_drbg.h>
+#include "psa/crypto.h"
 
 #include <protocomm_security.h>
 #include <protocomm_security2.h>
@@ -65,7 +66,9 @@ typedef struct session {
     uint16_t session_key_len;
     uint8_t iv[AES_GCM_IV_SIZE];
     /* mbedtls context data for AES-GCM */
-    mbedtls_gcm_context ctx_gcm;
+    // mbedtls_gcm_context ctx_gcm;
+    psa_cipher_operation_t ctx_gcm;
+    psa_key_id_t key_id;
     esp_srp_handle_t *srp_hd;
 } session_t;
 
@@ -215,7 +218,6 @@ static esp_err_t handle_session_command1(session_t *cur_session,
 {
     ESP_LOGD(TAG, "Request to handle setup1_command");
     Sec2Payload *in = (Sec2Payload *) req->sec2;
-    int mbed_err = -0x0001;
 
     if (cur_session->state != SESSION_STATE_CMD1) {
         ESP_LOGE(TAG, "Invalid state of session %d (expected %d)", SESSION_STATE_CMD1, cur_session->state);
@@ -242,24 +244,11 @@ static esp_err_t handle_session_command1(session_t *cur_session,
     }
     hexdump("Device proof", device_proof, CLIENT_PROOF_LEN);
 
-    mbedtls_entropy_context entropy;
-    mbedtls_ctr_drbg_context ctr_drbg;
-
-    mbedtls_entropy_init(&entropy);
-    mbedtls_ctr_drbg_init(&ctr_drbg);
-
-    int ret;
-    ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy, NULL, 0);
-    if (ret != 0) {
-        ESP_LOGE(TAG, "Failed to seed random number generator");
-        free(device_proof);
-        return ESP_FAIL;
-    }
-
     aes_gcm_iv_t *iv = (aes_gcm_iv_t *) cur_session->iv;
-    ret = mbedtls_ctr_drbg_random(&ctr_drbg, iv->session_id, SESSION_ID_LEN);
-    if (ret != 0) {
-        ESP_LOGE(TAG, "Failed to generate random number");
+    psa_status_t status;
+    status = psa_generate_random(iv->session_id, SESSION_ID_LEN);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_generate_random failed with status=%d", status);
         free(device_proof);
         return ESP_FAIL;
     }
@@ -269,16 +258,30 @@ static esp_err_t handle_session_command1(session_t *cur_session,
     hexdump("Initialization vector", (char *)cur_session->iv, AES_GCM_IV_SIZE);
 
     /* Initialize crypto context */
-    mbedtls_gcm_init(&cur_session->ctx_gcm);
-
-    mbed_err = mbedtls_gcm_setkey(&cur_session->ctx_gcm, MBEDTLS_CIPHER_ID_AES, (unsigned char *)cur_session->session_key, AES_GCM_KEY_LEN);
-    if (mbed_err != 0) {
-        ESP_LOGE(TAG, "Failure at mbedtls_gcm_setkey_enc with error code : -0x%x", -mbed_err);
+    cur_session->ctx_gcm = (psa_cipher_operation_t) {0};
+    psa_algorithm_t alg = PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM, AES_GCM_TAG_LEN);
+    psa_key_id_t key_id = 0;
+    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_set_key_type(&key_attributes, PSA_KEY_TYPE_AES);
+    psa_set_key_bits(&key_attributes, AES_GCM_KEY_LEN);
+    psa_set_key_usage_flags(&key_attributes, PSA_KEY_USAGE_ENCRYPT);
+    psa_set_key_algorithm(&key_attributes, alg);
+    psa_set_key_lifetime(&key_attributes, PSA_KEY_LIFETIME_VOLATILE);
+    status = psa_import_key(&key_attributes, (uint8_t *)cur_session->session_key, cur_session->session_key_len, &key_id);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_import_key failed with status=%d", status);
         free(device_proof);
-        mbedtls_gcm_free(&cur_session->ctx_gcm);
         return ESP_FAIL;
     }
 
+    status = psa_cipher_encrypt_setup(&cur_session->ctx_gcm, key_id, alg);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_cipher_encrypt_setup failed with status=%d", status);
+        free(device_proof);
+        return ESP_FAIL;
+    }
+    cur_session->key_id = key_id;
+
     Sec2Payload *out = (Sec2Payload *) malloc(sizeof(Sec2Payload));
     S2SessionResp1 *out_resp = (S2SessionResp1 *) malloc(sizeof(S2SessionResp1));
     if (!out || !out_resp) {
@@ -286,7 +289,9 @@ static esp_err_t handle_session_command1(session_t *cur_session,
         free(device_proof);
         free(out);
         free(out_resp);
-        mbedtls_gcm_free(&cur_session->ctx_gcm);
+        psa_cipher_abort(&cur_session->ctx_gcm);
+        psa_destroy_key(key_id);
+        // mbedtls_gcm_free(&cur_session->ctx_gcm);
         return ESP_ERR_NO_MEM;
     }
 
@@ -391,7 +396,14 @@ static esp_err_t sec2_close_session(protocomm_security_handle_t handle, uint32_t
 
     if (cur_session->state == SESSION_STATE_DONE) {
         /* Free GCM context data */
-        mbedtls_gcm_free(&cur_session->ctx_gcm);
+        // mbedtls_gcm_free(&cur_session->ctx_gcm);
+        psa_status_t status = psa_cipher_abort(&cur_session->ctx_gcm);
+        if (status != PSA_SUCCESS) {
+            ESP_LOGE(TAG, "psa_cipher_abort failed with status=%d", status);
+            return ESP_FAIL;
+        }
+        psa_destroy_key(cur_session->key_id);
+        cur_session->key_id = 0;
     }
 
     free(cur_session->username);
@@ -482,13 +494,41 @@ static esp_err_t sec2_encrypt(protocomm_security_handle_t handle,
     }
     uint8_t gcm_tag[AES_GCM_TAG_LEN];
 
-    int ret = mbedtls_gcm_crypt_and_tag(&cur_session->ctx_gcm, MBEDTLS_GCM_ENCRYPT, inlen, cur_session->iv,
-                                        AES_GCM_IV_SIZE, NULL, 0, inbuf,
-                                        *outbuf, AES_GCM_TAG_LEN, gcm_tag);
-    if (ret != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_gcm_crypt_and_tag with error code : %d", ret);
+    psa_status_t status;
+    status = psa_cipher_set_iv(&cur_session->ctx_gcm, cur_session->iv, AES_GCM_IV_SIZE);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_cipher_set_iv failed with status=%d", status);
+        free(*outbuf);
         return ESP_FAIL;
     }
+
+    size_t out_len = 0;
+    status = psa_cipher_update(&cur_session->ctx_gcm, inbuf, inlen, *outbuf, *outlen , &out_len);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_cipher_update failed with status=%d", status);
+        free(*outbuf);
+        return ESP_FAIL;
+    }
+
+    if (out_len != inlen) {
+        ESP_LOGE(TAG, "psa_cipher_update output length mismatch: expected %zd, got %zu", inlen, out_len);
+        free(*outbuf);
+        return ESP_FAIL;
+    }
+
+    status = psa_cipher_finish(&cur_session->ctx_gcm, gcm_tag, AES_GCM_TAG_LEN, &out_len);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_cipher_finish failed with status=%d", status);
+        free(*outbuf);
+        return ESP_FAIL;
+    }
+
+    if (out_len != AES_GCM_TAG_LEN) {
+        ESP_LOGE(TAG, "psa_cipher_finish output length mismatch: expected %d, got %zu", AES_GCM_TAG_LEN, out_len);
+        free(*outbuf);
+        return ESP_FAIL;
+    }
+
     memcpy(*outbuf + inlen, gcm_tag, AES_GCM_TAG_LEN);
 
     /* Increment counter value for next operation */
@@ -531,13 +571,48 @@ static esp_err_t sec2_decrypt(protocomm_security_handle_t handle,
         return ESP_ERR_NO_MEM;
     }
 
-    int ret = mbedtls_gcm_auth_decrypt(&cur_session->ctx_gcm, inlen - AES_GCM_TAG_LEN, cur_session->iv,
-                                       AES_GCM_IV_SIZE, NULL, 0, inbuf + (inlen - AES_GCM_TAG_LEN), AES_GCM_TAG_LEN, inbuf, *outbuf);
-    if (ret != 0) {
-        ESP_LOGE(TAG, "Failed at mbedtls_gcm_auth_decrypt : %d", ret);
+    psa_status_t status;
+    status = psa_cipher_set_iv(&cur_session->ctx_gcm, cur_session->iv, AES_GCM_IV_SIZE);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_cipher_set_iv failed with status=%d", status);
+        free(*outbuf);
         return ESP_FAIL;
     }
 
+    size_t out_len = 0;
+    status = psa_cipher_update(&cur_session->ctx_gcm, inbuf, inlen - AES_GCM_TAG_LEN, *outbuf, *outlen, &out_len);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_cipher_update failed with status=%d", status);
+        free(*outbuf);
+        return ESP_FAIL;
+    }
+
+    if (out_len != *outlen) {
+        ESP_LOGE(TAG, "psa_cipher_update output length mismatch: expected %zd, got %zu", *outlen, out_len);
+        free(*outbuf);
+        return ESP_FAIL;
+    }
+
+    uint8_t gcm_tag[AES_GCM_TAG_LEN];
+    memcpy(gcm_tag, inbuf + (inlen - AES_GCM_TAG_LEN), AES_GCM_TAG_LEN);
+    status = psa_cipher_finish(&cur_session->ctx_gcm, gcm_tag, AES_GCM_TAG_LEN, &out_len);
+    if (status != PSA_SUCCESS) {
+        ESP_LOGE(TAG, "psa_cipher_finish failed with status=%d", status);
+        free(*outbuf);
+        return ESP_FAIL;
+    }
+
+    if (out_len != 0) {
+        ESP_LOGE(TAG, "psa_cipher_finish output length mismatch: expected 0, got %zu", out_len);
+        free(*outbuf);
+        return ESP_FAIL;
+    }
+
+    if (*outbuf == NULL) {
+        ESP_LOGE(TAG, "Output buffer is NULL");
+        return ESP_ERR_INVALID_ARG;
+    }
+
     /* Increment counter value for next operation */
     sec2_gcm_iv_counter_increment(cur_session->iv);