/* BEGIN_HEADER */
#include "mbedtls/chachapoly.h"
/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_CHACHAPOLY_C
 * END_DEPENDENCIES
 */

/* BEGIN_CASE */
void mbedtls_chachapoly_enc(data_t *key_str,
                            data_t *nonce_str,
                            data_t *aad_str,
                            data_t *input_str,
                            data_t *output_str,
                            data_t *mac_str)
{
    unsigned char output[265];
    unsigned char mac[16]; /* size set by the standard */
    mbedtls_chachapoly_context ctx;

    TEST_ASSERT(key_str->len   == 32);
    TEST_ASSERT(nonce_str->len == 12);
    TEST_ASSERT(mac_str->len   == 16);

    mbedtls_chachapoly_init(&ctx);

    TEST_ASSERT(mbedtls_chachapoly_setkey(&ctx, key_str->x) == 0);

    TEST_ASSERT(mbedtls_chachapoly_encrypt_and_tag(&ctx,
                                                   input_str->len, nonce_str->x,
                                                   aad_str->x, aad_str->len,
                                                   input_str->x, output, mac) == 0);

    TEST_ASSERT(memcmp(output_str->x, output, output_str->len) == 0);
    TEST_ASSERT(memcmp(mac_str->x, mac, 16U) == 0);

exit:
    mbedtls_chachapoly_free(&ctx);
}
/* END_CASE */

/* BEGIN_CASE */
void mbedtls_chachapoly_dec(data_t *key_str,
                            data_t *nonce_str,
                            data_t *aad_str,
                            data_t *input_str,
                            data_t *output_str,
                            data_t *mac_str,
                            int ret_exp)
{
    unsigned char output[265];
    int ret;
    mbedtls_chachapoly_context ctx;

    TEST_ASSERT(key_str->len   == 32);
    TEST_ASSERT(nonce_str->len == 12);
    TEST_ASSERT(mac_str->len   == 16);

    mbedtls_chachapoly_init(&ctx);

    TEST_ASSERT(mbedtls_chachapoly_setkey(&ctx, key_str->x) == 0);

    ret = mbedtls_chachapoly_auth_decrypt(&ctx,
                                          input_str->len, nonce_str->x,
                                          aad_str->x, aad_str->len,
                                          mac_str->x, input_str->x, output);

    TEST_ASSERT(ret == ret_exp);
    if (ret_exp == 0) {
        TEST_ASSERT(memcmp(output_str->x, output, output_str->len) == 0);
    }

exit:
    mbedtls_chachapoly_free(&ctx);
}
/* END_CASE */

/* BEGIN_CASE */
void chachapoly_update_too_long()
{
    unsigned char key[32] = { 0 };
    unsigned char nonce[12] = { 0 };
    unsigned char input[65] = { 0 };
    unsigned char output[65] = { 0 };
    mbedtls_chachapoly_context ctx;
    size_t i;

    mbedtls_chachapoly_init(&ctx);

    TEST_EQUAL(mbedtls_chachapoly_setkey(&ctx, key), 0);
    TEST_EQUAL(mbedtls_chachapoly_starts(&ctx, nonce,
                                         MBEDTLS_CHACHAPOLY_ENCRYPT), 0);

    /* Force the next update to be too long without processing a huge buffer. */
    ctx.chacha20_ctx.state[12] = UINT32_MAX;
    /* Test that the failed update does not modify the output. This is a
     * defense in depth assertion, however is not currently gauranteed by
     * the API */
    memset(output, 0x2a, sizeof(output));
    TEST_EQUAL(mbedtls_chachapoly_update(&ctx, sizeof(input), input, output),
               MBEDTLS_ERR_CHACHA20_BAD_INPUT_DATA);
    for (i = 0; i < sizeof(output); i++) {
        TEST_EQUAL(output[i], 0x2a);
    }

exit:
    mbedtls_chachapoly_free(&ctx);
}
/* END_CASE */

/* BEGIN_CASE */
void chachapoly_state()
{
    unsigned char key[32];
    unsigned char nonce[12];
    unsigned char aad[1];
    unsigned char input[1];
    unsigned char output[1];
    unsigned char mac[16];
    size_t input_len = sizeof(input);
    size_t aad_len = sizeof(aad);
    mbedtls_chachapoly_context ctx;

    memset(key,    0x00, sizeof(key));
    memset(nonce,  0x00, sizeof(nonce));
    memset(aad,    0x00, sizeof(aad));
    memset(input,  0x00, sizeof(input));
    memset(output, 0x00, sizeof(output));
    memset(mac,    0x00, sizeof(mac));

    /* Initial state: finish, update, update_aad forbidden */
    mbedtls_chachapoly_init(&ctx);

    TEST_ASSERT(mbedtls_chachapoly_finish(&ctx, mac)
                == MBEDTLS_ERR_CHACHAPOLY_BAD_STATE);
    TEST_ASSERT(mbedtls_chachapoly_update(&ctx, input_len, input, output)
                == MBEDTLS_ERR_CHACHAPOLY_BAD_STATE);
    TEST_ASSERT(mbedtls_chachapoly_update_aad(&ctx, aad, aad_len)
                == MBEDTLS_ERR_CHACHAPOLY_BAD_STATE);

    /* Still initial state: finish, update, update_aad forbidden */
    TEST_ASSERT(mbedtls_chachapoly_setkey(&ctx, key)
                == 0);

    TEST_ASSERT(mbedtls_chachapoly_finish(&ctx, mac)
                == MBEDTLS_ERR_CHACHAPOLY_BAD_STATE);
    TEST_ASSERT(mbedtls_chachapoly_update(&ctx, input_len, input, output)
                == MBEDTLS_ERR_CHACHAPOLY_BAD_STATE);
    TEST_ASSERT(mbedtls_chachapoly_update_aad(&ctx, aad, aad_len)
                == MBEDTLS_ERR_CHACHAPOLY_BAD_STATE);

    /* Starts -> finish OK */
    TEST_ASSERT(mbedtls_chachapoly_starts(&ctx, nonce, MBEDTLS_CHACHAPOLY_ENCRYPT)
                == 0);
    TEST_ASSERT(mbedtls_chachapoly_finish(&ctx, mac)
                == 0);

    /* After finish: update, update_aad forbidden */
    TEST_ASSERT(mbedtls_chachapoly_update(&ctx, input_len, input, output)
                == MBEDTLS_ERR_CHACHAPOLY_BAD_STATE);
    TEST_ASSERT(mbedtls_chachapoly_update_aad(&ctx, aad, aad_len)
                == MBEDTLS_ERR_CHACHAPOLY_BAD_STATE);

    /* Starts -> update* OK */
    TEST_ASSERT(mbedtls_chachapoly_starts(&ctx, nonce, MBEDTLS_CHACHAPOLY_ENCRYPT)
                == 0);
    TEST_ASSERT(mbedtls_chachapoly_update(&ctx, input_len, input, output)
                == 0);
    TEST_ASSERT(mbedtls_chachapoly_update(&ctx, input_len, input, output)
                == 0);

    /* After update: update_aad forbidden */
    TEST_ASSERT(mbedtls_chachapoly_update_aad(&ctx, aad, aad_len)
                == MBEDTLS_ERR_CHACHAPOLY_BAD_STATE);

    /* Starts -> update_aad* -> finish OK */
    TEST_ASSERT(mbedtls_chachapoly_starts(&ctx, nonce, MBEDTLS_CHACHAPOLY_ENCRYPT)
                == 0);
    TEST_ASSERT(mbedtls_chachapoly_update_aad(&ctx, aad, aad_len)
                == 0);
    TEST_ASSERT(mbedtls_chachapoly_update_aad(&ctx, aad, aad_len)
                == 0);
    TEST_ASSERT(mbedtls_chachapoly_finish(&ctx, mac)
                == 0);

exit:
    mbedtls_chachapoly_free(&ctx);
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_SELF_TEST */
void chachapoly_selftest()
{
    TEST_ASSERT(mbedtls_chachapoly_self_test(1) == 0);
}
/* END_CASE */