/* BEGIN_HEADER */

#include "psa/crypto.h"
#include "test/psa_crypto_helpers.h"

static int test_equal_status(const char *test,
                             int line_no, const char *filename,
                             psa_status_t value1,
                             psa_status_t value2)
{
    if ((value1 == PSA_ERROR_INVALID_ARGUMENT &&
         value2 == PSA_ERROR_NOT_SUPPORTED) ||
        (value1 == PSA_ERROR_NOT_SUPPORTED &&
         value2 == PSA_ERROR_INVALID_ARGUMENT)) {
        return 1;
    }
    return mbedtls_test_equal(test, line_no, filename, value1, value2);
}

/** Like #TEST_EQUAL, but expects #psa_status_t values and treats
 * #PSA_ERROR_INVALID_ARGUMENT and #PSA_ERROR_NOT_SUPPORTED as
 * interchangeable.
 *
 * This test suite currently allows NOT_SUPPORTED and INVALID_ARGUMENT
 * to be interchangeable in places where the library's behavior does not
 * match the strict expectations of the test case generator. In the long
 * run, it would be better to clarify the expectations and reconcile the
 * library and the test case generator.
 */
#define TEST_STATUS(expr1, expr2)                                     \
    do {                                                                \
        if (!test_equal_status( #expr1 " == " #expr2, __LINE__, __FILE__, \
                                expr1, expr2))                       \
        goto exit;                                                  \
    } while (0)

/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_PSA_CRYPTO_C
 * END_DEPENDENCIES
 */

/* BEGIN_CASE */
void hash_fail(int alg_arg, int expected_status_arg)
{
    psa_status_t expected_status = expected_status_arg;
    psa_algorithm_t alg = alg_arg;
    psa_hash_operation_t operation = PSA_HASH_OPERATION_INIT;
    uint8_t input[1] = { 'A' };
    uint8_t output[PSA_HASH_MAX_SIZE] = { 0 };
    size_t length = SIZE_MAX;

    PSA_INIT();

    TEST_EQUAL(expected_status,
               psa_hash_setup(&operation, alg));
    TEST_EQUAL(expected_status,
               psa_hash_compute(alg, input, sizeof(input),
                                output, sizeof(output), &length));
    TEST_EQUAL(expected_status,
               psa_hash_compare(alg, input, sizeof(input),
                                output, sizeof(output)));

exit:
    psa_hash_abort(&operation);
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE */
void mac_fail(int key_type_arg, data_t *key_data,
              int alg_arg, int expected_status_arg)
{
    psa_status_t expected_status = expected_status_arg;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
    uint8_t input[1] = { 'A' };
    uint8_t output[PSA_MAC_MAX_SIZE] = { 0 };
    size_t length = SIZE_MAX;

    PSA_INIT();

    psa_set_key_type(&attributes, key_type);
    psa_set_key_usage_flags(&attributes,
                            PSA_KEY_USAGE_SIGN_HASH |
                            PSA_KEY_USAGE_VERIFY_HASH);
    psa_set_key_algorithm(&attributes, alg);
    PSA_ASSERT(psa_import_key(&attributes,
                              key_data->x, key_data->len,
                              &key_id));

    TEST_STATUS(expected_status,
                psa_mac_sign_setup(&operation, key_id, alg));
    TEST_STATUS(expected_status,
                psa_mac_verify_setup(&operation, key_id, alg));
    TEST_STATUS(expected_status,
                psa_mac_compute(key_id, alg,
                                input, sizeof(input),
                                output, sizeof(output), &length));
    TEST_STATUS(expected_status,
                psa_mac_verify(key_id, alg,
                               input, sizeof(input),
                               output, sizeof(output)));

exit:
    psa_mac_abort(&operation);
    psa_destroy_key(key_id);
    psa_reset_key_attributes(&attributes);
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE */
void cipher_fail(int key_type_arg, data_t *key_data,
                 int alg_arg, int expected_status_arg)
{
    psa_status_t expected_status = expected_status_arg;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
    uint8_t input[1] = { 'A' };
    uint8_t output[64] = { 0 };
    size_t length = SIZE_MAX;

    PSA_INIT();

    psa_set_key_type(&attributes, key_type);
    psa_set_key_usage_flags(&attributes,
                            PSA_KEY_USAGE_ENCRYPT |
                            PSA_KEY_USAGE_DECRYPT);
    psa_set_key_algorithm(&attributes, alg);
    PSA_ASSERT(psa_import_key(&attributes,
                              key_data->x, key_data->len,
                              &key_id));

    TEST_STATUS(expected_status,
                psa_cipher_encrypt_setup(&operation, key_id, alg));
    TEST_STATUS(expected_status,
                psa_cipher_decrypt_setup(&operation, key_id, alg));
    TEST_STATUS(expected_status,
                psa_cipher_encrypt(key_id, alg,
                                   input, sizeof(input),
                                   output, sizeof(output), &length));
    TEST_STATUS(expected_status,
                psa_cipher_decrypt(key_id, alg,
                                   input, sizeof(input),
                                   output, sizeof(output), &length));

exit:
    psa_cipher_abort(&operation);
    psa_destroy_key(key_id);
    psa_reset_key_attributes(&attributes);
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE */
void aead_fail(int key_type_arg, data_t *key_data,
               int alg_arg, int expected_status_arg)
{
    psa_status_t expected_status = expected_status_arg;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_aead_operation_t operation = PSA_AEAD_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
    uint8_t input[16] = "ABCDEFGHIJKLMNO";
    uint8_t output[64] = { 0 };
    size_t length = SIZE_MAX;

    PSA_INIT();

    psa_set_key_type(&attributes, key_type);
    psa_set_key_usage_flags(&attributes,
                            PSA_KEY_USAGE_ENCRYPT |
                            PSA_KEY_USAGE_DECRYPT);
    psa_set_key_algorithm(&attributes, alg);
    PSA_ASSERT(psa_import_key(&attributes,
                              key_data->x, key_data->len,
                              &key_id));

    TEST_STATUS(expected_status,
                psa_aead_encrypt_setup(&operation, key_id, alg));
    TEST_STATUS(expected_status,
                psa_aead_decrypt_setup(&operation, key_id, alg));
    TEST_STATUS(expected_status,
                psa_aead_encrypt(key_id, alg,
                                 input, sizeof(input),
                                 NULL, 0, input, sizeof(input),
                                 output, sizeof(output), &length));
    TEST_STATUS(expected_status,
                psa_aead_decrypt(key_id, alg,
                                 input, sizeof(input),
                                 NULL, 0, input, sizeof(input),
                                 output, sizeof(output), &length));

exit:
    psa_aead_abort(&operation);
    psa_destroy_key(key_id);
    psa_reset_key_attributes(&attributes);
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE */
void sign_fail(int key_type_arg, data_t *key_data,
               int alg_arg, int private_only,
               int expected_status_arg)
{
    psa_status_t expected_status = expected_status_arg;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
    uint8_t input[1] = { 'A' };
    uint8_t output[PSA_SIGNATURE_MAX_SIZE] = { 0 };
    size_t length = SIZE_MAX;

    PSA_INIT();

    psa_set_key_type(&attributes, key_type);
    psa_set_key_usage_flags(&attributes,
                            PSA_KEY_USAGE_SIGN_HASH |
                            PSA_KEY_USAGE_VERIFY_HASH);
    psa_set_key_algorithm(&attributes, alg);
    PSA_ASSERT(psa_import_key(&attributes,
                              key_data->x, key_data->len,
                              &key_id));

    TEST_STATUS(expected_status,
                psa_sign_hash(key_id, alg,
                              input, sizeof(input),
                              output, sizeof(output), &length));
    if (!private_only) {
        /* Determine a plausible signature size to avoid an INVALID_SIGNATURE
         * error based on this. */
        PSA_ASSERT(psa_get_key_attributes(key_id, &attributes));
        size_t key_bits = psa_get_key_bits(&attributes);
        size_t output_length = sizeof(output);
        if (PSA_KEY_TYPE_IS_RSA(key_type)) {
            output_length = PSA_BITS_TO_BYTES(key_bits);
        } else if (PSA_KEY_TYPE_IS_ECC(key_type)) {
            output_length = 2 * PSA_BITS_TO_BYTES(key_bits);
        }
        TEST_ASSERT(output_length <= sizeof(output));
        TEST_STATUS(expected_status,
                    psa_verify_hash(key_id, alg,
                                    input, sizeof(input),
                                    output, output_length));
    }

exit:
    psa_destroy_key(key_id);
    psa_reset_key_attributes(&attributes);
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE */
void asymmetric_encryption_fail(int key_type_arg, data_t *key_data,
                                int alg_arg, int private_only,
                                int expected_status_arg)
{
    psa_status_t expected_status = expected_status_arg;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
    uint8_t plaintext[PSA_ASYMMETRIC_DECRYPT_OUTPUT_MAX_SIZE] = { 0 };
    uint8_t ciphertext[PSA_ASYMMETRIC_ENCRYPT_OUTPUT_MAX_SIZE] = { 0 };
    size_t length = SIZE_MAX;

    PSA_INIT();

    psa_set_key_type(&attributes, key_type);
    psa_set_key_usage_flags(&attributes,
                            PSA_KEY_USAGE_ENCRYPT |
                            PSA_KEY_USAGE_DECRYPT);
    psa_set_key_algorithm(&attributes, alg);
    PSA_ASSERT(psa_import_key(&attributes,
                              key_data->x, key_data->len,
                              &key_id));

    if (!private_only) {
        TEST_STATUS(expected_status,
                    psa_asymmetric_encrypt(key_id, alg,
                                           plaintext, 1,
                                           NULL, 0,
                                           ciphertext, sizeof(ciphertext),
                                           &length));
    }
    TEST_STATUS(expected_status,
                psa_asymmetric_decrypt(key_id, alg,
                                       ciphertext, sizeof(ciphertext),
                                       NULL, 0,
                                       plaintext, sizeof(plaintext),
                                       &length));

exit:
    psa_destroy_key(key_id);
    psa_reset_key_attributes(&attributes);
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE */
void key_derivation_fail(int alg_arg, int expected_status_arg)
{
    psa_status_t expected_status = expected_status_arg;
    psa_algorithm_t alg = alg_arg;
    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;

    PSA_INIT();

    TEST_EQUAL(expected_status,
               psa_key_derivation_setup(&operation, alg));

exit:
    psa_key_derivation_abort(&operation);
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE */
void key_agreement_fail(int key_type_arg, data_t *key_data,
                        int alg_arg, int private_only,
                        int expected_status_arg)
{
    psa_status_t expected_status = expected_status_arg;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
    uint8_t public_key[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE] = { 0 };
    size_t public_key_length = SIZE_MAX;
    uint8_t output[PSA_SIGNATURE_MAX_SIZE] = { 0 };
    size_t length = SIZE_MAX;
    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;

    PSA_INIT();

    psa_set_key_type(&attributes, key_type);
    psa_set_key_usage_flags(&attributes,
                            PSA_KEY_USAGE_DERIVE);
    psa_set_key_algorithm(&attributes, alg);
    PSA_ASSERT(psa_import_key(&attributes,
                              key_data->x, key_data->len,
                              &key_id));
    if (PSA_KEY_TYPE_IS_KEY_PAIR(key_type) ||
        PSA_KEY_TYPE_IS_PUBLIC_KEY(key_type)) {
        PSA_ASSERT(psa_export_public_key(key_id,
                                         public_key, sizeof(public_key),
                                         &public_key_length));
    }

    TEST_STATUS(expected_status,
                psa_raw_key_agreement(alg, key_id,
                                      public_key, public_key_length,
                                      output, sizeof(output), &length));

#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)));
    TEST_STATUS(expected_status,
                psa_key_derivation_key_agreement(
                    &operation,
                    PSA_KEY_DERIVATION_INPUT_SECRET,
                    key_id,
                    public_key, public_key_length));
#endif

    /* There are no public-key operations. */
    (void) private_only;

exit:
    psa_key_derivation_abort(&operation);
    psa_destroy_key(key_id);
    psa_reset_key_attributes(&attributes);
    PSA_DONE();
}
/* END_CASE */