/* BEGIN_HEADER */
#include <stdint.h>

#include "common.h"

#include "psa/crypto.h"

#include "psa_crypto_core.h"

#include "test/psa_crypto_helpers.h"

/* Helper to fill a buffer with a data pattern. The pattern is not
 * important, it just allows a basic check that the correct thing has
 * been written, in a way that will detect an error in offset. */
static void fill_buffer_pattern(uint8_t *buffer, size_t len)
{
    uint8_t data[] = { 0x12, 0x34, 0x56, 0x78 };

    for (size_t i = 0; i < len; i++) {
        buffer[i] = data[i % sizeof(data)];
    }
}

/* Helper to get 2 buffers that overlap by the specified amount.
 * The parameter ptr_diff may be negative, in which case the start of
 * buf2 is allocated before the start of buf1. */
static int setup_overlapping_buffers(size_t buf1_len, size_t buf2_len,
                                     int ptr_diff,
                                     uint8_t **full_buffer,
                                     uint8_t **buf1, uint8_t **buf2)
{
    size_t total_len;
    int buf1_offset;
    int buf2_offset;

    *full_buffer = NULL;
    *buf1 = NULL;
    *buf2 = NULL;

    if (ptr_diff >= 0) {
        /*
         * |---------- buf1 ----------|
         *  <-- ptr_diff -->|---------- buf2 ----------|
         */
        total_len = ptr_diff + buf2_len;
        if (buf1_len > total_len) {
            total_len = buf1_len;
        }
        buf1_offset = 0;
        buf2_offset = ptr_diff;
    } else {
        /*
         *  <-- (-ptr_diff) -->|---------- buf1 ----------|
         * |---------- buf2 ----------|
         */
        total_len = (-ptr_diff) + buf1_len;
        if (buf2_len > total_len) {
            total_len = buf2_len;
        }
        buf1_offset = -ptr_diff;
        buf2_offset = 0;
    }

    /* Edge case: if the length is zero, allocate a 1-byte buffer to avoid
     * calloc returning NULL. */
    if (total_len == 0) {
        total_len = 1;
    }

    TEST_CALLOC(*full_buffer, total_len);

    *buf1 = *full_buffer + buf1_offset;
    *buf2 = *full_buffer + buf2_offset;

    return 0;

exit:
    return -1;
}
/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_PSA_CRYPTO_C
 * END_DEPENDENCIES
 */

/* BEGIN_CASE */
void copy_input(int src_len, int dst_len, psa_status_t exp_status)
{
    uint8_t *src_buffer = NULL;
    uint8_t *dst_buffer = NULL;
    psa_status_t status;

    TEST_CALLOC_NONNULL(src_buffer, src_len);
    TEST_CALLOC_NONNULL(dst_buffer, dst_len);

    fill_buffer_pattern(src_buffer, src_len);

    status = psa_crypto_copy_input(src_buffer, src_len, dst_buffer, dst_len);
    TEST_EQUAL(status, exp_status);

    if (exp_status == PSA_SUCCESS) {
        /* Note: We compare the first src_len bytes of each buffer, as this is what was copied. */
        TEST_MEMORY_COMPARE(src_buffer, src_len, dst_buffer, src_len);
    }

exit:
    mbedtls_free(src_buffer);
    mbedtls_free(dst_buffer);
}
/* END_CASE */

/* BEGIN_CASE */
void copy_output(int src_len, int dst_len, psa_status_t exp_status)
{
    uint8_t *src_buffer = NULL;
    uint8_t *dst_buffer = NULL;
    psa_status_t status;

    TEST_CALLOC_NONNULL(src_buffer, src_len);
    TEST_CALLOC_NONNULL(dst_buffer, dst_len);

    fill_buffer_pattern(src_buffer, src_len);

    status = psa_crypto_copy_output(src_buffer, src_len, dst_buffer, dst_len);
    TEST_EQUAL(status, exp_status);

    if (exp_status == PSA_SUCCESS) {
        /* Note: We compare the first src_len bytes of each buffer, as this is what was copied. */
        TEST_MEMORY_COMPARE(src_buffer, src_len, dst_buffer, src_len);
    }

exit:
    mbedtls_free(src_buffer);
    mbedtls_free(dst_buffer);
}
/* END_CASE */

/* BEGIN_CASE */
void psa_crypto_alloc_and_copy(int input_null, int input_len,
                               int output_null, int output_len,
                               int exp_ret)
{
    uint8_t *input_buffer = NULL;
    uint8_t *output_buffer = NULL;
    psa_status_t ret;

    psa_crypto_buffer_copy_t buffer_copies;
    memset(&buffer_copies, 0, sizeof(buffer_copies));

    if (!input_null) {
        TEST_CALLOC(input_buffer, input_len);
        fill_buffer_pattern(input_buffer, input_len);
    }
    if (!output_null) {
        TEST_CALLOC(output_buffer, output_len);
        fill_buffer_pattern(output_buffer, output_len);
    }
    ret = psa_crypto_alloc_and_copy(input_buffer, input_len,
                                    output_buffer, output_len,
                                    &buffer_copies);
    TEST_EQUAL(exp_ret, (int) ret);

    if (exp_ret == PSA_SUCCESS) {
        TEST_MEMORY_COMPARE(input_buffer, input_len, buffer_copies.input, buffer_copies.input_len);
        TEST_EQUAL(output_len, buffer_copies.output_len);
    }

exit:
    mbedtls_free(input_buffer);
    mbedtls_free(output_buffer);
    mbedtls_free(buffer_copies.input);
    mbedtls_free(buffer_copies.output);
}
/* END_CASE */

/* BEGIN_CASE */
void psa_crypto_alloc_and_copy_zero_length(int input_zero_length,
                                           int output_zero_length)
{
    uint8_t input_buffer[] = { 0x12 };
    uint8_t output_buffer[] = { 0x34 };

    size_t input_len = input_zero_length ? 0 : 1;
    size_t output_len = output_zero_length ? 0 : 1;

    psa_crypto_buffer_copy_t buffer_copies;
    memset(&buffer_copies, 0, sizeof(buffer_copies));

    psa_status_t ret = psa_crypto_alloc_and_copy(input_buffer, input_len,
                                                 output_buffer, output_len,
                                                 &buffer_copies);
    TEST_EQUAL(ret, PSA_SUCCESS);

    if (input_zero_length) {
        TEST_ASSERT(buffer_copies.input == NULL);
    } else {
        TEST_MEMORY_COMPARE(input_buffer, input_len, buffer_copies.input, buffer_copies.input_len);
    }
    if (output_zero_length) {
        TEST_ASSERT(buffer_copies.output == NULL);
    } else {
        TEST_EQUAL(output_len, buffer_copies.output_len);
    }

exit:
    mbedtls_free(buffer_copies.input);
    mbedtls_free(buffer_copies.output);
}
/* END_CASE */

/* BEGIN_CASE */
/* Ensure that overlapping buffers can be copied correctly. */
void psa_crypto_alloc_and_copy_overlapping(int input_len, int output_len,
                                           int ptr_diff, int exp_ret)
{
    uint8_t *full_buffer = NULL;
    uint8_t *input = NULL;
    uint8_t *output = NULL;

    psa_status_t ret;

    psa_crypto_buffer_copy_t buffer_copies;
    memset(&buffer_copies, 0, sizeof(buffer_copies));

    TEST_EQUAL(setup_overlapping_buffers(input_len, output_len, ptr_diff,
                                         &full_buffer, &input, &output), 0);

    fill_buffer_pattern(input, input_len);

    ret = psa_crypto_alloc_and_copy(input, input_len, output, output_len,
                                    &buffer_copies);

    TEST_EQUAL((int) ret, exp_ret);

    if (exp_ret == PSA_SUCCESS) {
        if (input_len == 0) {
            TEST_ASSERT(buffer_copies.input == NULL);
        } else {
            TEST_MEMORY_COMPARE(input, input_len, buffer_copies.input, buffer_copies.input_len);
        }
        if (output_len == 0) {
            TEST_ASSERT(buffer_copies.output == NULL);
        } else {
            TEST_EQUAL(output_len, buffer_copies.output_len);
        }
    }

exit:
    mbedtls_free(full_buffer);
    mbedtls_free(buffer_copies.input);
    mbedtls_free(buffer_copies.output);
}
/* END_CASE */

/* BEGIN_CASE */
void psa_crypto_copy_and_free(int input_null, int input_len,
                              int output_null, int output_len,
                              int orig_output_null,
                              int exp_ret)
{
    uint8_t *input = NULL;
    uint8_t *output = NULL;
    uint8_t *output_for_comparison = NULL;
    uint8_t *orig_output = NULL;
    size_t calloc_len;
    psa_status_t ret;

    psa_crypto_buffer_copy_t buffer_copies;
    memset(&buffer_copies, 0, sizeof(buffer_copies));

    if (!input_null) {
        /* If zero-length, ensure we actually allocate something
         * rather than getting NULL. */
        calloc_len = input_len == 0 ? 1 : input_len;
        TEST_CALLOC(input, calloc_len);
    }
    if (!output_null) {
        /* If zero-length, ensure we actually allocate something
         * rather than getting NULL. */
        calloc_len = output_len == 0 ? 1 : output_len;
        TEST_CALLOC(output, calloc_len);
        TEST_CALLOC(output_for_comparison, calloc_len);

        fill_buffer_pattern(output, output_len);
        /* We expect the output buffer to be freed, so keep a copy
         * for comparison. */
        memcpy(output_for_comparison, output, output_len);
    }
    if (!orig_output_null) {
        /* If zero-length, ensure we actually allocate something
         * rather than getting NULL. */
        calloc_len = output_len == 0 ? 1 : output_len;
        TEST_CALLOC(orig_output, calloc_len);
    }

    buffer_copies.input = input;
    buffer_copies.input_len = input_len;
    buffer_copies.output = output;
    buffer_copies.output_len = output_len;
    buffer_copies.output_original = orig_output;

    ret = psa_crypto_copy_and_free(&buffer_copies);

    TEST_EQUAL((int) ret, exp_ret);

    if (exp_ret == PSA_SUCCESS) {
        TEST_ASSERT(buffer_copies.input == NULL);
        TEST_ASSERT(buffer_copies.output == NULL);

        if (!output_null) {
            TEST_MEMORY_COMPARE(output_for_comparison, output_len,
                                buffer_copies.output_original, buffer_copies.output_len);
        }
    }

exit:
    mbedtls_free(buffer_copies.input);
    mbedtls_free(buffer_copies.output);
    mbedtls_free(output_for_comparison);
    mbedtls_free(orig_output);
}
/* END_CASE */

/* BEGIN_CASE */
void psa_crypto_buffer_copy_round_trip()
{
    uint8_t input[100];
    uint8_t output[100];
    uint8_t output_for_comparison[100];
    psa_status_t ret;

    psa_crypto_buffer_copy_t buffer_copies;
    memset(&buffer_copies, 0, sizeof(buffer_copies));

    fill_buffer_pattern(input, sizeof(input));

    ret = psa_crypto_alloc_and_copy(input, sizeof(input),
                                    output, sizeof(output),
                                    &buffer_copies);

    TEST_ASSERT(ret == PSA_SUCCESS);
    TEST_MEMORY_COMPARE(input, sizeof(input),
                        buffer_copies.input, buffer_copies.input_len);
    TEST_EQUAL(sizeof(output), buffer_copies.output_len);

    /* Simulate the PSA function filling the (internal) output buffer. */
    fill_buffer_pattern(buffer_copies.output, buffer_copies.output_len);

    /* Make a copy of output to compare the copy-back */
    memcpy(output_for_comparison, buffer_copies.output,
           sizeof(output_for_comparison));

    ret = psa_crypto_copy_and_free(&buffer_copies);

    TEST_EQUAL(ret, PSA_SUCCESS);
    /* Check that the output was copied back correctly. */
    TEST_MEMORY_COMPARE(output_for_comparison, sizeof(output_for_comparison),
                        output, sizeof(output));

exit:
    mbedtls_free(buffer_copies.input);
    mbedtls_free(buffer_copies.output);
}
/* END_CASE */