Add key_destroyable parameter to non-raw key agreement smoke tests

All current usages have this parameter set to 0 (this means the tests are unchanged).
Remove the GENERIC_ERROR return behaviour, in favour of returning the actual status.

Signed-off-by: Ryan Everett <ryan.everett@arm.com>
diff --git a/tests/include/test/psa_exercise_key.h b/tests/include/test/psa_exercise_key.h
index 2334916..f656b95 100644
--- a/tests/include/test/psa_exercise_key.h
+++ b/tests/include/test/psa_exercise_key.h
@@ -168,12 +168,15 @@
  *                          \p key.
  * \param key               A key pair object that is suitable for a key
  *                          agreement with \p operation.
+ * \param key_destroyable   If set to 1, a failure due to the key not existing
+ *                          or the key being destroyed mid-operation will only
+ *                          be reported if the error code is unexpected.
  *
  * \return                  \c 1 on success, \c 0 on failure.
  */
 psa_status_t mbedtls_test_psa_key_agreement_with_self(
     psa_key_derivation_operation_t *operation,
-    mbedtls_svc_key_id_t key);
+    mbedtls_svc_key_id_t key, int key_destroyable);
 
 /** Perform sanity checks on the given key representation.
  *
diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index b62a34b..1cf45ac 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -628,31 +628,45 @@
  * private key against its own public key. */
 psa_status_t mbedtls_test_psa_key_agreement_with_self(
     psa_key_derivation_operation_t *operation,
-    mbedtls_svc_key_id_t key)
+    mbedtls_svc_key_id_t key, int key_destroyable)
 {
     psa_key_type_t private_key_type;
     psa_key_type_t public_key_type;
     size_t key_bits;
     uint8_t *public_key = NULL;
     size_t public_key_length;
-    /* Return GENERIC_ERROR if something other than the final call to
-     * psa_key_derivation_key_agreement fails. This isn't fully satisfactory,
-     * but it's good enough: callers will report it as a failed test anyway. */
-    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    PSA_ASSERT(psa_get_key_attributes(key, &attributes));
+    psa_status_t status = psa_get_key_attributes(key, &attributes);
+    if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+        /* The key has been destroyed. */
+        psa_reset_key_attributes(&attributes);
+        return PSA_SUCCESS;
+    }
+    PSA_ASSERT(status);
+
     private_key_type = psa_get_key_type(&attributes);
     key_bits = psa_get_key_bits(&attributes);
     public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(private_key_type);
     public_key_length = PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE(public_key_type, key_bits);
     TEST_CALLOC(public_key, public_key_length);
-    PSA_ASSERT(psa_export_public_key(key, public_key, public_key_length,
-                                     &public_key_length));
+    status = psa_export_public_key(key, public_key, public_key_length,
+                                   &public_key_length);
+    if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+        /* The key has been destroyed. */
+        status = PSA_SUCCESS;
+        goto exit;
+    }
+    PSA_ASSERT(status);
 
     status = psa_key_derivation_key_agreement(
         operation, PSA_KEY_DERIVATION_INPUT_SECRET, key,
         public_key, public_key_length);
+    if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+        /* The key has been destroyed. */
+        status = PSA_SUCCESS;
+        goto exit;
+    }
 exit:
     /*
      * Key attributes may have been returned by psa_get_key_attributes()
@@ -750,7 +764,8 @@
 
 static int exercise_key_agreement_key(mbedtls_svc_key_id_t key,
                                       psa_key_usage_t usage,
-                                      psa_algorithm_t alg)
+                                      psa_algorithm_t alg,
+                                      int key_destroyable)
 {
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     unsigned char input[1] = { 0 };
@@ -781,7 +796,12 @@
            hash length. Otherwise test should fail with INVALID_ARGUMENT. */
         if (PSA_ALG_IS_HKDF_EXPAND(kdf_alg)) {
             psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-            PSA_ASSERT(psa_get_key_attributes(key, &attributes));
+            psa_status_t status = psa_get_key_attributes(key, &attributes);
+            if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+                /* The key has been destroyed. */
+                ok = 1;
+            }
+            PSA_ASSERT(status);
             size_t key_bits = psa_get_key_bits(&attributes);
             psa_algorithm_t hash_alg = PSA_ALG_HKDF_GET_HASH(kdf_alg);
 
@@ -790,7 +810,8 @@
             }
         }
 
-        TEST_EQUAL(mbedtls_test_psa_key_agreement_with_self(&operation, key),
+        TEST_EQUAL(mbedtls_test_psa_key_agreement_with_self(&operation, key,
+                                                            key_destroyable),
                    expected_key_agreement_status);
 
         if (expected_key_agreement_status != PSA_SUCCESS) {
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 9390958..8fb7d44 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -2490,7 +2490,7 @@
                               &key));
 
     PSA_ASSERT(psa_key_derivation_setup(&operation, exercise_alg));
-    status = mbedtls_test_psa_key_agreement_with_self(&operation, key);
+    status = mbedtls_test_psa_key_agreement_with_self(&operation, key, 0);
 
     TEST_EQUAL(status, expected_status);
 
@@ -8681,7 +8681,7 @@
                 // When taking a private key as secret input, use key agreement
                 // to add the shared secret to the derivation
                 TEST_EQUAL(mbedtls_test_psa_key_agreement_with_self(
-                               &operation, keys[i]),
+                               &operation, keys[i], 0),
                            expected_statuses[i]);
             } else {
                 TEST_EQUAL(psa_key_derivation_input_key(&operation, steps[i],