From 2ec141a429635e2a764be7393a57e1ebf00b0545 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Thu, 15 Feb 2024 17:22:37 +0100
Subject: [PATCH] After pk_import_into_psa, test that the keys match

We were testing the internal consistency of the resulting key, and that the
resulting key had the right metadata, but we were not testing that the PSA
key had the expected key material. Comparing the public keys fixes that.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
---
 tests/suites/test_suite_pk.function      | 16 ++++++++++++++++
 tests/suites/test_suite_pkparse.function |  3 +++
 2 files changed, 19 insertions(+)

diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 4ea7a70157..39e3846524 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -19,6 +19,8 @@
 #include "psa/crypto.h"
 #include "mbedtls/psa_util.h"
 
+#include <test/psa_exercise_key.h>
+
 /* Used for properly sizing the key buffer in pk_genkey_ec() */
 #include "psa_util_internal.h"
 
@@ -1893,6 +1895,9 @@ void pk_get_psa_attributes(int pk_type, int from_pair,
     TEST_EQUAL(psa_get_key_enrollment_algorithm(&attributes), PSA_ALG_NONE);
 
     TEST_EQUAL(mbedtls_pk_import_into_psa(&pk, &attributes, &new_key_id), 0);
+    if (!mbedtls_test_key_consistency_psa_pk(new_key_id, &pk)) {
+        goto exit;
+    }
 
 exit:
     mbedtls_pk_free(&pk);
@@ -1938,6 +1943,9 @@ void pk_rsa_v21_get_psa_attributes(int md_type, int from_pair,
     TEST_EQUAL(psa_get_key_enrollment_algorithm(&attributes), PSA_ALG_NONE);
 
     TEST_EQUAL(mbedtls_pk_import_into_psa(&pk, &attributes, &new_key_id), 0);
+    if (!mbedtls_test_key_consistency_psa_pk(new_key_id, &pk)) {
+        goto exit;
+    }
 
 exit:
     mbedtls_pk_free(&pk);
@@ -2026,6 +2034,11 @@ void pk_get_psa_attributes_opaque(int from_type_arg, int from_bits_arg,
         }
         TEST_EQUAL(mbedtls_pk_import_into_psa(&pk, &attributes, &new_key_id),
                    expected_import_ret);
+        if (expected_import_ret == 0) {
+            if (!mbedtls_test_key_consistency_psa_pk(new_key_id, &pk)) {
+                goto exit;
+            }
+        }
     }
 
 exit:
@@ -2117,6 +2130,9 @@ void pk_import_into_psa_opaque(int from_type, int from_bits,
             expected_usage |= PSA_KEY_USAGE_VERIFY_MESSAGE;
         }
         TEST_EQUAL(expected_usage, psa_get_key_usage_flags(&actual_attributes));
+        if (!mbedtls_test_key_consistency_psa_pk(to_key_id, &pk)) {
+            goto exit;
+        }
     } else {
         TEST_ASSERT(mbedtls_svc_key_id_equal(to_key_id, MBEDTLS_SVC_KEY_ID_INIT));
     }
diff --git a/tests/suites/test_suite_pkparse.function b/tests/suites/test_suite_pkparse.function
index add225d735..f275c324d5 100644
--- a/tests/suites/test_suite_pkparse.function
+++ b/tests/suites/test_suite_pkparse.function
@@ -48,6 +48,9 @@ static int test_psa_bridge(const mbedtls_pk_context *ctx,
 
     TEST_EQUAL(mbedtls_pk_get_psa_attributes(ctx, usage_flag, &attributes), 0);
     TEST_EQUAL(mbedtls_pk_import_into_psa(ctx, &attributes, &psa_key), 0);
+    if (!mbedtls_test_key_consistency_psa_pk(psa_key, ctx)) {
+        goto exit;
+    }
 
     psa_algorithm_t exercise_usage = psa_get_key_usage_flags(&attributes);
     psa_algorithm_t exercise_alg = psa_get_key_algorithm(&attributes);