Adapt names to more suitable and fix conditional compilation flags

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c57583a..2c1a910 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7407,7 +7407,7 @@
 }
 
 /* Auxiliary function to convert core computation stage(step, sequence, state) to single driver step. */
-static psa_pake_driver_step_t convert_jpake_computation_stage_to_driver_step(
+static psa_crypto_driver_pake_step_t convert_jpake_computation_stage_to_driver_step(
     psa_jpake_computation_stage_t *stage)
 {
     switch (stage->state) {
@@ -7843,7 +7843,7 @@
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_status_t abort_status = PSA_ERROR_CORRUPTION_DETECTED;
-    uint8_t shared_key[MBEDTLS_PSA_PAKE_BUFFER_SIZE];
+    uint8_t shared_key[MBEDTLS_PSA_JPAKE_BUFFER_SIZE];
     size_t shared_key_len = 0;
 
     if (operation->stage != PSA_PAKE_OPERATION_STAGE_COMPUTATION) {
@@ -7874,7 +7874,7 @@
                                             shared_key,
                                             shared_key_len);
 
-    mbedtls_platform_zeroize(shared_key, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+    mbedtls_platform_zeroize(shared_key, MBEDTLS_PSA_JPAKE_BUFFER_SIZE);
 exit:
     abort_status = psa_pake_abort(operation);
     return status == PSA_SUCCESS ? abort_status : status;
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index 11a95e3..65d0d3f 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -421,14 +421,14 @@
 
 psa_status_t psa_driver_wrapper_pake_output(
     psa_pake_operation_t *operation,
-    psa_pake_driver_step_t step,
+    psa_crypto_driver_pake_step_t step,
     uint8_t *output,
     size_t output_size,
     size_t *output_length);
 
 psa_status_t psa_driver_wrapper_pake_input(
     psa_pake_operation_t *operation,
-    psa_pake_driver_step_t step,
+    psa_crypto_driver_pake_step_t step,
     const uint8_t *input,
     size_t input_length);
 
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index fdfbd16..73032c6 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -163,6 +163,7 @@
 }
 #endif
 
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
 static psa_status_t psa_pake_ecjpake_setup(mbedtls_psa_pake_operation_t *operation)
 {
@@ -187,6 +188,7 @@
 
     return PSA_SUCCESS;
 }
+#endif
 
 psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
                                     const psa_crypto_driver_pake_inputs_t *inputs)
@@ -237,7 +239,7 @@
         operation->role = role;
         operation->alg = cipher_suite.algorithm;
 
-        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_JPAKE_BUFFER_SIZE);
         operation->buffer_length = 0;
         operation->buffer_offset = 0;
 
@@ -259,7 +261,7 @@
 
 static psa_status_t mbedtls_psa_pake_output_internal(
     mbedtls_psa_pake_operation_t *operation,
-    psa_pake_driver_step_t step,
+    psa_crypto_driver_pake_step_t step,
     uint8_t *output,
     size_t output_size,
     size_t *output_length)
@@ -288,7 +290,7 @@
         if (step == PSA_JPAKE_X1_STEP_KEY_SHARE) {
             ret = mbedtls_ecjpake_write_round_one(&operation->ctx.pake,
                                                   operation->buffer,
-                                                  MBEDTLS_PSA_PAKE_BUFFER_SIZE,
+                                                  MBEDTLS_PSA_JPAKE_BUFFER_SIZE,
                                                   &operation->buffer_length,
                                                   mbedtls_psa_get_random,
                                                   MBEDTLS_PSA_RANDOM_STATE);
@@ -300,7 +302,7 @@
         } else if (step == PSA_JPAKE_X2S_STEP_KEY_SHARE) {
             ret = mbedtls_ecjpake_write_round_two(&operation->ctx.pake,
                                                   operation->buffer,
-                                                  MBEDTLS_PSA_PAKE_BUFFER_SIZE,
+                                                  MBEDTLS_PSA_JPAKE_BUFFER_SIZE,
                                                   &operation->buffer_length,
                                                   mbedtls_psa_get_random,
                                                   MBEDTLS_PSA_RANDOM_STATE);
@@ -350,7 +352,7 @@
         /* Reset buffer after ZK_PROOF sequence */
         if ((step == PSA_JPAKE_X2_STEP_ZK_PROOF) ||
             (step == PSA_JPAKE_X2S_STEP_ZK_PROOF)) {
-            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_JPAKE_BUFFER_SIZE);
             operation->buffer_length = 0;
             operation->buffer_offset = 0;
         }
@@ -367,7 +369,7 @@
 }
 
 psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
-                                     psa_pake_driver_step_t step,
+                                     psa_crypto_driver_pake_step_t step,
                                      uint8_t *output,
                                      size_t output_size,
                                      size_t *output_length)
@@ -380,7 +382,7 @@
 
 static psa_status_t mbedtls_psa_pake_input_internal(
     mbedtls_psa_pake_operation_t *operation,
-    psa_pake_driver_step_t step,
+    psa_crypto_driver_pake_step_t step,
     const uint8_t *input,
     size_t input_length)
 {
@@ -441,7 +443,7 @@
                                                  operation->buffer,
                                                  operation->buffer_length);
 
-            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_JPAKE_BUFFER_SIZE);
             operation->buffer_length = 0;
 
             if (ret != 0) {
@@ -452,7 +454,7 @@
                                                  operation->buffer,
                                                  operation->buffer_length);
 
-            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_JPAKE_BUFFER_SIZE);
             operation->buffer_length = 0;
 
             if (ret != 0) {
@@ -471,7 +473,7 @@
 }
 
 psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
-                                    psa_pake_driver_step_t step,
+                                    psa_crypto_driver_pake_step_t step,
                                     const uint8_t *input,
                                     size_t input_length)
 {
@@ -491,7 +493,7 @@
     if (operation->alg == PSA_ALG_JPAKE) {
         ret = mbedtls_ecjpake_write_shared_key(&operation->ctx.pake,
                                                operation->buffer,
-                                               MBEDTLS_PSA_PAKE_BUFFER_SIZE,
+                                               MBEDTLS_PSA_JPAKE_BUFFER_SIZE,
                                                &operation->buffer_length,
                                                mbedtls_psa_get_random,
                                                MBEDTLS_PSA_RANDOM_STATE);
@@ -520,7 +522,7 @@
         operation->password = NULL;
         operation->password_len = 0;
         operation->role = PSA_PAKE_ROLE_NONE;
-        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_JPAKE_BUFFER_SIZE);
         operation->buffer_length = 0;
         operation->buffer_offset = 0;
         mbedtls_ecjpake_free(&operation->ctx.pake);
diff --git a/library/psa_crypto_pake.h b/library/psa_crypto_pake.h
index dc6ad7b..3658556 100644
--- a/library/psa_crypto_pake.h
+++ b/library/psa_crypto_pake.h
@@ -96,7 +96,7 @@
  *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
-                                     psa_pake_driver_step_t step,
+                                     psa_crypto_driver_pake_step_t step,
                                      uint8_t *output,
                                      size_t output_size,
                                      size_t *output_length);
@@ -143,7 +143,7 @@
  *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
-                                    psa_pake_driver_step_t step,
+                                    psa_crypto_driver_pake_step_t step,
                                     const uint8_t *input,
                                     size_t input_length);