Merge pull request #9602 from paul-elliott-arm/add_key_agreement_negative_tests

Add key agreement negative tests
diff --git a/tests/include/test/psa_exercise_key.h b/tests/include/test/psa_exercise_key.h
index f6be307..f44608b 100644
--- a/tests/include/test/psa_exercise_key.h
+++ b/tests/include/test/psa_exercise_key.h
@@ -138,7 +138,7 @@
     size_t capacity, int key_destroyable);
 
 /** Perform a key agreement using the given key pair against its public key
- * using psa_raw_key_agreement().
+ * using psa_raw_key_agreement() and psa_key_agreement().
  *
  * The result is discarded. The purpose of this function is to smoke-test a key.
  *
diff --git a/tests/include/test/psa_test_wrappers.h b/tests/include/test/psa_test_wrappers.h
index 4d674e7..d7e7b43 100644
--- a/tests/include/test/psa_test_wrappers.h
+++ b/tests/include/test/psa_test_wrappers.h
@@ -370,7 +370,7 @@
 
 psa_status_t mbedtls_test_wrap_psa_generate_key_iop_complete(
     psa_generate_key_iop_t *arg0_operation,
-    psa_key_id_t *arg1_key);
+    mbedtls_svc_key_id_t *arg1_key);
 #define psa_generate_key_iop_complete(arg0_operation, arg1_key) \
     mbedtls_test_wrap_psa_generate_key_iop_complete(arg0_operation, arg1_key)
 
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 1301ff2..4be53c5 100644
--- a/tests/psa-client-server/psasim/src/psa_functions_codes.h
+++ b/tests/psa-client-server/psasim/src/psa_functions_codes.h
@@ -41,6 +41,10 @@
     PSA_EXPORT_PUBLIC_KEY,
     PSA_GENERATE_KEY,
     PSA_GENERATE_KEY_CUSTOM,
+    PSA_GENERATE_KEY_IOP_ABORT,
+    PSA_GENERATE_KEY_IOP_COMPLETE,
+    PSA_GENERATE_KEY_IOP_GET_NUM_OPS,
+    PSA_GENERATE_KEY_IOP_SETUP,
     PSA_GENERATE_RANDOM,
     PSA_GET_KEY_ATTRIBUTES,
     PSA_HASH_ABORT,
@@ -54,6 +58,11 @@
     PSA_IMPORT_KEY,
     PSA_INTERRUPTIBLE_GET_MAX_OPS,
     PSA_INTERRUPTIBLE_SET_MAX_OPS,
+    PSA_KEY_AGREEMENT,
+    PSA_KEY_AGREEMENT_IOP_ABORT,
+    PSA_KEY_AGREEMENT_IOP_COMPLETE,
+    PSA_KEY_AGREEMENT_IOP_GET_NUM_OPS,
+    PSA_KEY_AGREEMENT_IOP_SETUP,
     PSA_KEY_DERIVATION_ABORT,
     PSA_KEY_DERIVATION_GET_CAPACITY,
     PSA_KEY_DERIVATION_INPUT_BYTES,
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 d4d9d60..f6efd62 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
@@ -2897,6 +2897,309 @@
 }
 
 
+psa_status_t psa_generate_key_iop_abort(
+    psa_generate_key_iop_t *operation
+    )
+{
+    uint8_t *ser_params = NULL;
+    uint8_t *ser_result = NULL;
+    size_t result_length;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    size_t needed =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_generate_key_iop_t_needs(*operation);
+
+    ser_params = malloc(needed);
+    if (ser_params == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto fail;
+    }
+
+    uint8_t *pos = ser_params;
+    size_t remaining = needed;
+    int ok;
+    ok = psasim_serialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_generate_key_iop_t(
+        &pos, &remaining,
+        *operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psa_crypto_call(PSA_GENERATE_KEY_IOP_ABORT,
+                         ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
+    if (!ok) {
+        printf("PSA_GENERATE_KEY_IOP_ABORT server call failed\n");
+        goto fail;
+    }
+
+    uint8_t *rpos = ser_result;
+    size_t rremain = result_length;
+
+    ok = psasim_deserialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_status_t(
+        &rpos, &rremain,
+        &status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_generate_key_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+fail:
+    free(ser_params);
+    free(ser_result);
+
+    return status;
+}
+
+
+psa_status_t psa_generate_key_iop_complete(
+    psa_generate_key_iop_t *operation,
+    mbedtls_svc_key_id_t *key
+    )
+{
+    uint8_t *ser_params = NULL;
+    uint8_t *ser_result = NULL;
+    size_t result_length;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    size_t needed =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_generate_key_iop_t_needs(*operation) +
+        psasim_serialise_mbedtls_svc_key_id_t_needs(*key);
+
+    ser_params = malloc(needed);
+    if (ser_params == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto fail;
+    }
+
+    uint8_t *pos = ser_params;
+    size_t remaining = needed;
+    int ok;
+    ok = psasim_serialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_generate_key_iop_t(
+        &pos, &remaining,
+        *operation);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        *key);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psa_crypto_call(PSA_GENERATE_KEY_IOP_COMPLETE,
+                         ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
+    if (!ok) {
+        printf("PSA_GENERATE_KEY_IOP_COMPLETE server call failed\n");
+        goto fail;
+    }
+
+    uint8_t *rpos = ser_result;
+    size_t rremain = result_length;
+
+    ok = psasim_deserialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_status_t(
+        &rpos, &rremain,
+        &status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_generate_key_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_mbedtls_svc_key_id_t(
+        &rpos, &rremain,
+        key);
+    if (!ok) {
+        goto fail;
+    }
+
+fail:
+    free(ser_params);
+    free(ser_result);
+
+    return status;
+}
+
+
+uint32_t psa_generate_key_iop_get_num_ops(
+    psa_generate_key_iop_t *operation
+    )
+{
+    uint8_t *ser_params = NULL;
+    uint8_t *ser_result = NULL;
+    size_t result_length;
+    uint32_t value = 0;
+
+    size_t needed =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_generate_key_iop_t_needs(*operation);
+
+    ser_params = malloc(needed);
+    if (ser_params == NULL) {
+        value = 0;
+        goto fail;
+    }
+
+    uint8_t *pos = ser_params;
+    size_t remaining = needed;
+    int ok;
+    ok = psasim_serialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_generate_key_iop_t(
+        &pos, &remaining,
+        *operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psa_crypto_call(PSA_GENERATE_KEY_IOP_GET_NUM_OPS,
+                         ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
+    if (!ok) {
+        printf("PSA_GENERATE_KEY_IOP_GET_NUM_OPS server call failed\n");
+        goto fail;
+    }
+
+    uint8_t *rpos = ser_result;
+    size_t rremain = result_length;
+
+    ok = psasim_deserialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_uint32_t(
+        &rpos, &rremain,
+        &value);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_generate_key_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+fail:
+    free(ser_params);
+    free(ser_result);
+
+    return value;
+}
+
+
+psa_status_t psa_generate_key_iop_setup(
+    psa_generate_key_iop_t *operation,
+    const psa_key_attributes_t *attributes
+    )
+{
+    uint8_t *ser_params = NULL;
+    uint8_t *ser_result = NULL;
+    size_t result_length;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    size_t needed =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_generate_key_iop_t_needs(*operation) +
+        psasim_serialise_psa_key_attributes_t_needs(*attributes);
+
+    ser_params = malloc(needed);
+    if (ser_params == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto fail;
+    }
+
+    uint8_t *pos = ser_params;
+    size_t remaining = needed;
+    int ok;
+    ok = psasim_serialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_generate_key_iop_t(
+        &pos, &remaining,
+        *operation);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_key_attributes_t(
+        &pos, &remaining,
+        *attributes);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psa_crypto_call(PSA_GENERATE_KEY_IOP_SETUP,
+                         ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
+    if (!ok) {
+        printf("PSA_GENERATE_KEY_IOP_SETUP server call failed\n");
+        goto fail;
+    }
+
+    uint8_t *rpos = ser_result;
+    size_t rremain = result_length;
+
+    ok = psasim_deserialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_status_t(
+        &rpos, &rremain,
+        &status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_generate_key_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+fail:
+    free(ser_params);
+    free(ser_result);
+
+    return status;
+}
+
+
 psa_status_t psa_generate_random(
     uint8_t *output, size_t  output_size
     )
@@ -3902,6 +4205,435 @@
 }
 
 
+psa_status_t psa_key_agreement(
+    mbedtls_svc_key_id_t private_key,
+    const uint8_t *peer_key, size_t  peer_key_length,
+    psa_algorithm_t alg,
+    const psa_key_attributes_t *attributes,
+    mbedtls_svc_key_id_t *key
+    )
+{
+    uint8_t *ser_params = NULL;
+    uint8_t *ser_result = NULL;
+    size_t result_length;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    size_t needed =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_mbedtls_svc_key_id_t_needs(private_key) +
+        psasim_serialise_buffer_needs(peer_key, peer_key_length) +
+        psasim_serialise_psa_algorithm_t_needs(alg) +
+        psasim_serialise_psa_key_attributes_t_needs(*attributes) +
+        psasim_serialise_mbedtls_svc_key_id_t_needs(*key);
+
+    ser_params = malloc(needed);
+    if (ser_params == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto fail;
+    }
+
+    uint8_t *pos = ser_params;
+    size_t remaining = needed;
+    int ok;
+    ok = psasim_serialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        private_key);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_buffer(
+        &pos, &remaining,
+        peer_key, peer_key_length);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_algorithm_t(
+        &pos, &remaining,
+        alg);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_key_attributes_t(
+        &pos, &remaining,
+        *attributes);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        *key);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psa_crypto_call(PSA_KEY_AGREEMENT,
+                         ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
+    if (!ok) {
+        printf("PSA_KEY_AGREEMENT server call failed\n");
+        goto fail;
+    }
+
+    uint8_t *rpos = ser_result;
+    size_t rremain = result_length;
+
+    ok = psasim_deserialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_status_t(
+        &rpos, &rremain,
+        &status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_mbedtls_svc_key_id_t(
+        &rpos, &rremain,
+        key);
+    if (!ok) {
+        goto fail;
+    }
+
+fail:
+    free(ser_params);
+    free(ser_result);
+
+    return status;
+}
+
+
+psa_status_t psa_key_agreement_iop_abort(
+    psa_key_agreement_iop_t *operation
+    )
+{
+    uint8_t *ser_params = NULL;
+    uint8_t *ser_result = NULL;
+    size_t result_length;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    size_t needed =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_key_agreement_iop_t_needs(*operation);
+
+    ser_params = malloc(needed);
+    if (ser_params == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto fail;
+    }
+
+    uint8_t *pos = ser_params;
+    size_t remaining = needed;
+    int ok;
+    ok = psasim_serialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_key_agreement_iop_t(
+        &pos, &remaining,
+        *operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psa_crypto_call(PSA_KEY_AGREEMENT_IOP_ABORT,
+                         ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
+    if (!ok) {
+        printf("PSA_KEY_AGREEMENT_IOP_ABORT server call failed\n");
+        goto fail;
+    }
+
+    uint8_t *rpos = ser_result;
+    size_t rremain = result_length;
+
+    ok = psasim_deserialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_status_t(
+        &rpos, &rremain,
+        &status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_agreement_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+fail:
+    free(ser_params);
+    free(ser_result);
+
+    return status;
+}
+
+
+psa_status_t psa_key_agreement_iop_complete(
+    psa_key_agreement_iop_t *operation,
+    mbedtls_svc_key_id_t *key
+    )
+{
+    uint8_t *ser_params = NULL;
+    uint8_t *ser_result = NULL;
+    size_t result_length;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    size_t needed =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_key_agreement_iop_t_needs(*operation) +
+        psasim_serialise_mbedtls_svc_key_id_t_needs(*key);
+
+    ser_params = malloc(needed);
+    if (ser_params == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto fail;
+    }
+
+    uint8_t *pos = ser_params;
+    size_t remaining = needed;
+    int ok;
+    ok = psasim_serialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_key_agreement_iop_t(
+        &pos, &remaining,
+        *operation);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        *key);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psa_crypto_call(PSA_KEY_AGREEMENT_IOP_COMPLETE,
+                         ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
+    if (!ok) {
+        printf("PSA_KEY_AGREEMENT_IOP_COMPLETE server call failed\n");
+        goto fail;
+    }
+
+    uint8_t *rpos = ser_result;
+    size_t rremain = result_length;
+
+    ok = psasim_deserialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_status_t(
+        &rpos, &rremain,
+        &status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_agreement_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_mbedtls_svc_key_id_t(
+        &rpos, &rremain,
+        key);
+    if (!ok) {
+        goto fail;
+    }
+
+fail:
+    free(ser_params);
+    free(ser_result);
+
+    return status;
+}
+
+
+uint32_t psa_key_agreement_iop_get_num_ops(
+    psa_key_agreement_iop_t *operation
+    )
+{
+    uint8_t *ser_params = NULL;
+    uint8_t *ser_result = NULL;
+    size_t result_length;
+    uint32_t value = 0;
+
+    size_t needed =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_key_agreement_iop_t_needs(*operation);
+
+    ser_params = malloc(needed);
+    if (ser_params == NULL) {
+        value = 0;
+        goto fail;
+    }
+
+    uint8_t *pos = ser_params;
+    size_t remaining = needed;
+    int ok;
+    ok = psasim_serialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_key_agreement_iop_t(
+        &pos, &remaining,
+        *operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psa_crypto_call(PSA_KEY_AGREEMENT_IOP_GET_NUM_OPS,
+                         ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
+    if (!ok) {
+        printf("PSA_KEY_AGREEMENT_IOP_GET_NUM_OPS server call failed\n");
+        goto fail;
+    }
+
+    uint8_t *rpos = ser_result;
+    size_t rremain = result_length;
+
+    ok = psasim_deserialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_uint32_t(
+        &rpos, &rremain,
+        &value);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_agreement_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+fail:
+    free(ser_params);
+    free(ser_result);
+
+    return value;
+}
+
+
+psa_status_t psa_key_agreement_iop_setup(
+    psa_key_agreement_iop_t *operation,
+    mbedtls_svc_key_id_t private_key,
+    const uint8_t *peer_key, size_t  peer_key_length,
+    psa_algorithm_t alg,
+    const psa_key_attributes_t *attributes
+    )
+{
+    uint8_t *ser_params = NULL;
+    uint8_t *ser_result = NULL;
+    size_t result_length;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    size_t needed =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_key_agreement_iop_t_needs(*operation) +
+        psasim_serialise_mbedtls_svc_key_id_t_needs(private_key) +
+        psasim_serialise_buffer_needs(peer_key, peer_key_length) +
+        psasim_serialise_psa_algorithm_t_needs(alg) +
+        psasim_serialise_psa_key_attributes_t_needs(*attributes);
+
+    ser_params = malloc(needed);
+    if (ser_params == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto fail;
+    }
+
+    uint8_t *pos = ser_params;
+    size_t remaining = needed;
+    int ok;
+    ok = psasim_serialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_key_agreement_iop_t(
+        &pos, &remaining,
+        *operation);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        private_key);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_buffer(
+        &pos, &remaining,
+        peer_key, peer_key_length);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_algorithm_t(
+        &pos, &remaining,
+        alg);
+    if (!ok) {
+        goto fail;
+    }
+    ok = psasim_serialise_psa_key_attributes_t(
+        &pos, &remaining,
+        *attributes);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psa_crypto_call(PSA_KEY_AGREEMENT_IOP_SETUP,
+                         ser_params, (size_t) (pos - ser_params), &ser_result, &result_length);
+    if (!ok) {
+        printf("PSA_KEY_AGREEMENT_IOP_SETUP server call failed\n");
+        goto fail;
+    }
+
+    uint8_t *rpos = ser_result;
+    size_t rremain = result_length;
+
+    ok = psasim_deserialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_status_t(
+        &rpos, &rremain,
+        &status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_agreement_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+fail:
+    free(ser_params);
+    free(ser_result);
+
+    return status;
+}
+
+
 psa_status_t psa_key_derivation_abort(
     psa_key_derivation_operation_t *operation
     )
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 a88fc51..599e55f 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
@@ -3227,6 +3227,332 @@
 }
 
 // Returns 1 for success, 0 for failure
+int psa_generate_key_iop_abort_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_generate_key_iop_t operation;
+
+    uint8_t *pos = in_params;
+    size_t remaining = in_params_len;
+    uint8_t *result = NULL;
+    int ok;
+
+    ok = psasim_deserialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_generate_key_iop_t(
+        &pos, &remaining,
+        &operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    // Now we call the actual target function
+
+    status = psa_generate_key_iop_abort(
+        &operation
+        );
+
+    // NOTE: Should really check there is no overflow as we go along.
+    size_t result_size =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_status_t_needs(status) +
+        psasim_serialise_psa_generate_key_iop_t_needs(operation);
+
+    result = malloc(result_size);
+    if (result == NULL) {
+        goto fail;
+    }
+
+    uint8_t *rpos = result;
+    size_t rremain = result_size;
+
+    ok = psasim_serialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_status_t(
+        &rpos, &rremain,
+        status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_generate_key_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    *out_params = result;
+    *out_params_len = result_size;
+
+    return 1;   // success
+
+fail:
+    free(result);
+
+    return 0;       // This shouldn't happen!
+}
+
+// Returns 1 for success, 0 for failure
+int psa_generate_key_iop_complete_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_generate_key_iop_t operation;
+    mbedtls_svc_key_id_t key;
+
+    uint8_t *pos = in_params;
+    size_t remaining = in_params_len;
+    uint8_t *result = NULL;
+    int ok;
+
+    ok = psasim_deserialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_generate_key_iop_t(
+        &pos, &remaining,
+        &operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        &key);
+    if (!ok) {
+        goto fail;
+    }
+
+    // Now we call the actual target function
+
+    status = psa_generate_key_iop_complete(
+        &operation,
+        &key
+        );
+
+    // NOTE: Should really check there is no overflow as we go along.
+    size_t result_size =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_status_t_needs(status) +
+        psasim_serialise_psa_generate_key_iop_t_needs(operation) +
+        psasim_serialise_mbedtls_svc_key_id_t_needs(key);
+
+    result = malloc(result_size);
+    if (result == NULL) {
+        goto fail;
+    }
+
+    uint8_t *rpos = result;
+    size_t rremain = result_size;
+
+    ok = psasim_serialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_status_t(
+        &rpos, &rremain,
+        status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_generate_key_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_mbedtls_svc_key_id_t(
+        &rpos, &rremain,
+        key);
+    if (!ok) {
+        goto fail;
+    }
+
+    *out_params = result;
+    *out_params_len = result_size;
+
+    return 1;   // success
+
+fail:
+    free(result);
+
+    return 0;       // This shouldn't happen!
+}
+
+// Returns 1 for success, 0 for failure
+int psa_generate_key_iop_get_num_ops_wrapper(
+    uint8_t *in_params, size_t in_params_len,
+    uint8_t **out_params, size_t *out_params_len)
+{
+    uint32_t value = 0;
+    psa_generate_key_iop_t operation;
+
+    uint8_t *pos = in_params;
+    size_t remaining = in_params_len;
+    uint8_t *result = NULL;
+    int ok;
+
+    ok = psasim_deserialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_generate_key_iop_t(
+        &pos, &remaining,
+        &operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    // Now we call the actual target function
+
+    value = psa_generate_key_iop_get_num_ops(
+        &operation
+        );
+
+    // NOTE: Should really check there is no overflow as we go along.
+    size_t result_size =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_uint32_t_needs(value) +
+        psasim_serialise_psa_generate_key_iop_t_needs(operation);
+
+    result = malloc(result_size);
+    if (result == NULL) {
+        goto fail;
+    }
+
+    uint8_t *rpos = result;
+    size_t rremain = result_size;
+
+    ok = psasim_serialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_uint32_t(
+        &rpos, &rremain,
+        value);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_generate_key_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    *out_params = result;
+    *out_params_len = result_size;
+
+    return 1;   // success
+
+fail:
+    free(result);
+
+    return 0;       // This shouldn't happen!
+}
+
+// Returns 1 for success, 0 for failure
+int psa_generate_key_iop_setup_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_generate_key_iop_t operation;
+    psa_key_attributes_t attributes;
+
+    uint8_t *pos = in_params;
+    size_t remaining = in_params_len;
+    uint8_t *result = NULL;
+    int ok;
+
+    ok = psasim_deserialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_generate_key_iop_t(
+        &pos, &remaining,
+        &operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_attributes_t(
+        &pos, &remaining,
+        &attributes);
+    if (!ok) {
+        goto fail;
+    }
+
+    // Now we call the actual target function
+
+    status = psa_generate_key_iop_setup(
+        &operation,
+        &attributes
+        );
+
+    // NOTE: Should really check there is no overflow as we go along.
+    size_t result_size =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_status_t_needs(status) +
+        psasim_serialise_psa_generate_key_iop_t_needs(operation);
+
+    result = malloc(result_size);
+    if (result == NULL) {
+        goto fail;
+    }
+
+    uint8_t *rpos = result;
+    size_t rremain = result_size;
+
+    ok = psasim_serialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_status_t(
+        &rpos, &rremain,
+        status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_generate_key_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    *out_params = result;
+    *out_params_len = result_size;
+
+    return 1;   // success
+
+fail:
+    free(result);
+
+    return 0;       // This shouldn't happen!
+}
+
+// Returns 1 for success, 0 for failure
 int psa_generate_random_wrapper(
     uint8_t *in_params, size_t in_params_len,
     uint8_t **out_params, size_t *out_params_len)
@@ -4344,6 +4670,480 @@
 }
 
 // Returns 1 for success, 0 for failure
+int psa_key_agreement_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;
+    mbedtls_svc_key_id_t private_key;
+    uint8_t *peer_key = NULL;
+    size_t peer_key_length;
+    psa_algorithm_t alg;
+    psa_key_attributes_t attributes;
+    mbedtls_svc_key_id_t key;
+
+    uint8_t *pos = in_params;
+    size_t remaining = in_params_len;
+    uint8_t *result = NULL;
+    int ok;
+
+    ok = psasim_deserialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        &private_key);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_buffer(
+        &pos, &remaining,
+        &peer_key, &peer_key_length);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_algorithm_t(
+        &pos, &remaining,
+        &alg);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_attributes_t(
+        &pos, &remaining,
+        &attributes);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        &key);
+    if (!ok) {
+        goto fail;
+    }
+
+    // Now we call the actual target function
+
+    status = psa_key_agreement(
+        private_key,
+        peer_key, peer_key_length,
+        alg,
+        &attributes,
+        &key
+        );
+
+    // NOTE: Should really check there is no overflow as we go along.
+    size_t result_size =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_status_t_needs(status) +
+        psasim_serialise_mbedtls_svc_key_id_t_needs(key);
+
+    result = malloc(result_size);
+    if (result == NULL) {
+        goto fail;
+    }
+
+    uint8_t *rpos = result;
+    size_t rremain = result_size;
+
+    ok = psasim_serialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_status_t(
+        &rpos, &rremain,
+        status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_mbedtls_svc_key_id_t(
+        &rpos, &rremain,
+        key);
+    if (!ok) {
+        goto fail;
+    }
+
+    *out_params = result;
+    *out_params_len = result_size;
+
+    free(peer_key);
+
+    return 1;   // success
+
+fail:
+    free(result);
+
+    free(peer_key);
+
+    return 0;       // This shouldn't happen!
+}
+
+// Returns 1 for success, 0 for failure
+int psa_key_agreement_iop_abort_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_agreement_iop_t operation;
+
+    uint8_t *pos = in_params;
+    size_t remaining = in_params_len;
+    uint8_t *result = NULL;
+    int ok;
+
+    ok = psasim_deserialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_agreement_iop_t(
+        &pos, &remaining,
+        &operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    // Now we call the actual target function
+
+    status = psa_key_agreement_iop_abort(
+        &operation
+        );
+
+    // NOTE: Should really check there is no overflow as we go along.
+    size_t result_size =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_status_t_needs(status) +
+        psasim_serialise_psa_key_agreement_iop_t_needs(operation);
+
+    result = malloc(result_size);
+    if (result == NULL) {
+        goto fail;
+    }
+
+    uint8_t *rpos = result;
+    size_t rremain = result_size;
+
+    ok = psasim_serialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_status_t(
+        &rpos, &rremain,
+        status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_key_agreement_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    *out_params = result;
+    *out_params_len = result_size;
+
+    return 1;   // success
+
+fail:
+    free(result);
+
+    return 0;       // This shouldn't happen!
+}
+
+// Returns 1 for success, 0 for failure
+int psa_key_agreement_iop_complete_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_agreement_iop_t operation;
+    mbedtls_svc_key_id_t key;
+
+    uint8_t *pos = in_params;
+    size_t remaining = in_params_len;
+    uint8_t *result = NULL;
+    int ok;
+
+    ok = psasim_deserialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_agreement_iop_t(
+        &pos, &remaining,
+        &operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        &key);
+    if (!ok) {
+        goto fail;
+    }
+
+    // Now we call the actual target function
+
+    status = psa_key_agreement_iop_complete(
+        &operation,
+        &key
+        );
+
+    // NOTE: Should really check there is no overflow as we go along.
+    size_t result_size =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_status_t_needs(status) +
+        psasim_serialise_psa_key_agreement_iop_t_needs(operation) +
+        psasim_serialise_mbedtls_svc_key_id_t_needs(key);
+
+    result = malloc(result_size);
+    if (result == NULL) {
+        goto fail;
+    }
+
+    uint8_t *rpos = result;
+    size_t rremain = result_size;
+
+    ok = psasim_serialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_status_t(
+        &rpos, &rremain,
+        status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_key_agreement_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_mbedtls_svc_key_id_t(
+        &rpos, &rremain,
+        key);
+    if (!ok) {
+        goto fail;
+    }
+
+    *out_params = result;
+    *out_params_len = result_size;
+
+    return 1;   // success
+
+fail:
+    free(result);
+
+    return 0;       // This shouldn't happen!
+}
+
+// Returns 1 for success, 0 for failure
+int psa_key_agreement_iop_get_num_ops_wrapper(
+    uint8_t *in_params, size_t in_params_len,
+    uint8_t **out_params, size_t *out_params_len)
+{
+    uint32_t value = 0;
+    psa_key_agreement_iop_t operation;
+
+    uint8_t *pos = in_params;
+    size_t remaining = in_params_len;
+    uint8_t *result = NULL;
+    int ok;
+
+    ok = psasim_deserialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_agreement_iop_t(
+        &pos, &remaining,
+        &operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    // Now we call the actual target function
+
+    value = psa_key_agreement_iop_get_num_ops(
+        &operation
+        );
+
+    // NOTE: Should really check there is no overflow as we go along.
+    size_t result_size =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_uint32_t_needs(value) +
+        psasim_serialise_psa_key_agreement_iop_t_needs(operation);
+
+    result = malloc(result_size);
+    if (result == NULL) {
+        goto fail;
+    }
+
+    uint8_t *rpos = result;
+    size_t rremain = result_size;
+
+    ok = psasim_serialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_uint32_t(
+        &rpos, &rremain,
+        value);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_key_agreement_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    *out_params = result;
+    *out_params_len = result_size;
+
+    return 1;   // success
+
+fail:
+    free(result);
+
+    return 0;       // This shouldn't happen!
+}
+
+// Returns 1 for success, 0 for failure
+int psa_key_agreement_iop_setup_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_agreement_iop_t operation;
+    mbedtls_svc_key_id_t private_key;
+    uint8_t *peer_key = NULL;
+    size_t peer_key_length;
+    psa_algorithm_t alg;
+    psa_key_attributes_t attributes;
+
+    uint8_t *pos = in_params;
+    size_t remaining = in_params_len;
+    uint8_t *result = NULL;
+    int ok;
+
+    ok = psasim_deserialise_begin(&pos, &remaining);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_agreement_iop_t(
+        &pos, &remaining,
+        &operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_mbedtls_svc_key_id_t(
+        &pos, &remaining,
+        &private_key);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_buffer(
+        &pos, &remaining,
+        &peer_key, &peer_key_length);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_algorithm_t(
+        &pos, &remaining,
+        &alg);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_deserialise_psa_key_attributes_t(
+        &pos, &remaining,
+        &attributes);
+    if (!ok) {
+        goto fail;
+    }
+
+    // Now we call the actual target function
+
+    status = psa_key_agreement_iop_setup(
+        &operation,
+        private_key,
+        peer_key, peer_key_length,
+        alg,
+        &attributes
+        );
+
+    // NOTE: Should really check there is no overflow as we go along.
+    size_t result_size =
+        psasim_serialise_begin_needs() +
+        psasim_serialise_psa_status_t_needs(status) +
+        psasim_serialise_psa_key_agreement_iop_t_needs(operation);
+
+    result = malloc(result_size);
+    if (result == NULL) {
+        goto fail;
+    }
+
+    uint8_t *rpos = result;
+    size_t rremain = result_size;
+
+    ok = psasim_serialise_begin(&rpos, &rremain);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_status_t(
+        &rpos, &rremain,
+        status);
+    if (!ok) {
+        goto fail;
+    }
+
+    ok = psasim_serialise_psa_key_agreement_iop_t(
+        &rpos, &rremain,
+        operation);
+    if (!ok) {
+        goto fail;
+    }
+
+    *out_params = result;
+    *out_params_len = result_size;
+
+    free(peer_key);
+
+    return 1;   // success
+
+fail:
+    free(result);
+
+    free(peer_key);
+
+    return 0;       // This shouldn't happen!
+}
+
+// Returns 1 for success, 0 for failure
 int psa_key_derivation_abort_wrapper(
     uint8_t *in_params, size_t in_params_len,
     uint8_t **out_params, size_t *out_params_len)
@@ -7738,6 +8538,22 @@
             ok = psa_generate_key_custom_wrapper(in_params, in_params_len,
                                                  &out_params, &out_params_len);
             break;
+        case PSA_GENERATE_KEY_IOP_ABORT:
+            ok = psa_generate_key_iop_abort_wrapper(in_params, in_params_len,
+                                                    &out_params, &out_params_len);
+            break;
+        case PSA_GENERATE_KEY_IOP_COMPLETE:
+            ok = psa_generate_key_iop_complete_wrapper(in_params, in_params_len,
+                                                       &out_params, &out_params_len);
+            break;
+        case PSA_GENERATE_KEY_IOP_GET_NUM_OPS:
+            ok = psa_generate_key_iop_get_num_ops_wrapper(in_params, in_params_len,
+                                                          &out_params, &out_params_len);
+            break;
+        case PSA_GENERATE_KEY_IOP_SETUP:
+            ok = psa_generate_key_iop_setup_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,
                                              &out_params, &out_params_len);
@@ -7790,6 +8606,26 @@
             ok = psa_interruptible_set_max_ops_wrapper(in_params, in_params_len,
                                                        &out_params, &out_params_len);
             break;
+        case PSA_KEY_AGREEMENT:
+            ok = psa_key_agreement_wrapper(in_params, in_params_len,
+                                           &out_params, &out_params_len);
+            break;
+        case PSA_KEY_AGREEMENT_IOP_ABORT:
+            ok = psa_key_agreement_iop_abort_wrapper(in_params, in_params_len,
+                                                     &out_params, &out_params_len);
+            break;
+        case PSA_KEY_AGREEMENT_IOP_COMPLETE:
+            ok = psa_key_agreement_iop_complete_wrapper(in_params, in_params_len,
+                                                        &out_params, &out_params_len);
+            break;
+        case PSA_KEY_AGREEMENT_IOP_GET_NUM_OPS:
+            ok = psa_key_agreement_iop_get_num_ops_wrapper(in_params, in_params_len,
+                                                           &out_params, &out_params_len);
+            break;
+        case PSA_KEY_AGREEMENT_IOP_SETUP:
+            ok = psa_key_agreement_iop_setup_wrapper(in_params, in_params_len,
+                                                     &out_params, &out_params_len);
+            break;
         case PSA_KEY_DERIVATION_ABORT:
             ok = psa_key_derivation_abort_wrapper(in_params, in_params_len,
                                                   &out_params, &out_params_len);
diff --git a/tests/psa-client-server/psasim/src/psa_sim_generate.pl b/tests/psa-client-server/psasim/src/psa_sim_generate.pl
index e0e9b19..fbceddf 100755
--- a/tests/psa-client-server/psasim/src/psa_sim_generate.pl
+++ b/tests/psa-client-server/psasim/src/psa_sim_generate.pl
@@ -302,6 +302,10 @@
 #if !defined(MBEDTLS_PSA_CRYPTO_C)
 #error "Error: MBEDTLS_PSA_CRYPTO_C must be enabled on server build"
 #endif
+
+#if defined(MBEDTLS_TEST_HOOKS)
+void (*mbedtls_test_hook_error_add)(int, int, const char *, int);
+#endif
 EOF
 }
 
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 44d87d6..cd081e4 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_serialise.c
+++ b/tests/psa-client-server/psasim/src/psa_sim_serialise.c
@@ -1660,6 +1660,42 @@
     return 1;
 }
 
+size_t psasim_serialise_psa_generate_key_iop_t_needs(
+    psa_generate_key_iop_t value)
+{
+    return sizeof(value);
+}
+
+int psasim_serialise_psa_generate_key_iop_t(uint8_t **pos,
+                                            size_t *remaining,
+                                            psa_generate_key_iop_t value)
+{
+    if (*remaining < sizeof(value)) {
+        return 0;
+    }
+
+    memcpy(*pos, &value, sizeof(value));
+    *pos += sizeof(value);
+
+    return 1;
+}
+
+int psasim_deserialise_psa_generate_key_iop_t(uint8_t **pos,
+                                              size_t *remaining,
+                                              psa_generate_key_iop_t *value)
+{
+    if (*remaining < sizeof(*value)) {
+        return 0;
+    }
+
+    memcpy(value, *pos, sizeof(*value));
+
+    *pos += sizeof(*value);
+    *remaining -= sizeof(*value);
+
+    return 1;
+}
+
 void psa_sim_serialize_reset(void)
 {
     memset(hash_operation_handles, 0,
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 02f6bcb..a224d82 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_serialise.h
+++ b/tests/psa-client-server/psasim/src/psa_sim_serialise.h
@@ -1344,3 +1344,46 @@
 int psasim_deserialise_psa_key_agreement_iop_t(uint8_t **pos,
                                                size_t *remaining,
                                                psa_key_agreement_iop_t *value);
+
+/** Return how much buffer space is needed by \c psasim_serialise_psa_generate_key_iop_t()
+ *  to serialise a `psa_generate_key_iop_t`.
+ *
+ * \param value              The value that will be serialised into the buffer
+ *                           (needed in case some serialisations are value-
+ *                           dependent).
+ *
+ * \return                   The number of bytes needed in the buffer by
+ *                           \c psasim_serialise_psa_generate_key_iop_t() to serialise
+ *                           the given value.
+ */
+size_t psasim_serialise_psa_generate_key_iop_t_needs(
+    psa_generate_key_iop_t value);
+
+/** Serialise a `psa_generate_key_iop_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 value              The value to serialise into the buffer.
+ *
+ * \return                   \c 1 on success ("okay"), \c 0 on error.
+ */
+int psasim_serialise_psa_generate_key_iop_t(uint8_t **pos,
+                                            size_t *remaining,
+                                            psa_generate_key_iop_t value);
+
+/** Deserialise a `psa_generate_key_iop_t` from 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 value              Pointer to a `psa_generate_key_iop_t` to receive the value
+ *                           deserialised from the buffer.
+ *
+ * \return                   \c 1 on success ("okay"), \c 0 on error.
+ */
+int psasim_deserialise_psa_generate_key_iop_t(uint8_t **pos,
+                                              size_t *remaining,
+                                              psa_generate_key_iop_t *value);
diff --git a/tests/psa-client-server/psasim/src/psa_sim_serialise.pl b/tests/psa-client-server/psasim/src/psa_sim_serialise.pl
index 054841e..0dba81e 100755
--- a/tests/psa-client-server/psasim/src/psa_sim_serialise.pl
+++ b/tests/psa-client-server/psasim/src/psa_sim_serialise.pl
@@ -49,7 +49,8 @@
                psa_sign_hash_interruptible_operation_t
                psa_verify_hash_interruptible_operation_t
                mbedtls_svc_key_id_t
-               psa_key_agreement_iop_t);
+               psa_key_agreement_iop_t
+               sa_generate_key_iop_t);
 
 grep(s/-/ /g, @types);
 
diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index 032c489..7b082a3 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -693,7 +693,16 @@
     size_t public_key_length;
     uint8_t output[1024];
     size_t output_length;
+
+    uint8_t *exported = NULL;
+    size_t exported_size = 0;
+    size_t exported_length = 0;
+
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_attributes_t export_attributes = PSA_KEY_ATTRIBUTES_INIT;
+
+    mbedtls_svc_key_id_t shared_secret_id = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_attributes_t shared_secret_attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     psa_status_t status = psa_get_key_attributes(key, &attributes);
     if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
@@ -734,14 +743,53 @@
                     PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE);
     }
 
+    psa_set_key_type(&shared_secret_attributes, PSA_KEY_TYPE_DERIVE);
+    psa_set_key_usage_flags(&shared_secret_attributes, PSA_KEY_USAGE_DERIVE | PSA_KEY_USAGE_EXPORT);
+
+    status = psa_key_agreement(key, public_key, public_key_length, alg,
+                               &shared_secret_attributes, &shared_secret_id);
+
+    if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+        /* The key has been destroyed. */
+        status = PSA_SUCCESS;
+        goto exit;
+    } else if (status == PSA_SUCCESS) {
+
+        status = psa_get_key_attributes(shared_secret_id, &export_attributes);
+        if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+            /* The key has been destroyed. */
+            status = PSA_SUCCESS;
+            goto exit;
+        }
+
+        exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(psa_get_key_type(&export_attributes),
+                                                   psa_get_key_bits(&export_attributes));
+        TEST_CALLOC(exported, exported_size);
+
+        status = psa_export_key(shared_secret_id, exported, exported_size, &exported_length);
+
+        if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+            /* The key has been destroyed. */
+            status = PSA_SUCCESS;
+        }
+
+        PSA_ASSERT(status);
+    }
+
 exit:
     /*
      * Key attributes may have been returned by psa_get_key_attributes()
      * thus reset them as required.
      */
     psa_reset_key_attributes(&attributes);
+    psa_reset_key_attributes(&export_attributes);
+
+    /* Make sure to reset and free derived key attributes and slot. */
+    psa_reset_key_attributes(&shared_secret_attributes);
+    psa_destroy_key(shared_secret_id);
 
     mbedtls_free(public_key);
+    mbedtls_free(exported);
     return status;
 }
 
diff --git a/tests/src/psa_test_wrappers.c b/tests/src/psa_test_wrappers.c
index 1277df0..47d73da 100644
--- a/tests/src/psa_test_wrappers.c
+++ b/tests/src/psa_test_wrappers.c
@@ -633,7 +633,7 @@
 /* Wrapper for psa_generate_key_iop_complete */
 psa_status_t mbedtls_test_wrap_psa_generate_key_iop_complete(
     psa_generate_key_iop_t *arg0_operation,
-    psa_key_id_t *arg1_key)
+    mbedtls_svc_key_id_t *arg1_key)
 {
     psa_status_t status = (psa_generate_key_iop_complete)(arg0_operation, arg1_key);
     return status;
diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index ceaa243..d82e36a 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -7737,7 +7737,7 @@
     psa_key_type_t key_type;
 
     if (!PSA_ALG_IS_RAW_KEY_AGREEMENT(alg)) {
-        return PSA_ERROR_INVALID_ARGUMENT;
+        return PSA_ERROR_NOT_SUPPORTED;
     }
 
     key_type = psa_get_key_type(attributes);
@@ -8348,7 +8348,7 @@
 
 psa_status_t psa_generate_key_iop_complete(
     psa_generate_key_iop_t *operation,
-    psa_key_id_t *key)
+    mbedtls_svc_key_id_t *key)
 {
     (void) operation;
     (void) key;
diff --git a/tf-psa-crypto/include/psa/crypto.h b/tf-psa-crypto/include/psa/crypto.h
index 62f1dca..5a106bc 100644
--- a/tf-psa-crypto/include/psa/crypto.h
+++ b/tf-psa-crypto/include/psa/crypto.h
@@ -5501,7 +5501,7 @@
  */
 psa_status_t psa_generate_key_iop_complete(
     psa_generate_key_iop_t *operation,
-    psa_key_id_t *key);
+    mbedtls_svc_key_id_t *key);
 
 /**
  * \brief                       Abort a key generation operation.
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_op_fail.function b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_op_fail.function
index 9289869..d88b4fa 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_op_fail.function
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_op_fail.function
@@ -364,6 +364,9 @@
     size_t length = 0;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
 
+    mbedtls_svc_key_id_t shared_secret_id = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_attributes_t shared_secret_attributes = PSA_KEY_ATTRIBUTES_INIT;
+
     PSA_INIT();
 
     psa_set_key_type(&attributes, key_type);
@@ -385,6 +388,14 @@
                                       public_key, public_key_length,
                                       output, sizeof(output), &length));
 
+    psa_set_key_type(&shared_secret_attributes, PSA_KEY_TYPE_DERIVE);
+    psa_set_key_usage_flags(&shared_secret_attributes, PSA_KEY_USAGE_DERIVE | PSA_KEY_USAGE_EXPORT);
+
+    TEST_STATUS(expected_status, psa_key_agreement(key_id, public_key,
+                                                   public_key_length, alg,
+                                                   &shared_secret_attributes,
+                                                   &shared_secret_id));
+
 #if defined(PSA_WANT_ALG_HKDF) && defined(PSA_WANT_ALG_SHA_256)
     PSA_ASSERT(psa_key_derivation_setup(&operation,
                                         PSA_ALG_HKDF(PSA_ALG_SHA_256)));
@@ -403,6 +414,8 @@
     psa_key_derivation_abort(&operation);
     psa_destroy_key(key_id);
     psa_reset_key_attributes(&attributes);
+    psa_destroy_key(shared_secret_id);
+    psa_reset_key_attributes(&shared_secret_attributes);
     PSA_DONE();
 }
 /* END_CASE */