tests/src/psa_exercise_key.c

/** Code to exercise a PSA key object, i.e. validate that it seems well-formed
 * and can do what it is supposed to do.
 */

/*
 *  Copyright The Mbed TLS Contributors
 *  SPDX-License-Identifier: Apache-2.0
 *
 *  Licensed under the Apache License, Version 2.0 (the "License"); you may
 *  not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#include <test/helpers.h>
#include <test/macros.h>
#include <test/psa_exercise_key.h>

#if defined(MBEDTLS_PSA_CRYPTO_C)

#    include <mbedtls/asn1.h>
#    include <psa/crypto.h>

#    include <test/asn1_helpers.h>
#    include <test/psa_crypto_helpers.h>

#    if defined(MBEDTLS_PSA_CRYPTO_SE_C)
static int lifetime_is_dynamic_secure_element(psa_key_lifetime_t lifetime)
{
    return (PSA_KEY_LIFETIME_GET_LOCATION(lifetime) !=
            PSA_KEY_LOCATION_LOCAL_STORAGE);
}
#    endif

static int check_key_attributes_sanity(mbedtls_svc_key_id_t key)
{
    int ok = 0;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_key_lifetime_t lifetime;
    mbedtls_svc_key_id_t id;
    psa_key_type_t type;
    size_t bits;

    PSA_ASSERT(psa_get_key_attributes(key, &attributes));
    lifetime = psa_get_key_lifetime(&attributes);
    id = psa_get_key_id(&attributes);
    type = psa_get_key_type(&attributes);
    bits = psa_get_key_bits(&attributes);

    /* Persistence */
    if (PSA_KEY_LIFETIME_IS_VOLATILE(lifetime)) {
        TEST_ASSERT(
            (PSA_KEY_ID_VOLATILE_MIN <= MBEDTLS_SVC_KEY_ID_GET_KEY_ID(id)) &&
            (MBEDTLS_SVC_KEY_ID_GET_KEY_ID(id) <= PSA_KEY_ID_VOLATILE_MAX));
    } else {
        TEST_ASSERT(
            (PSA_KEY_ID_USER_MIN <= MBEDTLS_SVC_KEY_ID_GET_KEY_ID(id)) &&
            (MBEDTLS_SVC_KEY_ID_GET_KEY_ID(id) <= PSA_KEY_ID_USER_MAX));
    }
#    if defined(MBEDTLS_PSA_CRYPTO_SE_C)
    /* randomly-generated 64-bit constant, should never appear in test data */
    psa_key_slot_number_t slot_number = 0xec94d4a5058a1a21;
    psa_status_t status = psa_get_key_slot_number(&attributes, &slot_number);
    if (lifetime_is_dynamic_secure_element(lifetime)) {
        /* Mbed Crypto currently always exposes the slot number to
         * applications. This is not mandated by the PSA specification
         * and may change in future versions. */
        TEST_EQUAL(status, 0);
        TEST_ASSERT(slot_number != 0xec94d4a5058a1a21);
    } else {
        TEST_EQUAL(status, PSA_ERROR_INVALID_ARGUMENT);
    }
#    endif

    /* Type and size */
    TEST_ASSERT(type != 0);
    TEST_ASSERT(bits != 0);
    TEST_ASSERT(bits <= PSA_MAX_KEY_BITS);
    if (PSA_KEY_TYPE_IS_UNSTRUCTURED(type))
        TEST_ASSERT(bits % 8 == 0);

    /* MAX macros concerning specific key types */
    if (PSA_KEY_TYPE_IS_ECC(type))
        TEST_ASSERT(bits <= PSA_VENDOR_ECC_MAX_CURVE_BITS);
    else if (PSA_KEY_TYPE_IS_RSA(type))
        TEST_ASSERT(bits <= PSA_VENDOR_RSA_MAX_KEY_BITS);
    TEST_ASSERT(PSA_BLOCK_CIPHER_BLOCK_LENGTH(type) <=
                PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE);

    ok = 1;

exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes(&attributes);

    return ok;
}

static int exercise_mac_key(mbedtls_svc_key_id_t key,
                            psa_key_usage_t usage,
                            psa_algorithm_t alg)
{
    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
    const unsigned char input[] = "foo";
    unsigned char mac[PSA_MAC_MAX_SIZE] = { 0 };
    size_t mac_length = sizeof(mac);

    /* Convert wildcard algorithm to exercisable algorithm */
    if (alg & PSA_ALG_MAC_AT_LEAST_THIS_LENGTH_FLAG) {
        alg = PSA_ALG_TRUNCATED_MAC(alg, PSA_MAC_TRUNCATED_LENGTH(alg));
    }

    if (usage & PSA_KEY_USAGE_SIGN_HASH) {
        PSA_ASSERT(psa_mac_sign_setup(&operation, key, alg));
        PSA_ASSERT(psa_mac_update(&operation, input, sizeof(input)));
        PSA_ASSERT(
            psa_mac_sign_finish(&operation, mac, sizeof(mac), &mac_length));
    }

    if (usage & PSA_KEY_USAGE_VERIFY_HASH) {
        psa_status_t verify_status = (usage & PSA_KEY_USAGE_SIGN_HASH ?
                                          PSA_SUCCESS :
                                          PSA_ERROR_INVALID_SIGNATURE);
        PSA_ASSERT(psa_mac_verify_setup(&operation, key, alg));
        PSA_ASSERT(psa_mac_update(&operation, input, sizeof(input)));
        TEST_EQUAL(psa_mac_verify_finish(&operation, mac, mac_length),
                   verify_status);
    }

    return 1;

exit:
    psa_mac_abort(&operation);
    return 0;
}

static int exercise_cipher_key(mbedtls_svc_key_id_t key,
                               psa_key_usage_t usage,
                               psa_algorithm_t alg)
{
    psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
    unsigned char iv[16] = { 0 };
    size_t iv_length = sizeof(iv);
    const unsigned char plaintext[16] = "Hello, world...";
    unsigned char ciphertext[32] = "(wabblewebblewibblewobblewubble)";
    size_t ciphertext_length = sizeof(ciphertext);
    unsigned char decrypted[sizeof(ciphertext)];
    size_t part_length;

    if (usage & PSA_KEY_USAGE_ENCRYPT) {
        PSA_ASSERT(psa_cipher_encrypt_setup(&operation, key, alg));
        PSA_ASSERT(
            psa_cipher_generate_iv(&operation, iv, sizeof(iv), &iv_length));
        PSA_ASSERT(psa_cipher_update(&operation, plaintext, sizeof(plaintext),
                                     ciphertext, sizeof(ciphertext),
                                     &ciphertext_length));
        PSA_ASSERT(psa_cipher_finish(&operation, ciphertext + ciphertext_length,
                                     sizeof(ciphertext) - ciphertext_length,
                                     &part_length));
        ciphertext_length += part_length;
    }

    if (usage & PSA_KEY_USAGE_DECRYPT) {
        psa_status_t status;
        int maybe_invalid_padding = 0;
        if (!(usage & PSA_KEY_USAGE_ENCRYPT)) {
            psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
            PSA_ASSERT(psa_get_key_attributes(key, &attributes));
            /* This should be PSA_CIPHER_GET_IV_SIZE but the API doesn't
             * have this macro yet. */
            iv_length =
                PSA_BLOCK_CIPHER_BLOCK_LENGTH(psa_get_key_type(&attributes));
            maybe_invalid_padding = !PSA_ALG_IS_STREAM_CIPHER(alg);
            psa_reset_key_attributes(&attributes);
        }
        PSA_ASSERT(psa_cipher_decrypt_setup(&operation, key, alg));
        PSA_ASSERT(psa_cipher_set_iv(&operation, iv, iv_length));
        PSA_ASSERT(psa_cipher_update(&operation, ciphertext, ciphertext_length,
                                     decrypted, sizeof(decrypted),
                                     &part_length));
        status = psa_cipher_finish(&operation, decrypted + part_length,
                                   sizeof(decrypted) - part_length,
                                   &part_length);
        /* For a stream cipher, all inputs are valid. For a block cipher,
         * if the input is some aribtrary data rather than an actual
         ciphertext, a padding error is likely.  */
        if (maybe_invalid_padding)
            TEST_ASSERT(status == PSA_SUCCESS ||
                        status == PSA_ERROR_INVALID_PADDING);
        else
            PSA_ASSERT(status);
    }

    return 1;

exit:
    psa_cipher_abort(&operation);
    return 0;
}

static int exercise_aead_key(mbedtls_svc_key_id_t key,
                             psa_key_usage_t usage,
                             psa_algorithm_t alg)
{
    unsigned char nonce[16] = { 0 };
    size_t nonce_length = sizeof(nonce);
    unsigned char plaintext[16] = "Hello, world...";
    unsigned char ciphertext[48] = "(wabblewebblewibblewobblewubble)";
    size_t ciphertext_length = sizeof(ciphertext);
    size_t plaintext_length = sizeof(ciphertext);

    /* Convert wildcard algorithm to exercisable algorithm */
    if (alg & PSA_ALG_AEAD_AT_LEAST_THIS_LENGTH_FLAG) {
        alg = PSA_ALG_AEAD_WITH_SHORTENED_TAG(alg,
                                              PSA_ALG_AEAD_GET_TAG_LENGTH(alg));
    }

    /* Default IV length for AES-GCM is 12 bytes */
    if (PSA_ALG_AEAD_WITH_SHORTENED_TAG(alg, 0) ==
        PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM, 0)) {
        nonce_length = 12;
    }

    /* IV length for CCM needs to be between 7 and 13 bytes */
    if (PSA_ALG_AEAD_WITH_SHORTENED_TAG(alg, 0) ==
        PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_CCM, 0)) {
        nonce_length = 12;
    }

    if (usage & PSA_KEY_USAGE_ENCRYPT) {
        PSA_ASSERT(psa_aead_encrypt(key, alg, nonce, nonce_length, NULL, 0,
                                    plaintext, sizeof(plaintext), ciphertext,
                                    sizeof(ciphertext), &ciphertext_length));
    }

    if (usage & PSA_KEY_USAGE_DECRYPT) {
        psa_status_t verify_status = (usage & PSA_KEY_USAGE_ENCRYPT ?
                                          PSA_SUCCESS :
                                          PSA_ERROR_INVALID_SIGNATURE);
        TEST_EQUAL(psa_aead_decrypt(key, alg, nonce, nonce_length, NULL, 0,
                                    ciphertext, ciphertext_length, plaintext,
                                    sizeof(plaintext), &plaintext_length),
                   verify_status);
    }

    return 1;

exit:
    return 0;
}

static int exercise_signature_key(mbedtls_svc_key_id_t key,
                                  psa_key_usage_t usage,
                                  psa_algorithm_t alg)
{
    if (usage & (PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH)) {
        unsigned char payload[PSA_HASH_MAX_SIZE] = { 1 };
        size_t payload_length = 16;
        unsigned char signature[PSA_SIGNATURE_MAX_SIZE] = { 0 };
        size_t signature_length = sizeof(signature);
        psa_algorithm_t hash_alg = PSA_ALG_SIGN_GET_HASH(alg);

        /* If the policy allows signing with any hash, just pick one. */
        if (PSA_ALG_IS_HASH_AND_SIGN(alg) && hash_alg == PSA_ALG_ANY_HASH) {
#    if defined(KNOWN_SUPPORTED_HASH_ALG)
            hash_alg = KNOWN_SUPPORTED_HASH_ALG;
            alg ^= PSA_ALG_ANY_HASH ^ hash_alg;
#    else
            TEST_ASSERT(!"No hash algorithm for hash-and-sign testing");
#    endif
        }

        /* Some algorithms require the payload to have the size of
         * the hash encoded in the algorithm. Use this input size
         * even for algorithms that allow other input sizes. */
        if (hash_alg != 0)
            payload_length = PSA_HASH_LENGTH(hash_alg);

        if (usage & PSA_KEY_USAGE_SIGN_HASH) {
            PSA_ASSERT(psa_sign_hash(key, alg, payload, payload_length,
                                     signature, sizeof(signature),
                                     &signature_length));
        }

        if (usage & PSA_KEY_USAGE_VERIFY_HASH) {
            psa_status_t verify_status = (usage & PSA_KEY_USAGE_SIGN_HASH ?
                                              PSA_SUCCESS :
                                              PSA_ERROR_INVALID_SIGNATURE);
            TEST_EQUAL(psa_verify_hash(key, alg, payload, payload_length,
                                       signature, signature_length),
                       verify_status);
        }
    }

    if (usage & (PSA_KEY_USAGE_SIGN_MESSAGE | PSA_KEY_USAGE_VERIFY_MESSAGE)) {
        unsigned char message[256] = "Hello, world...";
        unsigned char signature[PSA_SIGNATURE_MAX_SIZE] = { 0 };
        size_t message_length = 16;
        size_t signature_length = sizeof(signature);

        if (usage & PSA_KEY_USAGE_SIGN_MESSAGE) {
            PSA_ASSERT(psa_sign_message(key, alg, message, message_length,
                                        signature, sizeof(signature),
                                        &signature_length));
        }

        if (usage & PSA_KEY_USAGE_VERIFY_MESSAGE) {
            psa_status_t verify_status = (usage & PSA_KEY_USAGE_SIGN_MESSAGE ?
                                              PSA_SUCCESS :
                                              PSA_ERROR_INVALID_SIGNATURE);
            TEST_EQUAL(psa_verify_message(key, alg, message, message_length,
                                          signature, signature_length),
                       verify_status);
        }
    }

    return 1;

exit:
    return 0;
}

static int exercise_asymmetric_encryption_key(mbedtls_svc_key_id_t key,
                                              psa_key_usage_t usage,
                                              psa_algorithm_t alg)
{
    unsigned char plaintext[256] = "Hello, world...";
    unsigned char ciphertext[256] = "(wabblewebblewibblewobblewubble)";
    size_t ciphertext_length = sizeof(ciphertext);
    size_t plaintext_length = 16;

    if (usage & PSA_KEY_USAGE_ENCRYPT) {
        PSA_ASSERT(psa_asymmetric_encrypt(
            key, alg, plaintext, plaintext_length, NULL, 0, ciphertext,
            sizeof(ciphertext), &ciphertext_length));
    }

    if (usage & PSA_KEY_USAGE_DECRYPT) {
        psa_status_t status = psa_asymmetric_decrypt(
            key, alg, ciphertext, ciphertext_length, NULL, 0, plaintext,
            sizeof(plaintext), &plaintext_length);
        TEST_ASSERT(status == PSA_SUCCESS ||
                    ((usage & PSA_KEY_USAGE_ENCRYPT) == 0 &&
                     (status == PSA_ERROR_INVALID_ARGUMENT ||
                      status == PSA_ERROR_INVALID_PADDING)));
    }

    return 1;

exit:
    return 0;
}

int mbedtls_test_psa_setup_key_derivation_wrap(
    psa_key_derivation_operation_t *operation,
    mbedtls_svc_key_id_t key,
    psa_algorithm_t alg,
    const unsigned char *input1,
    size_t input1_length,
    const unsigned char *input2,
    size_t input2_length,
    size_t capacity)
{
    PSA_ASSERT(psa_key_derivation_setup(operation, alg));
    if (PSA_ALG_IS_HKDF(alg)) {
        PSA_ASSERT(psa_key_derivation_input_bytes(
            operation, PSA_KEY_DERIVATION_INPUT_SALT, input1, input1_length));
        PSA_ASSERT(psa_key_derivation_input_key(
            operation, PSA_KEY_DERIVATION_INPUT_SECRET, key));
        PSA_ASSERT(psa_key_derivation_input_bytes(
            operation, PSA_KEY_DERIVATION_INPUT_INFO, input2, input2_length));
    } else if (PSA_ALG_IS_TLS12_PRF(alg) || PSA_ALG_IS_TLS12_PSK_TO_MS(alg)) {
        PSA_ASSERT(psa_key_derivation_input_bytes(
            operation, PSA_KEY_DERIVATION_INPUT_SEED, input1, input1_length));
        PSA_ASSERT(psa_key_derivation_input_key(
            operation, PSA_KEY_DERIVATION_INPUT_SECRET, key));
        PSA_ASSERT(psa_key_derivation_input_bytes(
            operation, PSA_KEY_DERIVATION_INPUT_LABEL, input2, input2_length));
    } else {
        TEST_ASSERT(!"Key derivation algorithm not supported");
    }

    if (capacity != SIZE_MAX)
        PSA_ASSERT(psa_key_derivation_set_capacity(operation, capacity));

    return 1;

exit:
    return 0;
}

static int exercise_key_derivation_key(mbedtls_svc_key_id_t key,
                                       psa_key_usage_t usage,
                                       psa_algorithm_t alg)
{
    psa_key_derivation_operation_t operation =
        PSA_KEY_DERIVATION_OPERATION_INIT;
    unsigned char input1[] = "Input 1";
    size_t input1_length = sizeof(input1);
    unsigned char input2[] = "Input 2";
    size_t input2_length = sizeof(input2);
    unsigned char output[1];
    size_t capacity = sizeof(output);

    if (usage & PSA_KEY_USAGE_DERIVE) {
        if (!mbedtls_test_psa_setup_key_derivation_wrap(
                &operation, key, alg, input1, input1_length, input2,
                input2_length, capacity))
            goto exit;

        PSA_ASSERT(
            psa_key_derivation_output_bytes(&operation, output, capacity));
        PSA_ASSERT(psa_key_derivation_abort(&operation));
    }

    return 1;

exit:
    return 0;
}

/* We need two keys to exercise key agreement. Exercise the
 * 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)
{
    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));
    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);
    ASSERT_ALLOC(public_key, public_key_length);
    PSA_ASSERT(psa_export_public_key(key, public_key, public_key_length,
                                     &public_key_length));

    status = psa_key_derivation_key_agreement(operation,
                                              PSA_KEY_DERIVATION_INPUT_SECRET,
                                              key, public_key,
                                              public_key_length);
exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes(&attributes);

    mbedtls_free(public_key);
    return status;
}

/* We need two keys to exercise key agreement. Exercise the
 * private key against its own public key. */
psa_status_t
mbedtls_test_psa_raw_key_agreement_with_self(psa_algorithm_t alg,
                                             mbedtls_svc_key_id_t key)
{
    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;
    uint8_t output[1024];
    size_t output_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));
    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);
    ASSERT_ALLOC(public_key, public_key_length);
    PSA_ASSERT(psa_export_public_key(key, public_key, public_key_length,
                                     &public_key_length));

    status = psa_raw_key_agreement(alg, key, public_key, public_key_length,
                                   output, sizeof(output), &output_length);
    if (status == PSA_SUCCESS) {
        TEST_ASSERT(output_length <= PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE(
                                         private_key_type, key_bits));
        TEST_ASSERT(output_length <= PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE);
    }

exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes(&attributes);

    mbedtls_free(public_key);
    return status;
}

static int exercise_raw_key_agreement_key(mbedtls_svc_key_id_t key,
                                          psa_key_usage_t usage,
                                          psa_algorithm_t alg)
{
    int ok = 0;

    if (usage & PSA_KEY_USAGE_DERIVE) {
        /* We need two keys to exercise key agreement. Exercise the
         * private key against its own public key. */
        PSA_ASSERT(mbedtls_test_psa_raw_key_agreement_with_self(alg, key));
    }
    ok = 1;

exit:
    return ok;
}

static int exercise_key_agreement_key(mbedtls_svc_key_id_t key,
                                      psa_key_usage_t usage,
                                      psa_algorithm_t alg)
{
    psa_key_derivation_operation_t operation =
        PSA_KEY_DERIVATION_OPERATION_INIT;
    unsigned char output[1];
    int ok = 0;

    if (usage & PSA_KEY_USAGE_DERIVE) {
        /* We need two keys to exercise key agreement. Exercise the
         * private key against its own public key. */
        PSA_ASSERT(psa_key_derivation_setup(&operation, alg));
        PSA_ASSERT(mbedtls_test_psa_key_agreement_with_self(&operation, key));
        PSA_ASSERT(psa_key_derivation_output_bytes(&operation, output,
                                                   sizeof(output)));
        PSA_ASSERT(psa_key_derivation_abort(&operation));
    }
    ok = 1;

exit:
    return ok;
}

int mbedtls_test_psa_exported_key_sanity_check(psa_key_type_t type,
                                               size_t bits,
                                               const uint8_t *exported,
                                               size_t exported_length)
{
    TEST_ASSERT(exported_length <= PSA_EXPORT_KEY_OUTPUT_SIZE(type, bits));

    if (PSA_KEY_TYPE_IS_UNSTRUCTURED(type))
        TEST_EQUAL(exported_length, PSA_BITS_TO_BYTES(bits));
    else
#    if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_PK_PARSE_C)
        if (type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
        uint8_t *p = (uint8_t *)exported;
        const uint8_t *end = exported + exported_length;
        size_t len;
        /*   RSAPrivateKey ::= SEQUENCE {
         *       version             INTEGER,  -- must be 0
         *       modulus             INTEGER,  -- n
         *       publicExponent      INTEGER,  -- e
         *       privateExponent     INTEGER,  -- d
         *       prime1              INTEGER,  -- p
         *       prime2              INTEGER,  -- q
         *       exponent1           INTEGER,  -- d mod (p-1)
         *       exponent2           INTEGER,  -- d mod (q-1)
         *       coefficient         INTEGER,  -- (inverse of q) mod p
         *   }
         */
        TEST_EQUAL(mbedtls_asn1_get_tag(&p, end, &len,
                                        MBEDTLS_ASN1_SEQUENCE |
                                            MBEDTLS_ASN1_CONSTRUCTED),
                   0);
        TEST_EQUAL(p + len, end);
        if (!mbedtls_test_asn1_skip_integer(&p, end, 0, 0, 0))
            goto exit;
        if (!mbedtls_test_asn1_skip_integer(&p, end, bits, bits, 1))
            goto exit;
        if (!mbedtls_test_asn1_skip_integer(&p, end, 2, bits, 1))
            goto exit;
        /* Require d to be at least half the size of n. */
        if (!mbedtls_test_asn1_skip_integer(&p, end, bits / 2, bits, 1))
            goto exit;
        /* Require p and q to be at most half the size of n, rounded up. */
        if (!mbedtls_test_asn1_skip_integer(&p, end, bits / 2, bits / 2 + 1, 1))
            goto exit;
        if (!mbedtls_test_asn1_skip_integer(&p, end, bits / 2, bits / 2 + 1, 1))
            goto exit;
        if (!mbedtls_test_asn1_skip_integer(&p, end, 1, bits / 2 + 1, 0))
            goto exit;
        if (!mbedtls_test_asn1_skip_integer(&p, end, 1, bits / 2 + 1, 0))
            goto exit;
        if (!mbedtls_test_asn1_skip_integer(&p, end, 1, bits / 2 + 1, 0))
            goto exit;
        TEST_EQUAL(p, end);

        TEST_ASSERT(exported_length <= PSA_EXPORT_KEY_PAIR_MAX_SIZE);
    } else
#    endif /* MBEDTLS_RSA_C */

#    if defined(MBEDTLS_ECP_C)
        if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(type)) {
        /* Just the secret value */
        TEST_EQUAL(exported_length, PSA_BITS_TO_BYTES(bits));

        TEST_ASSERT(exported_length <= PSA_EXPORT_KEY_PAIR_MAX_SIZE);
    } else
#    endif /* MBEDTLS_ECP_C */

#    if defined(MBEDTLS_RSA_C)
        if (type == PSA_KEY_TYPE_RSA_PUBLIC_KEY) {
        uint8_t *p = (uint8_t *)exported;
        const uint8_t *end = exported + exported_length;
        size_t len;
        /*   RSAPublicKey ::= SEQUENCE {
         *      modulus            INTEGER,    -- n
         *      publicExponent     INTEGER  }  -- e
         */
        TEST_EQUAL(mbedtls_asn1_get_tag(&p, end, &len,
                                        MBEDTLS_ASN1_SEQUENCE |
                                            MBEDTLS_ASN1_CONSTRUCTED),
                   0);
        TEST_EQUAL(p + len, end);
        if (!mbedtls_test_asn1_skip_integer(&p, end, bits, bits, 1))
            goto exit;
        if (!mbedtls_test_asn1_skip_integer(&p, end, 2, bits, 1))
            goto exit;
        TEST_EQUAL(p, end);

        TEST_ASSERT(exported_length <=
                    PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE(type, bits));
        TEST_ASSERT(exported_length <= PSA_EXPORT_PUBLIC_KEY_MAX_SIZE);
    } else
#    endif /* MBEDTLS_RSA_C */

#    if defined(MBEDTLS_ECP_C)
        if (PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY(type)) {

        TEST_ASSERT(exported_length <=
                    PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE(type, bits));
        TEST_ASSERT(exported_length <= PSA_EXPORT_PUBLIC_KEY_MAX_SIZE);

        if (PSA_KEY_TYPE_ECC_GET_FAMILY(type) == PSA_ECC_FAMILY_MONTGOMERY) {
            /* The representation of an ECC Montgomery public key is
             * the raw compressed point */
            TEST_EQUAL(PSA_BITS_TO_BYTES(bits), exported_length);
        } else {
            /* The representation of an ECC Weierstrass public key is:
             *      - The byte 0x04;
             *      - `x_P` as a `ceiling(m/8)`-byte string, big-endian;
             *      - `y_P` as a `ceiling(m/8)`-byte string, big-endian;
             *      - where m is the bit size associated with the curve.
             */
            TEST_EQUAL(1 + 2 * PSA_BITS_TO_BYTES(bits), exported_length);
            TEST_EQUAL(exported[0], 4);
        }
    } else
#    endif /* MBEDTLS_ECP_C */

    {
        TEST_ASSERT(!"Sanity check not implemented for this key type");
    }

#    if defined(MBEDTLS_DES_C)
    if (type == PSA_KEY_TYPE_DES) {
        /* Check the parity bits. */
        unsigned i;
        for (i = 0; i < bits / 8; i++) {
            unsigned bit_count = 0;
            unsigned m;
            for (m = 1; m <= 0x100; m <<= 1) {
                if (exported[i] & m)
                    ++bit_count;
            }
            TEST_ASSERT(bit_count % 2 != 0);
        }
    }
#    endif

    return 1;

exit:
    return 0;
}

static int exercise_export_key(mbedtls_svc_key_id_t key, psa_key_usage_t usage)
{
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    uint8_t *exported = NULL;
    size_t exported_size = 0;
    size_t exported_length = 0;
    int ok = 0;

    PSA_ASSERT(psa_get_key_attributes(key, &attributes));

    exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(psa_get_key_type(&attributes),
                                               psa_get_key_bits(&attributes));
    ASSERT_ALLOC(exported, exported_size);

    if ((usage & PSA_KEY_USAGE_EXPORT) == 0 &&
        !PSA_KEY_TYPE_IS_PUBLIC_KEY(psa_get_key_type(&attributes))) {
        TEST_EQUAL(psa_export_key(key, exported, exported_size,
                                  &exported_length),
                   PSA_ERROR_NOT_PERMITTED);
        ok = 1;
        goto exit;
    }

    PSA_ASSERT(psa_export_key(key, exported, exported_size, &exported_length));
    ok = mbedtls_test_psa_exported_key_sanity_check(
        psa_get_key_type(&attributes), psa_get_key_bits(&attributes), exported,
        exported_length);

exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes(&attributes);

    mbedtls_free(exported);
    return ok;
}

static int exercise_export_public_key(mbedtls_svc_key_id_t key)
{
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_key_type_t public_type;
    uint8_t *exported = NULL;
    size_t exported_size = 0;
    size_t exported_length = 0;
    int ok = 0;

    PSA_ASSERT(psa_get_key_attributes(key, &attributes));
    if (!PSA_KEY_TYPE_IS_ASYMMETRIC(psa_get_key_type(&attributes))) {
        exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(
            psa_get_key_type(&attributes), psa_get_key_bits(&attributes));
        ASSERT_ALLOC(exported, exported_size);

        TEST_EQUAL(psa_export_public_key(key, exported, exported_size,
                                         &exported_length),
                   PSA_ERROR_INVALID_ARGUMENT);
        ok = 1;
        goto exit;
    }

    public_type =
        PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(psa_get_key_type(&attributes));
    exported_size = PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE(
        public_type, psa_get_key_bits(&attributes));
    ASSERT_ALLOC(exported, exported_size);

    PSA_ASSERT(
        psa_export_public_key(key, exported, exported_size, &exported_length));
    ok = mbedtls_test_psa_exported_key_sanity_check(
        public_type, psa_get_key_bits(&attributes), exported, exported_length);

exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes(&attributes);

    mbedtls_free(exported);
    return ok;
}

int mbedtls_test_psa_exercise_key(mbedtls_svc_key_id_t key,
                                  psa_key_usage_t usage,
                                  psa_algorithm_t alg)
{
    int ok = 0;

    if (!check_key_attributes_sanity(key))
        return 0;

    if (alg == 0)
        ok = 1; /* If no algorihm, do nothing (used for raw data "keys"). */
    else if (PSA_ALG_IS_MAC(alg))
        ok = exercise_mac_key(key, usage, alg);
    else if (PSA_ALG_IS_CIPHER(alg))
        ok = exercise_cipher_key(key, usage, alg);
    else if (PSA_ALG_IS_AEAD(alg))
        ok = exercise_aead_key(key, usage, alg);
    else if (PSA_ALG_IS_SIGN(alg))
        ok = exercise_signature_key(key, usage, alg);
    else if (PSA_ALG_IS_ASYMMETRIC_ENCRYPTION(alg))
        ok = exercise_asymmetric_encryption_key(key, usage, alg);
    else if (PSA_ALG_IS_KEY_DERIVATION(alg))
        ok = exercise_key_derivation_key(key, usage, alg);
    else if (PSA_ALG_IS_RAW_KEY_AGREEMENT(alg))
        ok = exercise_raw_key_agreement_key(key, usage, alg);
    else if (PSA_ALG_IS_KEY_AGREEMENT(alg))
        ok = exercise_key_agreement_key(key, usage, alg);
    else
        TEST_ASSERT(!"No code to exercise this category of algorithm");

    ok = ok && exercise_export_key(key, usage);
    ok = ok && exercise_export_public_key(key);

exit:
    return ok;
}

psa_key_usage_t mbedtls_test_psa_usage_to_exercise(psa_key_type_t type,
                                                   psa_algorithm_t alg)
{
    if (PSA_ALG_IS_MAC(alg) || PSA_ALG_IS_SIGN(alg)) {
        if (PSA_ALG_IS_HASH_AND_SIGN(alg)) {
            if (PSA_ALG_SIGN_GET_HASH(alg))
                return (PSA_KEY_TYPE_IS_PUBLIC_KEY(type) ?
                            PSA_KEY_USAGE_VERIFY_HASH |
                                PSA_KEY_USAGE_VERIFY_MESSAGE :
                            PSA_KEY_USAGE_SIGN_HASH |
                                PSA_KEY_USAGE_VERIFY_HASH |
                                PSA_KEY_USAGE_SIGN_MESSAGE |
                                PSA_KEY_USAGE_VERIFY_MESSAGE);
        } else if (PSA_ALG_IS_SIGN_MESSAGE(alg))
            return (PSA_KEY_TYPE_IS_PUBLIC_KEY(type) ?
                        PSA_KEY_USAGE_VERIFY_MESSAGE :
                        PSA_KEY_USAGE_SIGN_MESSAGE |
                            PSA_KEY_USAGE_VERIFY_MESSAGE);

        return (PSA_KEY_TYPE_IS_PUBLIC_KEY(type) ?
                    PSA_KEY_USAGE_VERIFY_HASH :
                    PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH);
    } else if (PSA_ALG_IS_CIPHER(alg) || PSA_ALG_IS_AEAD(alg) ||
               PSA_ALG_IS_ASYMMETRIC_ENCRYPTION(alg)) {
        return (PSA_KEY_TYPE_IS_PUBLIC_KEY(type) ?
                    PSA_KEY_USAGE_ENCRYPT :
                    PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT);
    } else if (PSA_ALG_IS_KEY_DERIVATION(alg) ||
               PSA_ALG_IS_KEY_AGREEMENT(alg)) {
        return PSA_KEY_USAGE_DERIVE;
    } else {
        return 0;
    }
}

#endif /* MBEDTLS_PSA_CRYPTO_C */