/* BEGIN_HEADER */

/* Test random generation as a whole. */

#include "mbedtls/bignum.h"
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/ecdsa.h"
#include "mbedtls/entropy.h"
#include "mbedtls/hmac_drbg.h"
#include "mbedtls/psa_util.h"
#include "psa/crypto.h"

/* How many bytes to generate in each test case for repeated generation.
 * This must be high enough that the probability of generating the same
 * output twice is infinitesimal, but low enough that random generators
 * are willing to deliver that much. */
#define OUTPUT_SIZE 32

/* END_HEADER */

/* BEGIN_CASE depends_on:MBEDTLS_ENTROPY_C:!MBEDTLS_PSA_INJECT_ENTROPY:MBEDTLS_CTR_DRBG_C */
void random_twice_with_ctr_drbg()
{
    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_context drbg;
    unsigned char output1[OUTPUT_SIZE];
    unsigned char output2[OUTPUT_SIZE];

#if defined(MBEDTLS_AES_C)
    MD_PSA_INIT();
#else
    USE_PSA_INIT();
#endif


    /* First round */
    mbedtls_entropy_init(&entropy);
    mbedtls_ctr_drbg_init(&drbg);
    TEST_EQUAL(0, mbedtls_ctr_drbg_seed(&drbg,
                                        mbedtls_entropy_func, &entropy,
                                        NULL, 0));
    TEST_EQUAL(0, mbedtls_ctr_drbg_random(&drbg,
                                          output1, sizeof(output1)));
    mbedtls_ctr_drbg_free(&drbg);
    mbedtls_entropy_free(&entropy);

    /* Second round */
    mbedtls_entropy_init(&entropy);
    mbedtls_ctr_drbg_init(&drbg);
    TEST_EQUAL(0, mbedtls_ctr_drbg_seed(&drbg,
                                        mbedtls_entropy_func, &entropy,
                                        NULL, 0));
    TEST_EQUAL(0, mbedtls_ctr_drbg_random(&drbg,
                                          output2, sizeof(output2)));
    mbedtls_ctr_drbg_free(&drbg);
    mbedtls_entropy_free(&entropy);

    /* The two rounds must generate different random data. */
    TEST_ASSERT(memcmp(output1, output2, OUTPUT_SIZE) != 0);

exit:
    mbedtls_ctr_drbg_free(&drbg);
    mbedtls_entropy_free(&entropy);
#if defined(MBEDTLS_AES_C)
    MD_PSA_DONE();
#else
    USE_PSA_DONE();
#endif
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_ENTROPY_C:!MBEDTLS_PSA_INJECT_ENTROPY:MBEDTLS_HMAC_DRBG_C */
void random_twice_with_hmac_drbg(int md_type)
{
    mbedtls_entropy_context entropy;
    mbedtls_hmac_drbg_context drbg;
    unsigned char output1[OUTPUT_SIZE];
    unsigned char output2[OUTPUT_SIZE];
    const mbedtls_md_info_t *md_info = mbedtls_md_info_from_type(md_type);

    MD_PSA_INIT();

    /* First round */
    mbedtls_entropy_init(&entropy);
    mbedtls_hmac_drbg_init(&drbg);
    TEST_EQUAL(0, mbedtls_hmac_drbg_seed(&drbg, md_info,
                                         mbedtls_entropy_func, &entropy,
                                         NULL, 0));
    TEST_EQUAL(0, mbedtls_hmac_drbg_random(&drbg,
                                           output1, sizeof(output1)));
    mbedtls_hmac_drbg_free(&drbg);
    mbedtls_entropy_free(&entropy);

    /* Second round */
    mbedtls_entropy_init(&entropy);
    mbedtls_hmac_drbg_init(&drbg);
    TEST_EQUAL(0, mbedtls_hmac_drbg_seed(&drbg, md_info,
                                         mbedtls_entropy_func, &entropy,
                                         NULL, 0));
    TEST_EQUAL(0, mbedtls_hmac_drbg_random(&drbg,
                                           output2, sizeof(output2)));
    mbedtls_hmac_drbg_free(&drbg);
    mbedtls_entropy_free(&entropy);

    /* The two rounds must generate different random data. */
    TEST_ASSERT(memcmp(output1, output2, OUTPUT_SIZE) != 0);

exit:
    mbedtls_hmac_drbg_free(&drbg);
    mbedtls_entropy_free(&entropy);
    MD_PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_C:!MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
void random_twice_with_psa_from_classic()
{
    unsigned char output1[OUTPUT_SIZE];
    unsigned char output2[OUTPUT_SIZE];

    /* First round */
    PSA_ASSERT(psa_crypto_init());
    TEST_EQUAL(0, mbedtls_psa_get_random(MBEDTLS_PSA_RANDOM_STATE,
                                         output1, sizeof(output1)));
    PSA_DONE();

    /* Second round */
    PSA_ASSERT(psa_crypto_init());
    TEST_EQUAL(0, mbedtls_psa_get_random(MBEDTLS_PSA_RANDOM_STATE,
                                         output2, sizeof(output2)));
    PSA_DONE();

    /* The two rounds must generate different random data. */
    TEST_ASSERT(memcmp(output1, output2, OUTPUT_SIZE) != 0);

exit:
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_C:!MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
void random_twice_with_psa_from_psa()
{
    unsigned char output1[OUTPUT_SIZE];
    unsigned char output2[OUTPUT_SIZE];

    /* First round */
    PSA_ASSERT(psa_crypto_init());
    PSA_ASSERT(psa_generate_random(output1, sizeof(output1)));
    PSA_DONE();

    /* Second round */
    PSA_ASSERT(psa_crypto_init());
    PSA_ASSERT(psa_generate_random(output2, sizeof(output2)));
    PSA_DONE();

    /* The two rounds must generate different random data. */
    TEST_ASSERT(memcmp(output1, output2, OUTPUT_SIZE) != 0);

exit:
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_C */
void mbedtls_psa_get_random_no_init()
{
    unsigned char output[1];

    TEST_ASSERT(mbedtls_psa_get_random(MBEDTLS_PSA_RANDOM_STATE,
                                       output, sizeof(output)) != 0);
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_C */
void mbedtls_psa_get_random_length(int n)
{
    unsigned char *output = NULL;

    PSA_ASSERT(psa_crypto_init());
    TEST_CALLOC(output, n);

    TEST_EQUAL(0, mbedtls_psa_get_random(MBEDTLS_PSA_RANDOM_STATE,
                                         output, n));
exit:
    mbedtls_free(output);
    PSA_DONE();
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_C:MBEDTLS_ECDSA_C */
void mbedtls_psa_get_random_ecdsa_sign(int curve)
{
    mbedtls_ecp_group grp;
    mbedtls_mpi d, r, s;
    unsigned char buf[] = "This is not a hash.";

    mbedtls_ecp_group_init(&grp);
    mbedtls_mpi_init(&d);
    mbedtls_mpi_init(&r);
    mbedtls_mpi_init(&s);

    TEST_EQUAL(0, mbedtls_mpi_lset(&d, 123456789));
    TEST_EQUAL(0, mbedtls_ecp_group_load(&grp, curve));
    PSA_ASSERT(psa_crypto_init());
    TEST_EQUAL(0, mbedtls_ecdsa_sign(&grp, &r, &s, &d,
                                     buf, sizeof(buf),
                                     mbedtls_psa_get_random,
                                     MBEDTLS_PSA_RANDOM_STATE));
exit:
    mbedtls_mpi_free(&d);
    mbedtls_mpi_free(&r);
    mbedtls_mpi_free(&s);
    mbedtls_ecp_group_free(&grp);
    PSA_DONE();
}
/* END_CASE */