/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file 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 "stuffer/s2n_stuffer.h"

#include "api/s2n.h"
#include "s2n_test.h"
#include "utils/s2n_mem.h"

int main(int argc, char **argv)
{
    uint8_t entropy[2048] = { 0 };
    struct s2n_stuffer stuffer = { 0 };
    uint8_t u8 = 0;
    uint16_t u16 = 0;
    uint32_t u32 = 0;
    uint64_t u64 = 0;

    BEGIN_TEST();
    EXPECT_SUCCESS(s2n_disable_tls13_in_test());

    /* Create a 100 byte stuffer */
    EXPECT_SUCCESS(s2n_stuffer_alloc(&stuffer, 100));

    /* Try to write 101 bytes */
    struct s2n_blob in = { 0 };
    EXPECT_SUCCESS(s2n_blob_init(&in, entropy, 101));
    EXPECT_FAILURE(s2n_stuffer_write(&stuffer, &in));

    /* Try to write 101 1-byte ints bytes */
    for (uint64_t i = 0; i < 100; i++) {
        uint64_t value = i * (0xff / 100);
        EXPECT_SUCCESS(s2n_stuffer_write_uint8(&stuffer, value));
    }
    EXPECT_FAILURE(s2n_stuffer_write_uint8(&stuffer, 1));

    struct s2n_blob copy_of_bytes = { 0 };
    EXPECT_SUCCESS(s2n_stuffer_extract_blob(&stuffer, &copy_of_bytes));

    /* Read those back, and expect the same results */
    for (uint64_t i = 0; i < 100; i++) {
        uint64_t value = i * (0xff / 100);
        EXPECT_SUCCESS(s2n_stuffer_read_uint8(&stuffer, &u8));
        EXPECT_EQUAL(value, u8);
        EXPECT_EQUAL(copy_of_bytes.data[i], u8);
    }

    /* The copy_of_bytes should have the same values */
    for (uint64_t i = 0; i < 100; i++) {
        uint64_t value = i * (0xff / 100);
        EXPECT_EQUAL(copy_of_bytes.data[i], value);
    }

    EXPECT_FAILURE(s2n_stuffer_read_uint8(&stuffer, &u8));

    /* Try to write 51 2-byte ints bytes */
    EXPECT_SUCCESS(s2n_stuffer_wipe(&stuffer));
    for (uint64_t i = 0; i < 50; i++) {
        uint64_t value = i * (0xffff / 50);
        EXPECT_SUCCESS(s2n_stuffer_write_uint16(&stuffer, value));
    }
    EXPECT_FAILURE(s2n_stuffer_write_uint16(&stuffer, 1));

    /* Read those back, and expect the same results */
    for (uint64_t i = 0; i < 50; i++) {
        uint64_t value = i * (0xffff / 50);
        EXPECT_SUCCESS(s2n_stuffer_read_uint16(&stuffer, &u16));
        EXPECT_EQUAL(value, u16);
    }
    EXPECT_FAILURE(s2n_stuffer_read_uint16(&stuffer, &u16));

    /* Try to write 34 3-byte ints bytes */
    EXPECT_SUCCESS(s2n_stuffer_wipe(&stuffer));
    for (uint64_t i = 0; i < 33; i++) {
        uint64_t value = i * (0xffffff / 33);
        EXPECT_SUCCESS(s2n_stuffer_write_uint24(&stuffer, value));
    }
    EXPECT_FAILURE(s2n_stuffer_write_uint24(&stuffer, 1));

    /* Read those back, and expect the same results */
    for (uint64_t i = 0; i < 33; i++) {
        uint64_t value = i * (0xffffff / 33);
        EXPECT_SUCCESS(s2n_stuffer_read_uint24(&stuffer, &u32));
        EXPECT_EQUAL(value, u32);
    }
    EXPECT_FAILURE(s2n_stuffer_read_uint16(&stuffer, &u16));

    /* Try to write 26 4-byte ints bytes */
    EXPECT_SUCCESS(s2n_stuffer_wipe(&stuffer));
    for (uint64_t i = 0; i < 25; i++) {
        uint64_t value = i * (0xffffffff / 25);
        EXPECT_SUCCESS(s2n_stuffer_write_uint32(&stuffer, value));
    }
    EXPECT_FAILURE(s2n_stuffer_write_uint32(&stuffer, 1));

    /* Read those back, and expect the same results */
    for (uint64_t i = 0; i < 25; i++) {
        uint64_t value = i * (0xffffffff / 25);
        EXPECT_SUCCESS(s2n_stuffer_read_uint32(&stuffer, &u32));
        EXPECT_EQUAL(value, u32);
    }
    EXPECT_FAILURE(s2n_stuffer_read_uint32(&stuffer, &u32));

    /* Try to write 13 8-byte ints bytes */
    EXPECT_SUCCESS(s2n_stuffer_wipe(&stuffer));
    for (int i = 0; i < 12; i++) {
        uint64_t value = i * (0xffffffffffffffff / 12);
        EXPECT_SUCCESS(s2n_stuffer_write_uint64(&stuffer, value));
    }
    EXPECT_FAILURE(s2n_stuffer_write_uint64(&stuffer, 1));

    /* Read those back, and expect the same results */
    for (int i = 0; i < 12; i++) {
        uint64_t value = i * (0xffffffffffffffff / 12);
        EXPECT_SUCCESS(s2n_stuffer_read_uint64(&stuffer, &u64));
        EXPECT_EQUAL(value, u64);
    }
    EXPECT_FAILURE(s2n_stuffer_read_uint64(&stuffer, &u64));

    EXPECT_SUCCESS(s2n_stuffer_free(&stuffer));

    /* Can still read the copy_of_bytes even once the stuffer has been overwritten and freed */
    for (uint64_t i = 0; i < 100; i++) {
        uint64_t value = i * (0xff / 100);
        EXPECT_EQUAL(copy_of_bytes.data[i], value);
    }
    EXPECT_SUCCESS(s2n_free(&copy_of_bytes));

#ifndef NDEBUG
    /* Invalid blob should fail init */
    struct s2n_stuffer s1 = { 0 };
    struct s2n_blob b1 = { .data = 0, .size = 101 };
    EXPECT_FAILURE(s2n_stuffer_init(&s1, &b1));
#endif

    /* Valid empty blob should succeed init */
    struct s2n_stuffer s2 = { 0 };
    struct s2n_blob b2 = { 0 };
    EXPECT_SUCCESS(s2n_blob_init(&b2, 0, 0));
    EXPECT_SUCCESS(s2n_stuffer_init(&s2, &b2));

    /* Valid blob should succeed init */
    struct s2n_stuffer s3 = { 0 };
    uint8_t a3[12];
    struct s2n_blob b3 = { 0 };
    EXPECT_SUCCESS(s2n_blob_init(&b3, a3, sizeof(a3)));
    EXPECT_SUCCESS(s2n_stuffer_init(&s3, &b3));

    /* Null blob should fail init */
    struct s2n_stuffer s4 = { 0 };
    EXPECT_FAILURE(s2n_stuffer_init(&s4, NULL));

    /* Null stuffer should fail init */
    struct s2n_blob b5 = { 0 };
    EXPECT_SUCCESS(s2n_blob_init(&b5, 0, 0));
    EXPECT_FAILURE(s2n_stuffer_init(NULL, &b5));

    /* Check s2n_stuffer_validate() function */
    EXPECT_ERROR(s2n_stuffer_validate(NULL));
    uint8_t valid_blob_array[12];
    struct s2n_blob blob_valid = { 0 };
    EXPECT_SUCCESS(s2n_blob_init(&blob_valid, valid_blob_array, sizeof(valid_blob_array)));

    struct s2n_stuffer stuffer_valid = { 0 };
    EXPECT_SUCCESS(s2n_stuffer_init(&stuffer_valid, &blob_valid));
    EXPECT_OK(s2n_stuffer_validate(&stuffer));

#ifndef NDEBUG
    struct s2n_blob blob_invalid = { .data = 0, .size = sizeof(valid_blob_array) };
    struct s2n_stuffer stuffer_invalid1 = { .blob = blob_invalid };
    EXPECT_ERROR(s2n_stuffer_validate(&stuffer_invalid1));

    struct s2n_stuffer stuffer_invalid2 = { .blob = blob_valid, .write_cursor = 13 };
    EXPECT_ERROR(s2n_stuffer_validate(&stuffer_invalid2));

    struct s2n_stuffer stuffer_invalid3 = { .blob = blob_valid, .read_cursor = 13 };
    EXPECT_ERROR(s2n_stuffer_validate(&stuffer_invalid3));

    struct s2n_stuffer stuffer_invalid4 = { .blob = blob_valid, .read_cursor = 12, .write_cursor = 1 };
    EXPECT_ERROR(s2n_stuffer_validate(&stuffer_invalid4));
#endif

    struct s2n_stuffer reserve_test_stuffer = { 0 };
    EXPECT_SUCCESS(s2n_stuffer_alloc(&reserve_test_stuffer, 1024));
    EXPECT_EQUAL(s2n_stuffer_space_remaining(&reserve_test_stuffer), 1024);
    EXPECT_EQUAL(s2n_stuffer_data_available(&reserve_test_stuffer), 0);
    EXPECT_FAILURE(s2n_stuffer_reserve_space(&reserve_test_stuffer, 2048));
    EXPECT_EQUAL(s2n_stuffer_space_remaining(&reserve_test_stuffer), 1024);
    EXPECT_EQUAL(s2n_stuffer_data_available(&reserve_test_stuffer), 0);
    EXPECT_SUCCESS(s2n_stuffer_free(&reserve_test_stuffer));

    EXPECT_SUCCESS(s2n_stuffer_growable_alloc(&reserve_test_stuffer, 1024));
    EXPECT_EQUAL(s2n_stuffer_space_remaining(&reserve_test_stuffer), 1024);
    EXPECT_EQUAL(s2n_stuffer_data_available(&reserve_test_stuffer), 0);
    EXPECT_SUCCESS(s2n_stuffer_reserve_space(&reserve_test_stuffer, 2048));
    EXPECT_EQUAL(s2n_stuffer_space_remaining(&reserve_test_stuffer), 2048);
    EXPECT_EQUAL(s2n_stuffer_data_available(&reserve_test_stuffer), 0);
    EXPECT_SUCCESS(s2n_stuffer_free(&reserve_test_stuffer));

    /* Test: s2n_stuffer_init_written */
    {
        uint8_t data[] = "hello world";
        uint8_t input[sizeof(data)] = { 0 };
        uint8_t output[sizeof(data)] = { 0 };

        struct s2n_blob blob = { 0 };
        EXPECT_SUCCESS(s2n_blob_init(&blob, input, sizeof(input)));

        /* Repeat control to show behavior is consistent */
        for (size_t i = 0; i < 3; i++) {
            struct s2n_stuffer unwritten_stuffer = { 0 };
            EXPECT_SUCCESS(s2n_stuffer_init(&unwritten_stuffer, &blob));
            EXPECT_EQUAL(s2n_stuffer_data_available(&unwritten_stuffer), 0);
            EXPECT_FAILURE_WITH_ERRNO(s2n_stuffer_read_bytes(&unwritten_stuffer, output, sizeof(output)),
                    S2N_ERR_STUFFER_OUT_OF_DATA);
            EXPECT_SUCCESS(s2n_stuffer_write_bytes(&unwritten_stuffer, data, sizeof(data)));
        }

        /* Repeat test to show behavior is consistent */
        for (size_t i = 0; i < 3; i++) {
            struct s2n_stuffer written_stuffer = { 0 };
            EXPECT_SUCCESS(s2n_stuffer_init_written(&written_stuffer, &blob));
            EXPECT_EQUAL(s2n_stuffer_data_available(&written_stuffer), sizeof(data));
            EXPECT_SUCCESS(s2n_stuffer_read_bytes(&written_stuffer, output, sizeof(output)));
            EXPECT_BYTEARRAY_EQUAL(data, output, sizeof(output));
            EXPECT_FAILURE_WITH_ERRNO(s2n_stuffer_write_bytes(&written_stuffer, data, sizeof(data)),
                    S2N_ERR_STUFFER_IS_FULL);
        }
    };

    /* Test s2n_stuffer_shift */
    {
        /* Safety */
        EXPECT_FAILURE_WITH_ERRNO(s2n_stuffer_shift(NULL), S2N_ERR_NULL);

        const uint8_t test_data[] = "hello world";
        const uint32_t test_data_size = sizeof(test_data);
        const uint32_t test_offset = 10;

        /* Uninitialized stuffer: no shift */
        {
            struct s2n_stuffer test = { 0 };

            EXPECT_SUCCESS(s2n_stuffer_shift(&test));
            EXPECT_EQUAL(test.read_cursor, 0);
            EXPECT_EQUAL(test.write_cursor, 0);
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), 0);
        }

        /* No data available: no shift */
        {
            uint8_t data[100] = { 0 };
            struct s2n_stuffer test = { 0 };
            EXPECT_SUCCESS(s2n_blob_init(&test.blob, data, sizeof(data)));

            EXPECT_SUCCESS(s2n_stuffer_write_bytes(&test, test_data, sizeof(test_data)));
            EXPECT_SUCCESS(s2n_stuffer_skip_read(&test, sizeof(test_data)));
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), 0);

            EXPECT_SUCCESS(s2n_stuffer_shift(&test));
            EXPECT_EQUAL(test.read_cursor, 0);
            EXPECT_EQUAL(test.write_cursor, 0);
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), 0);
            EXPECT_BYTEARRAY_EQUAL(data, test_data, sizeof(test_data));
        }

        /* Data not offset: no shift */
        {
            uint8_t data[100] = { 0 };
            struct s2n_stuffer test = { 0 };
            EXPECT_SUCCESS(s2n_blob_init(&test.blob, data, sizeof(data)));

            EXPECT_SUCCESS(s2n_stuffer_write_bytes(&test, test_data, test_data_size));
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), test_data_size);

            EXPECT_SUCCESS(s2n_stuffer_shift(&test));
            EXPECT_EQUAL(test.read_cursor, 0);
            EXPECT_EQUAL(test.write_cursor, test_data_size);
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), test_data_size);
            EXPECT_BYTEARRAY_EQUAL(data, test_data, test_data_size);
        }

        /* Data at offset: shifted */
        {
            uint8_t data[100] = { 0 };
            struct s2n_stuffer test = { 0 };
            EXPECT_SUCCESS(s2n_blob_init(&test.blob, data, sizeof(data)));

            EXPECT_SUCCESS(s2n_stuffer_skip_write(&test, test_offset));
            EXPECT_SUCCESS(s2n_stuffer_skip_read(&test, test_offset));
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), 0);
            EXPECT_SUCCESS(s2n_stuffer_write_bytes(&test, test_data, test_data_size));
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), test_data_size);

            EXPECT_SUCCESS(s2n_stuffer_shift(&test));
            EXPECT_EQUAL(test.read_cursor, 0);
            EXPECT_EQUAL(test.write_cursor, test_data_size);
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), test_data_size);
            EXPECT_BYTEARRAY_EQUAL(data, test_data, test_data_size);
        }

        /* Data overlaps: shifted */
        {
            uint8_t data[100] = { 0 };
            struct s2n_stuffer test = { 0 };
            EXPECT_SUCCESS(s2n_blob_init(&test.blob, data, sizeof(data)));

            /* Allocate data large enough that it will overlap when shifted.
             * Allocate the entire block to distinctive data, not just all one character.
             */
            uint8_t overlap_test_data[sizeof(data) - 1] = { 0 };
            for (size_t i = 0; i < sizeof(overlap_test_data); i++) {
                overlap_test_data[i] = i;
            }
            size_t overlap_test_data_size = sizeof(overlap_test_data);
            EXPECT_TRUE(overlap_test_data_size > sizeof(data) / 2);

            EXPECT_SUCCESS(s2n_stuffer_skip_write(&test, 1));
            EXPECT_SUCCESS(s2n_stuffer_skip_read(&test, 1));
            EXPECT_SUCCESS(s2n_stuffer_write_bytes(&test,
                    overlap_test_data, overlap_test_data_size));
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), overlap_test_data_size);
            EXPECT_EQUAL(s2n_stuffer_space_remaining(&test), 0);

            EXPECT_SUCCESS(s2n_stuffer_shift(&test));
            EXPECT_EQUAL(test.read_cursor, 0);
            EXPECT_EQUAL(test.write_cursor, overlap_test_data_size);
            EXPECT_EQUAL(s2n_stuffer_data_available(&test), overlap_test_data_size);
            EXPECT_BYTEARRAY_EQUAL(data, overlap_test_data, overlap_test_data_size);
        }
    }

    END_TEST();
}
