Documentation fixes and code adaptation

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/docs/proposed/psa-driver-interface.md b/docs/proposed/psa-driver-interface.md
index 07f1989..ac6b8de 100644
--- a/docs/proposed/psa-driver-interface.md
+++ b/docs/proposed/psa-driver-interface.md
@@ -463,12 +463,14 @@
 ```
 psa_status_t acme_pake_get_implicit_key(
                             acme_pake_operation_t *operation,
-                            uint8_t *output, size_t *output_size );
+                            uint8_t *output, size_t output_size,
+                            size_t *output_length );
 ```
 
-* `operation` is an operation object
-* `output` output buffer for implicit key
-* `output_size` size of the returned implicit key
+* `operation` The driver PAKE operation object to use.
+* `output` Buffer where the implicit key is to be written.
+* `output_size` Size of the output buffer in bytes.
+* `output_length` On success, the number of bytes of the implicit key.
 
 ### Driver entry points for key management
 
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 39ef52c..5f86c3f 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -1318,8 +1318,8 @@
  *
  * \param[in]  inputs           Operation inputs.
  * \param[out] buffer           Return buffer for password.
- * \param[in]  buffer_size      Size of the return buffer in bytes.
- * \param[in]  buffer_length    Actual size of the password in bytes.
+ * \param      buffer_size      Size of the return buffer in bytes.
+ * \param[out] buffer_length    Actual size of the password in bytes.
  *
  * \retval #PSA_SUCCESS
  *         Success.
@@ -2034,11 +2034,10 @@
      * ID value zero means the context is not valid or not assigned to
      * any driver (i.e. none of the driver contexts are active). */
     unsigned int MBEDTLS_PRIVATE(id);
-    /* Algorithm used for PAKE operation */
+    /* Algorithm of the PAKE operation */
     psa_algorithm_t MBEDTLS_PRIVATE(alg);
-    /* Based on stage (collecting inputs/computation) we select active structure of data union.
-     * While switching stage (when driver setup is called) collected inputs
-       are copied to the corresponding operation context. */
+    /* Stage of the PAKE operation: waiting for the setup, collecting inputs
+     * or computing. */
     uint8_t MBEDTLS_PRIVATE(stage);
     /* Holds computation stage of the PAKE algorithms. */
     union {
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 2c1a910..1c066ce 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7863,6 +7863,7 @@
 
     status = psa_driver_wrapper_pake_get_implicit_key(operation,
                                                       shared_key,
+                                                      sizeof(shared_key),
                                                       &shared_key_len);
 
     if (status != PSA_SUCCESS) {
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index 65d0d3f..dd1b763 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -434,7 +434,8 @@
 
 psa_status_t psa_driver_wrapper_pake_get_implicit_key(
     psa_pake_operation_t *operation,
-    uint8_t *output, size_t *output_size);
+    uint8_t *output, size_t output_size,
+    size_t *output_length);
 
 psa_status_t psa_driver_wrapper_pake_abort(
     psa_pake_operation_t *operation);
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 73032c6..150270c 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -485,7 +485,8 @@
 
 psa_status_t mbedtls_psa_pake_get_implicit_key(
     mbedtls_psa_pake_operation_t *operation,
-    uint8_t *output, size_t *output_size)
+    uint8_t *output, size_t output_size,
+    size_t *output_length)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
@@ -493,7 +494,7 @@
     if (operation->alg == PSA_ALG_JPAKE) {
         ret = mbedtls_ecjpake_write_shared_key(&operation->ctx.pake,
                                                operation->buffer,
-                                               MBEDTLS_PSA_JPAKE_BUFFER_SIZE,
+                                               output_size,
                                                &operation->buffer_length,
                                                mbedtls_psa_get_random,
                                                MBEDTLS_PSA_RANDOM_STATE);
@@ -502,7 +503,7 @@
         }
 
         memcpy(output, operation->buffer, operation->buffer_length);
-        *output_size = operation->buffer_length;
+        *output_length = operation->buffer_length;
 
         return PSA_SUCCESS;
     } else
diff --git a/library/psa_crypto_pake.h b/library/psa_crypto_pake.h
index 3658556..9bdcc33 100644
--- a/library/psa_crypto_pake.h
+++ b/library/psa_crypto_pake.h
@@ -43,6 +43,8 @@
  *         compatible with the PAKE algorithm, or the hash algorithm in
  *         \p cipher_suite is not supported or not compatible with the PAKE
  *         algorithm and primitive.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
  */
 psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
                                     const psa_crypto_driver_pake_inputs_t *inputs);
@@ -59,10 +61,9 @@
  * \param step                 The step of the algorithm for which the output is
  *                             requested.
  * \param[out] output          Buffer where the output is to be written in the
- *                             format appropriate for this \p step. Refer to
- *                             the documentation of the individual
- *                             \c PSA_PAKE_STEP_XXX constants for more
- *                             information.
+ *                             format appropriate for this driver \p step. Refer to
+ *                             the documentation of psa_crypto_driver_pake_step_t for
+ *                             more information.
  * \param output_size          Size of the \p output buffer in bytes. This must
  *                             be at least #PSA_PAKE_OUTPUT_SIZE(\p alg, \p
  *                             primitive, \p step) where \p alg and
@@ -77,23 +78,10 @@
  *         Success.
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
  *         The size of the \p output buffer is too small.
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p step is not compatible with the operation's algorithm.
- * \retval #PSA_ERROR_NOT_SUPPORTED
- *         \p step is not supported with the operation's algorithm.
  * \retval #PSA_ERROR_INSUFFICIENT_ENTROPY
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_DATA_INVALID
- * \retval #PSA_ERROR_BAD_STATE
- *         The operation state is not valid (it must be active, and fully set
- *         up, and this call must conform to the algorithm's requirements
- *         for ordering of input and output steps).
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
                                      psa_crypto_driver_pake_step_t step,
@@ -104,43 +92,32 @@
 /** Provide input for a step of a password-authenticated key exchange.
  *
  * \note The signature of this function is that of a PSA driver
- *       key_agreement entry point. This function behaves as a key_agreement
+ *       pake_input entry point. This function behaves as a pake_input
  *       entry point as defined in the PSA driver interface specification for
  *       transparent drivers.
  *
  * \param[in,out] operation    Active PAKE operation.
- * \param step                 The step for which the input is provided.
+ * \param step                 The driver step for which the input is provided.
  * \param[in] input            Buffer containing the input in the format
  *                             appropriate for this \p step. Refer to the
- *                             documentation of the individual
- *                             \c PSA_PAKE_STEP_XXX constants for more
- *                             information.
+ *                             documentation of psa_crypto_driver_pake_step_t
+ *                             for more information.
  * \param input_length         Size of the \p input buffer in bytes.
  *
  * \retval #PSA_SUCCESS
  *         Success.
  * \retval #PSA_ERROR_INVALID_SIGNATURE
- *         The verification fails for a #PSA_PAKE_STEP_ZK_PROOF input step.
+ *         The verification fails for a zero-knowledge input step.
  * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p step is not compatible with the \p operation’s algorithm, or the
- *         \p input is not valid for the \p operation's algorithm, cipher suite
+ *         the \p input is not valid for the \p operation's algorithm, cipher suite
  *         or \p step.
  * \retval #PSA_ERROR_NOT_SUPPORTED
- *         \p step p is not supported with the \p operation's algorithm, or the
- *         \p input is not supported for the \p operation's algorithm, cipher
+ *         the \p input is not supported for the \p operation's algorithm, cipher
  *         suite or \p step.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_DATA_INVALID
- * \retval #PSA_ERROR_BAD_STATE
- *         The operation state is not valid (it must be active, and fully set
- *         up, and this call must conform to the algorithm's requirements
- *         for ordering of input and output steps).
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
                                     psa_crypto_driver_pake_step_t step,
@@ -155,8 +132,9 @@
  *       interface specification for transparent drivers.
  *
  * \param[in,out] operation    Active PAKE operation.
- * \param[out] output          Output buffer for implicit key
- * \param[out] output_size     Size of the returned implicit key
+ * \param[out] output          Output buffer for implicit key.
+ * \param      output_size     Size of the output buffer in bytes.
+ * \param[out] output_length   On success, the number of bytes of the implicit key.
  *
  * \retval #PSA_SUCCESS
  *         Success.
@@ -164,24 +142,14 @@
  *         Input from a PAKE is not supported by the algorithm in the \p output
  *         key derivation operation.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_DATA_INVALID
- * \retval #PSA_ERROR_BAD_STATE
- *         The PAKE operation state is not valid (it must be active, but beyond
- *         that validity is specific to the algorithm),
- *         or the state of \p output is not valid for
- *         the #PSA_KEY_DERIVATION_INPUT_SECRET step. This can happen if the
- *         step is out of order or the application has done this step already
- *         and it may not be repeated.
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_get_implicit_key(
     mbedtls_psa_pake_operation_t *operation,
-    uint8_t *output, size_t *output_size);
+    uint8_t *output, size_t output_size,
+    size_t *output_length);
 
 /** Abort a PAKE operation.
  *
@@ -194,11 +162,7 @@
  *
  * \retval #PSA_SUCCESS
  *         Success.
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_BAD_STATE
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_abort(mbedtls_psa_pake_operation_t *operation);
 
diff --git a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja
index b287b37..a34d9b0 100644
--- a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja
+++ b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja
@@ -2938,13 +2938,15 @@
 
 psa_status_t psa_driver_wrapper_pake_get_implicit_key(
     psa_pake_operation_t *operation,
-    uint8_t *output, size_t *output_size )
+    uint8_t *output, size_t output_size,
+    size_t *output_length )
 {
     switch( operation->id )
     {
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_pake_get_implicit_key( &operation->data.ctx.mbedtls_ctx, output, output_size ) );
+            return( mbedtls_psa_pake_get_implicit_key( &operation->data.ctx.mbedtls_ctx,
+                                                       output, output_size, output_length ) );
 #endif /* MBEDTLS_PSA_BUILTIN_PAKE */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
@@ -2952,11 +2954,11 @@
         case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
             return( mbedtls_test_transparent_pake_get_implicit_key(
                         &operation->data.ctx.transparent_test_driver_ctx,
-                        output, output_size ) );
+                        output, output_size, output_length ) );
         case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
             return( mbedtls_test_opaque_pake_get_implicit_key(
                         &operation->data.ctx.opaque_test_driver_ctx,
-                        output, output_size ) );
+                        output, output_size, output_length ) );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
         default:
diff --git a/tests/include/test/drivers/pake.h b/tests/include/test/drivers/pake.h
index d082d6e..4a2b7c4 100644
--- a/tests/include/test/drivers/pake.h
+++ b/tests/include/test/drivers/pake.h
@@ -70,7 +70,7 @@
 
 psa_status_t mbedtls_test_transparent_pake_get_implicit_key(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
-    uint8_t *output, size_t *output_size);
+    uint8_t *output, size_t output_size, size_t *output_length);
 
 psa_status_t mbedtls_test_transparent_pake_abort(
     mbedtls_transparent_test_driver_pake_operation_t *operation);
@@ -114,7 +114,7 @@
 
 psa_status_t mbedtls_test_opaque_pake_get_implicit_key(
     mbedtls_opaque_test_driver_pake_operation_t *operation,
-    uint8_t *output, size_t *output_size);
+    uint8_t *output, size_t output_size, size_t *output_length);
 
 psa_status_t mbedtls_test_opaque_pake_abort(
     mbedtls_opaque_test_driver_pake_operation_t *operation);
diff --git a/tests/src/drivers/test_driver_pake.c b/tests/src/drivers/test_driver_pake.c
index 615f7ef..3eaf38a 100644
--- a/tests/src/drivers/test_driver_pake.c
+++ b/tests/src/drivers/test_driver_pake.c
@@ -145,7 +145,7 @@
 
 psa_status_t mbedtls_test_transparent_pake_get_implicit_key(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
-    uint8_t *output, size_t *output_size)
+    uint8_t *output, size_t output_size, size_t *output_length)
 {
     mbedtls_test_driver_pake_hooks.hits++;
 
@@ -157,11 +157,11 @@
         defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
         mbedtls_test_driver_pake_hooks.driver_status =
             libtestdriver1_mbedtls_psa_pake_get_implicit_key(
-                operation,  output, output_size);
+                operation,  output, output_size, output_length);
 #elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
         mbedtls_test_driver_pake_hooks.driver_status =
             mbedtls_psa_pake_get_implicit_key(
-                operation, output, output_size);
+                operation, output, output_size, output_length);
 #else
         (void) operation;
         (void) output;
@@ -289,11 +289,12 @@
 
 psa_status_t mbedtls_test_opaque_pake_get_implicit_key(
     mbedtls_opaque_test_driver_pake_operation_t *operation,
-    uint8_t *output, size_t *output_size)
+    uint8_t *output, size_t output_size, size_t *output_length)
 {
     (void) operation;
     (void) output;
     (void) output_size;
+    (void) output_length;
     return PSA_ERROR_NOT_SUPPORTED;
 }