Update generated PSA wrappers

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/tests/include/test/psa_test_wrappers.h b/tests/include/test/psa_test_wrappers.h
index 134a547..b83f405 100644
--- a/tests/include/test/psa_test_wrappers.h
+++ b/tests/include/test/psa_test_wrappers.h
@@ -363,14 +363,6 @@
 #define psa_generate_key_custom(arg0_attributes, arg1_custom, arg2_custom_data, arg3_custom_data_length, arg4_key) \
     mbedtls_test_wrap_psa_generate_key_custom(arg0_attributes, arg1_custom, arg2_custom_data, arg3_custom_data_length, arg4_key)
 
-psa_status_t mbedtls_test_wrap_psa_generate_key_ext(
-    const psa_key_attributes_t *arg0_attributes,
-    const psa_key_production_parameters_t *arg1_params,
-    size_t arg2_params_data_length,
-    mbedtls_svc_key_id_t *arg3_key);
-#define psa_generate_key_ext(arg0_attributes, arg1_params, arg2_params_data_length, arg3_key) \
-    mbedtls_test_wrap_psa_generate_key_ext(arg0_attributes, arg1_params, arg2_params_data_length, arg3_key)
-
 psa_status_t mbedtls_test_wrap_psa_generate_random(
     uint8_t *arg0_output,
     size_t arg1_output_size);
@@ -515,15 +507,6 @@
 #define psa_key_derivation_output_key_custom(arg0_attributes, arg1_operation, arg2_custom, arg3_custom_data, arg4_custom_data_length, arg5_key) \
     mbedtls_test_wrap_psa_key_derivation_output_key_custom(arg0_attributes, arg1_operation, arg2_custom, arg3_custom_data, arg4_custom_data_length, arg5_key)
 
-psa_status_t mbedtls_test_wrap_psa_key_derivation_output_key_ext(
-    const psa_key_attributes_t *arg0_attributes,
-    psa_key_derivation_operation_t *arg1_operation,
-    const psa_key_production_parameters_t *arg2_params,
-    size_t arg3_params_data_length,
-    mbedtls_svc_key_id_t *arg4_key);
-#define psa_key_derivation_output_key_ext(arg0_attributes, arg1_operation, arg2_params, arg3_params_data_length, arg4_key) \
-    mbedtls_test_wrap_psa_key_derivation_output_key_ext(arg0_attributes, arg1_operation, arg2_params, arg3_params_data_length, arg4_key)
-
 psa_status_t mbedtls_test_wrap_psa_key_derivation_set_capacity(
     psa_key_derivation_operation_t *arg0_operation,
     size_t arg1_capacity);
diff --git a/tests/psa-client-server/psasim/src/psa_functions_codes.h b/tests/psa-client-server/psasim/src/psa_functions_codes.h
index bc1b844..1301ff2 100644
--- a/tests/psa-client-server/psasim/src/psa_functions_codes.h
+++ b/tests/psa-client-server/psasim/src/psa_functions_codes.h
@@ -40,7 +40,7 @@
     PSA_EXPORT_KEY,
     PSA_EXPORT_PUBLIC_KEY,
     PSA_GENERATE_KEY,
-    PSA_GENERATE_KEY_EXT,
+    PSA_GENERATE_KEY_CUSTOM,
     PSA_GENERATE_RANDOM,
     PSA_GET_KEY_ATTRIBUTES,
     PSA_HASH_ABORT,
@@ -62,7 +62,7 @@
     PSA_KEY_DERIVATION_KEY_AGREEMENT,
     PSA_KEY_DERIVATION_OUTPUT_BYTES,
     PSA_KEY_DERIVATION_OUTPUT_KEY,
-    PSA_KEY_DERIVATION_OUTPUT_KEY_EXT,
+    PSA_KEY_DERIVATION_OUTPUT_KEY_CUSTOM,
     PSA_KEY_DERIVATION_SET_CAPACITY,
     PSA_KEY_DERIVATION_SETUP,
     PSA_MAC_ABORT,
diff --git a/tests/psa-client-server/psasim/src/psa_sim_crypto_client.c b/tests/psa-client-server/psasim/src/psa_sim_crypto_client.c
index 4200f6c..d4d9d60 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_crypto_client.c
+++ b/tests/psa-client-server/psasim/src/psa_sim_crypto_client.c
@@ -2803,9 +2803,10 @@
 }
 
 
-psa_status_t psa_generate_key_ext(
+psa_status_t psa_generate_key_custom(
     const psa_key_attributes_t *attributes,
-    const psa_key_production_parameters_t *params, size_t  params_data_length,
+    const psa_custom_key_parameters_t *custom,
+    const uint8_t *custom_data, size_t  custom_data_length,
     mbedtls_svc_key_id_t *key
     )
 {
@@ -2817,7 +2818,8 @@
     size_t needed =
         psasim_serialise_begin_needs() +
         psasim_serialise_psa_key_attributes_t_needs(*attributes) +
-        psasim_serialise_psa_key_production_parameters_t_needs(params, params_data_length) +
+        psasim_serialise_psa_custom_key_parameters_t_needs(*custom) +
+        psasim_serialise_buffer_needs(custom_data, custom_data_length) +
         psasim_serialise_mbedtls_svc_key_id_t_needs(*key);
 
     ser_params = malloc(needed);
@@ -2839,9 +2841,15 @@
     if (!ok) {
         goto fail;
     }
-    ok = psasim_serialise_psa_key_production_parameters_t(
+    ok = psasim_serialise_psa_custom_key_parameters_t(
         &pos, &remaining,
-        params, params_data_length);
+        *custom);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_buffer(
+        &pos, &remaining,
+        custom_data, custom_data_length);
     if (!ok) {
         goto fail;
     }
@@ -2852,10 +2860,10 @@
         goto fail;
     }
 
-    ok = psa_crypto_call(PSA_GENERATE_KEY_EXT,
+    ok = psa_crypto_call(PSA_GENERATE_KEY_CUSTOM,
                          ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
     if (!ok) {
-        printf("PSA_GENERATE_KEY_EXT server call failed\n");
+        printf("PSA_GENERATE_KEY_CUSTOM server call failed\n");
         goto fail;
     }
 
@@ -4572,10 +4580,11 @@
 }
 
 
-psa_status_t psa_key_derivation_output_key_ext(
+psa_status_t psa_key_derivation_output_key_custom(
     const psa_key_attributes_t *attributes,
     psa_key_derivation_operation_t *operation,
-    const psa_key_production_parameters_t *params, size_t  params_data_length,
+    const psa_custom_key_parameters_t *custom,
+    const uint8_t *custom_data, size_t  custom_data_length,
     mbedtls_svc_key_id_t *key
     )
 {
@@ -4588,7 +4597,8 @@
         psasim_serialise_begin_needs() +
         psasim_serialise_psa_key_attributes_t_needs(*attributes) +
         psasim_serialise_psa_key_derivation_operation_t_needs(*operation) +
-        psasim_serialise_psa_key_production_parameters_t_needs(params, params_data_length) +
+        psasim_serialise_psa_custom_key_parameters_t_needs(*custom) +
+        psasim_serialise_buffer_needs(custom_data, custom_data_length) +
         psasim_serialise_mbedtls_svc_key_id_t_needs(*key);
 
     ser_params = malloc(needed);
@@ -4616,9 +4626,15 @@
     if (!ok) {
         goto fail;
     }
-    ok = psasim_serialise_psa_key_production_parameters_t(
+    ok = psasim_serialise_psa_custom_key_parameters_t(
         &pos, &remaining,
-        params, params_data_length);
+        *custom);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_buffer(
+        &pos, &remaining,
+        custom_data, custom_data_length);
     if (!ok) {
         goto fail;
     }
@@ -4629,10 +4645,10 @@
         goto fail;
     }
 
-    ok = psa_crypto_call(PSA_KEY_DERIVATION_OUTPUT_KEY_EXT,
+    ok = psa_crypto_call(PSA_KEY_DERIVATION_OUTPUT_KEY_CUSTOM,
                          ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
     if (!ok) {
-        printf("PSA_KEY_DERIVATION_OUTPUT_KEY_EXT server call failed\n");
+        printf("PSA_KEY_DERIVATION_OUTPUT_KEY_CUSTOM server call failed\n");
         goto fail;
     }
 
diff --git a/tests/psa-client-server/psasim/src/psa_sim_crypto_server.c b/tests/psa-client-server/psasim/src/psa_sim_crypto_server.c
index cab32c4..b2ed070 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_crypto_server.c
+++ b/tests/psa-client-server/psasim/src/psa_sim_crypto_server.c
@@ -3116,14 +3116,15 @@
 }
 
 // Returns 1 for success, 0 for failure
-int psa_generate_key_ext_wrapper(
+int psa_generate_key_custom_wrapper(
     uint8_t *in_params, size_t in_params_len,
     uint8_t **out_params, size_t *out_params_len)
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_attributes_t attributes;
-    psa_key_production_parameters_t *params = NULL;
-    size_t params_data_length;
+    psa_custom_key_parameters_t custom;
+    uint8_t *custom_data = NULL;
+    size_t custom_data_length;
     mbedtls_svc_key_id_t key;
 
     uint8_t *pos = in_params;
@@ -3143,9 +3144,16 @@
         goto fail;
     }
 
-    ok = psasim_deserialise_psa_key_production_parameters_t(
+    ok = psasim_deserialise_psa_custom_key_parameters_t(
         &pos, &remaining,
-        &params, &params_data_length);
+        &custom);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_buffer(
+        &pos, &remaining,
+        &custom_data, &custom_data_length);
     if (!ok) {
         goto fail;
     }
@@ -3159,9 +3167,10 @@
 
     // Now we call the actual target function
 
-    status = psa_generate_key_ext(
+    status = psa_generate_key_custom(
         &attributes,
-        params, params_data_length,
+        &custom,
+        custom_data, custom_data_length,
         &key
         );
 
@@ -3201,14 +3210,14 @@
     *out_params = result;
     *out_params_len = result_size;
 
-    free(params);
+    free(custom_data);
 
     return 1;   // success
 
 fail:
     free(result);
 
-    free(params);
+    free(custom_data);
 
     return 0;       // This shouldn't happen!
 }
@@ -5079,15 +5088,16 @@
 }
 
 // Returns 1 for success, 0 for failure
-int psa_key_derivation_output_key_ext_wrapper(
+int psa_key_derivation_output_key_custom_wrapper(
     uint8_t *in_params, size_t in_params_len,
     uint8_t **out_params, size_t *out_params_len)
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_attributes_t attributes;
     psa_key_derivation_operation_t *operation;
-    psa_key_production_parameters_t *params = NULL;
-    size_t params_data_length;
+    psa_custom_key_parameters_t custom;
+    uint8_t *custom_data = NULL;
+    size_t custom_data_length;
     mbedtls_svc_key_id_t key;
 
     uint8_t *pos = in_params;
@@ -5114,9 +5124,16 @@
         goto fail;
     }
 
-    ok = psasim_deserialise_psa_key_production_parameters_t(
+    ok = psasim_deserialise_psa_custom_key_parameters_t(
         &pos, &remaining,
-        &params, &params_data_length);
+        &custom);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_buffer(
+        &pos, &remaining,
+        &custom_data, &custom_data_length);
     if (!ok) {
         goto fail;
     }
@@ -5130,10 +5147,11 @@
 
     // Now we call the actual target function
 
-    status = psa_key_derivation_output_key_ext(
+    status = psa_key_derivation_output_key_custom(
         &attributes,
         operation,
-        params, params_data_length,
+        &custom,
+        custom_data, custom_data_length,
         &key
         );
 
@@ -5181,14 +5199,14 @@
     *out_params = result;
     *out_params_len = result_size;
 
-    free(params);
+    free(custom_data);
 
     return 1;   // success
 
 fail:
     free(result);
 
-    free(params);
+    free(custom_data);
 
     return 0;       // This shouldn't happen!
 }
@@ -7712,9 +7730,9 @@
             ok = psa_generate_key_wrapper(in_params, in_params_len,
                                           &out_params, &out_params_len);
             break;
-        case PSA_GENERATE_KEY_EXT:
-            ok = psa_generate_key_ext_wrapper(in_params, in_params_len,
-                                              &out_params, &out_params_len);
+        case PSA_GENERATE_KEY_CUSTOM:
+            ok = psa_generate_key_custom_wrapper(in_params, in_params_len,
+                                                 &out_params, &out_params_len);
             break;
         case PSA_GENERATE_RANDOM:
             ok = psa_generate_random_wrapper(in_params, in_params_len,
@@ -7800,9 +7818,9 @@
             ok = psa_key_derivation_output_key_wrapper(in_params, in_params_len,
                                                        &out_params, &out_params_len);
             break;
-        case PSA_KEY_DERIVATION_OUTPUT_KEY_EXT:
-            ok = psa_key_derivation_output_key_ext_wrapper(in_params, in_params_len,
-                                                           &out_params, &out_params_len);
+        case PSA_KEY_DERIVATION_OUTPUT_KEY_CUSTOM:
+            ok = psa_key_derivation_output_key_custom_wrapper(in_params, in_params_len,
+                                                              &out_params, &out_params_len);
             break;
         case PSA_KEY_DERIVATION_SET_CAPACITY:
             ok = psa_key_derivation_set_capacity_wrapper(in_params, in_params_len,
diff --git a/tests/psa-client-server/psasim/src/psa_sim_serialise.c b/tests/psa-client-server/psasim/src/psa_sim_serialise.c
index 92ecdd2..e5c7225 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_serialise.c
+++ b/tests/psa-client-server/psasim/src/psa_sim_serialise.c
@@ -735,96 +735,38 @@
     return 1;
 }
 
-#define SER_TAG_SIZE        4
-
-size_t psasim_serialise_psa_key_production_parameters_t_needs(
-    const psa_key_production_parameters_t *params,
-    size_t data_length)
+size_t psasim_serialise_psa_custom_key_parameters_t_needs(
+    psa_custom_key_parameters_t value)
 {
-    /* We will serialise with 4-byte tag = "PKPP" + 4-byte overall length at the beginning,
-     * followed by size_t data_length, then the actual data from the structure.
-     */
-    return SER_TAG_SIZE + sizeof(uint32_t) + sizeof(data_length) + sizeof(*params) + data_length;
+    return sizeof(value);
 }
 
-int psasim_serialise_psa_key_production_parameters_t(uint8_t **pos,
-                                                     size_t *remaining,
-                                                     const psa_key_production_parameters_t *params,
-                                                     size_t data_length)
+int psasim_serialise_psa_custom_key_parameters_t(uint8_t **pos,
+                                                 size_t *remaining,
+                                                 psa_custom_key_parameters_t value)
 {
-    if (data_length > UINT32_MAX / 2) {       /* arbitrary limit */
-        return 0;       /* too big to serialise */
-    }
-
-    /* We use 32-bit lengths, which should be enough for any reasonable usage :) */
-    /* (the UINT32_MAX / 2 above is an even more conservative check to avoid overflow here) */
-    uint32_t len = (uint32_t) (sizeof(data_length) + sizeof(*params) + data_length);
-    if (*remaining < SER_TAG_SIZE + sizeof(uint32_t) + len) {
+    if (*remaining < sizeof(value)) {
         return 0;
     }
 
-    char tag[SER_TAG_SIZE] = "PKPP";
-
-    memcpy(*pos, tag, sizeof(tag));
-    memcpy(*pos + sizeof(tag), &len, sizeof(len));
-    *pos += sizeof(tag) + sizeof(len);
-    *remaining -= sizeof(tag) + sizeof(len);
-
-    memcpy(*pos, &data_length, sizeof(data_length));
-    memcpy(*pos + sizeof(data_length), params, sizeof(*params) + data_length);
-    *pos += sizeof(data_length) + sizeof(*params) + data_length;
-    *remaining -= sizeof(data_length) + sizeof(*params) + data_length;
+    memcpy(*pos, &value, sizeof(value));
+    *pos += sizeof(value);
 
     return 1;
 }
 
-int psasim_deserialise_psa_key_production_parameters_t(uint8_t **pos,
-                                                       size_t *remaining,
-                                                       psa_key_production_parameters_t **params,
-                                                       size_t *data_length)
+int psasim_deserialise_psa_custom_key_parameters_t(uint8_t **pos,
+                                                   size_t *remaining,
+                                                   psa_custom_key_parameters_t *value)
 {
-    if (*remaining < SER_TAG_SIZE + sizeof(uint32_t)) {
-        return 0;       /* can't even be an empty serialisation */
+    if (*remaining < sizeof(*value)) {
+        return 0;
     }
 
-    char tag[SER_TAG_SIZE] = "PKPP";    /* expected */
-    uint32_t len;
+    memcpy(value, *pos, sizeof(*value));
 
-    memcpy(&len, *pos + sizeof(tag), sizeof(len));
-
-    if (memcmp(*pos, tag, sizeof(tag)) != 0) {
-        return 0;       /* wrong tag */
-    }
-
-    *pos += sizeof(tag) + sizeof(len);
-    *remaining -= sizeof(tag) + sizeof(len);
-
-    if (*remaining < sizeof(*data_length)) {
-        return 0;       /* missing data_length */
-    }
-    memcpy(data_length, *pos, sizeof(*data_length));
-
-    if ((size_t) len != (sizeof(data_length) + sizeof(**params) + *data_length)) {
-        return 0;       /* wrong length */
-    }
-
-    if (*remaining < sizeof(*data_length) + sizeof(**params) + *data_length) {
-        return 0;       /* not enough data provided */
-    }
-
-    *pos += sizeof(data_length);
-    *remaining -= sizeof(data_length);
-
-    psa_key_production_parameters_t *out = malloc(sizeof(**params) + *data_length);
-    if (out == NULL) {
-        return 0;       /* allocation failure */
-    }
-
-    memcpy(out, *pos, sizeof(*out) + *data_length);
-    *pos += sizeof(*out) + *data_length;
-    *remaining -= sizeof(*out) + *data_length;
-
-    *params = out;
+    *pos += sizeof(*value);
+    *remaining -= sizeof(*value);
 
     return 1;
 }
diff --git a/tests/psa-client-server/psasim/src/psa_sim_serialise.h b/tests/psa-client-server/psasim/src/psa_sim_serialise.h
index f60e371..523ce80 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_serialise.h
+++ b/tests/psa-client-server/psasim/src/psa_sim_serialise.h
@@ -421,55 +421,48 @@
 int psasim_deserialise_return_buffer(uint8_t **pos, size_t *remaining,
                                      uint8_t *buffer, size_t buffer_length);
 
-/** Return how much space is needed by \c psasim_serialise_psa_key_production_parameters_t()
- *  to serialise a psa_key_production_parameters_t (a structure with a flexible array member).
+/** Return how much buffer space is needed by \c psasim_serialise_psa_custom_key_parameters_t()
+ *  to serialise a `psa_custom_key_parameters_t`.
  *
- * \param params             Pointer to the struct to be serialised
+ * \param value              The value that will be serialised into the buffer
  *                           (needed in case some serialisations are value-
  *                           dependent).
- * \param data_length        Number of bytes in the data[] of the struct to be serialised.
  *
- * \return                   The number of bytes needed in the serialisation buffer by
- *                           \c psasim_serialise_psa_key_production_parameters_t() to serialise
- *                           the specified structure.
+ * \return                   The number of bytes needed in the buffer by
+ *                           \c psasim_serialise_psa_custom_key_parameters_t() to serialise
+ *                           the given value.
  */
-size_t psasim_serialise_psa_key_production_parameters_t_needs(
-    const psa_key_production_parameters_t *params,
-    size_t buffer_size);
+size_t psasim_serialise_psa_custom_key_parameters_t_needs(
+    psa_custom_key_parameters_t value);
 
-/** Serialise a psa_key_production_parameters_t.
+/** Serialise a `psa_custom_key_parameters_t` into a buffer.
  *
  * \param pos[in,out]        Pointer to a `uint8_t *` holding current position
  *                           in the buffer.
  * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
  *                           remaining in the buffer.
- * \param params             Pointer to the structure to be serialised.
- * \param data_length        Number of bytes in the data[] of the struct to be serialised.
+ * \param value              The value to serialise into the buffer.
  *
  * \return                   \c 1 on success ("okay"), \c 0 on error.
  */
-int psasim_serialise_psa_key_production_parameters_t(uint8_t **pos,
-                                                     size_t *remaining,
-                                                     const psa_key_production_parameters_t *params,
-                                                     size_t data_length);
+int psasim_serialise_psa_custom_key_parameters_t(uint8_t **pos,
+                                                 size_t *remaining,
+                                                 psa_custom_key_parameters_t value);
 
-/** Deserialise a psa_key_production_parameters_t.
+/** Deserialise a `psa_custom_key_parameters_t` from a buffer.
  *
  * \param pos[in,out]        Pointer to a `uint8_t *` holding current position
- *                           in the serialisation buffer.
+ *                           in the buffer.
  * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
- *                           remaining in the serialisation buffer.
- * \param params             Pointer to a `psa_key_production_parameters_t *` to
- *                           receive the address of a newly-allocated structure,
- *                           which the caller must `free()`.
- * \param data_length        Pointer to a `size_t` to receive the number of
- *                           bytes in the data[] member of the structure deserialised.
+ *                           remaining in the buffer.
+ * \param value              Pointer to a `psa_custom_key_parameters_t` to receive the value
+ *                           deserialised from the buffer.
  *
  * \return                   \c 1 on success ("okay"), \c 0 on error.
  */
-int psasim_deserialise_psa_key_production_parameters_t(uint8_t **pos, size_t *remaining,
-                                                       psa_key_production_parameters_t **params,
-                                                       size_t *buffer_length);
+int psasim_deserialise_psa_custom_key_parameters_t(uint8_t **pos,
+                                                   size_t *remaining,
+                                                   psa_custom_key_parameters_t *value);
 
 /** Return how much buffer space is needed by \c psasim_serialise_psa_status_t()
  *  to serialise a `psa_status_t`.
diff --git a/tests/src/psa_test_wrappers.c b/tests/src/psa_test_wrappers.c
index eceb40b..deac300 100644
--- a/tests/src/psa_test_wrappers.c
+++ b/tests/src/psa_test_wrappers.c
@@ -622,17 +622,6 @@
     return status;
 }
 
-/* Wrapper for psa_generate_key_ext */
-psa_status_t mbedtls_test_wrap_psa_generate_key_ext(
-    const psa_key_attributes_t *arg0_attributes,
-    const psa_key_production_parameters_t *arg1_params,
-    size_t arg2_params_data_length,
-    mbedtls_svc_key_id_t *arg3_key)
-{
-    psa_status_t status = (psa_generate_key_ext)(arg0_attributes, arg1_params, arg2_params_data_length, arg3_key);
-    return status;
-}
-
 /* Wrapper for psa_generate_random */
 psa_status_t mbedtls_test_wrap_psa_generate_random(
     uint8_t *arg0_output,
@@ -907,18 +896,6 @@
     return status;
 }
 
-/* Wrapper for psa_key_derivation_output_key_ext */
-psa_status_t mbedtls_test_wrap_psa_key_derivation_output_key_ext(
-    const psa_key_attributes_t *arg0_attributes,
-    psa_key_derivation_operation_t *arg1_operation,
-    const psa_key_production_parameters_t *arg2_params,
-    size_t arg3_params_data_length,
-    mbedtls_svc_key_id_t *arg4_key)
-{
-    psa_status_t status = (psa_key_derivation_output_key_ext)(arg0_attributes, arg1_operation, arg2_params, arg3_params_data_length, arg4_key);
-    return status;
-}
-
 /* Wrapper for psa_key_derivation_set_capacity */
 psa_status_t mbedtls_test_wrap_psa_key_derivation_set_capacity(
     psa_key_derivation_operation_t *arg0_operation,