#include <gtest/gtest.h>

#include "tcp_reassembly.h"

#if 1
TEST(TCP_REASSEMBLY, ORDER)
{
    struct tcp_segment *seg;
    struct tcp_reassembly *queue = tcp_reassembly_new(10, 16);
    EXPECT_TRUE(queue != NULL);

    tcp_reassembly_set_recv_next(queue, 90);

    /*
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     * |0|1|2|3|4|5|6|7|8|9|a|b|c|d|e|f|g|h|i|j|A|B|C|D|E|F|G|H|I|J|
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     *  |                   |                   |
     *  +---> 90            +---> 100           +---> 110
     */

    seg = tcp_segment_new(90, "0123456789", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_segment_new(100, "abcdefghij", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_segment_new(110, "ABCDEFGHIJ", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "0123456789", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "abcdefghij", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "ABCDEFGHIJ", seg->len) == 0);
    tcp_segment_free(seg);

    EXPECT_TRUE(tcp_reassembly_pop(queue) == NULL);

    tcp_reassembly_free(queue);
}
#endif

#if 1
TEST(TCP_REASSEMBLY, OUT_OF_ORDER)
{
    struct tcp_segment *seg;
    struct tcp_reassembly *queue = tcp_reassembly_new(10, 16);
    EXPECT_TRUE(queue != NULL);

    tcp_reassembly_set_recv_next(queue, 90);

    /*
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     * |0|1|2|3|4|5|6|7|8|9|a|b|c|d|e|f|g|h|i|j|A|B|C|D|E|F|G|H|I|J|
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     *  |                   |                   |
     *  +---> 90            +---> 100           +---> 110
     */

    seg = tcp_segment_new(110, "ABCDEFGHIJ", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_segment_new(100, "abcdefghij", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_segment_new(90, "0123456789", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "0123456789", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "abcdefghij", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "ABCDEFGHIJ", seg->len) == 0);
    tcp_segment_free(seg);

    EXPECT_TRUE(tcp_reassembly_pop(queue) == NULL);

    tcp_reassembly_free(queue);
}
#endif

#if 1
TEST(TCP_REASSEMBLY, REPEAT)
{
    struct tcp_segment *seg;
    struct tcp_reassembly *queue = tcp_reassembly_new(10, 16);
    EXPECT_TRUE(queue != NULL);

    tcp_reassembly_set_recv_next(queue, 100);

    /*
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     * |0|1|2|3|4|5|6|7|8|9|a|b|c|d|e|f|g|h|i|j|
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     *  |                  |A|B|C|D|E|F|G|H|I|J|
     *  |                  +-+-+-+-+-+-+-+-+-+-+
     *  |                   |                 |
     *  +---> 90            +---> 100         +---> 109
     */

    seg = tcp_segment_new(90, "0123456789", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_segment_new(100, "abcdefghij", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_segment_new(100, "ABCDEFGHIJ", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == -2); // repeat
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "abcdefghij", seg->len) == 0);
    tcp_segment_free(seg);

    EXPECT_TRUE(tcp_reassembly_pop(queue) == NULL);

    tcp_reassembly_free(queue);
}
#endif

#if 1
TEST(TCP_REASSEMBLY, OVERLAP)
{
    struct tcp_segment *seg;
    struct tcp_reassembly *queue = tcp_reassembly_new(10, 16);
    EXPECT_TRUE(queue != NULL);

    tcp_reassembly_set_recv_next(queue, 90);

    /*
     * +-+-+-+-+-+-+-+-+-+-+
     * |0|1|2|3|4|5|6|7|8|9|
     * +-+-+-+-+-+-+-+-+-+-+
     *  |          +-+-+-+-+-+-+-+-+-+-+
     *  |          |a|b|c|d|e|f|g|h|i|j|
     *  |          +-+-+-+-+-+-+-+-+-+-+
     *  |           |     |    +-+-+-+-+-+-+-+-+-+-+
     *  |           |     |    |A|B|C|D|E|F|G|H|I|J|
     *  |           |     |    +-+-+-+-+-+-+-+-+-+-+
     *  |           |     |     |     |           |
     *  |           |     |     |     |           +---> 111
     *  |           |     |     |     +---> 105
     *  |           |     |     +---> 102
     *  |           |     +---> 99
     *  |           +---> 96
     *  +---> 90
     *
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     * |0|1|2|3|4|5|6|7|8|9|e|f|g|h|i|j|E|F|G|H|I|J|
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     */

    seg = tcp_segment_new(90, "0123456789", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_segment_new(96, "abcdefghij", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 1);

    seg = tcp_segment_new(102, "ABCDEFGHIJ", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 1);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "0123456789", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 6);
    EXPECT_TRUE(memcmp(seg->data, "efghij", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 6);
    EXPECT_TRUE(memcmp(seg->data, "EFGHIJ", seg->len) == 0);
    tcp_segment_free(seg);

    EXPECT_TRUE(tcp_reassembly_pop(queue) == NULL);

    tcp_reassembly_free(queue);
}
#endif

#if 1
TEST(TCP_REASSEMBLY, SEQ_WRAPAROUND)
{
    struct tcp_segment *seg;
    struct tcp_reassembly *queue = tcp_reassembly_new(10, 16);
    EXPECT_TRUE(queue != NULL);

    // UINT32_MAX = 4294967295
    printf("UINT32_MAX = %u\n", UINT32_MAX);
    tcp_reassembly_set_recv_next(queue, UINT32_MAX - 10);

    /*
     * +-+-+-+-+-+-+-+-+-+-+
     * |0|1|2|3|4|5|6|7|8|9|
     * +-+-+-+-+-+-+-+-+-+-+
     *  |          +-+-+-+-+-+-+-+-+-+-+
     *  |          |a|b|c|d|e|f|g|h|i|j|
     *  |          +-+-+-+-+-+-+-+-+-+-+
     *  |           |     |    +-+-+-+-+-+-+-+-+-+-+
     *  |           |     |    |A|B|C|D|E|F|G|H|I|J|
     *  |           |     |    +-+-+-+-+-+-+-+-+-+-+
     *  |           |     |     |     |           |
     *  |           |     |     |     |           +---> 10
     *  |           |     |     |     +---> 4
     *  |           |     |     +---> 1
     *  |           |     +---> UINT32_MAX - 1
     *  |           +---> UINT32_MAX - 4
     *  +---> UINT32_MAX - 10
     *
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     * |0|1|2|3|4|5|6|7|8|9|e|f|g|h|i|j|E|F|G|H|I|J|
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     */

    seg = tcp_segment_new(UINT32_MAX - 10, "0123456789", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_segment_new(UINT32_MAX - 4, "abcdefghij", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 1);

    seg = tcp_segment_new(1, "ABCDEFGHIJ", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "0123456789", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 6);
    EXPECT_TRUE(memcmp(seg->data, "efghij", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 6);
    EXPECT_TRUE(memcmp(seg->data, "EFGHIJ", seg->len) == 0);
    tcp_segment_free(seg);

    EXPECT_TRUE(tcp_reassembly_pop(queue) == NULL);

    tcp_reassembly_free(queue);
}
#endif

#if 1
TEST(TCP_REASSEMBLY, MAX_TIMEOUT)
{
    struct tcp_segment *seg;
    struct tcp_reassembly *queue = tcp_reassembly_new(10, 16);
    EXPECT_TRUE(queue != NULL);

    tcp_reassembly_set_recv_next(queue, 90);

    /*
     * +-+-+-+-+-+-+-+-+-+-+
     * |0|1|2|3|4|5|6|7|8|9|
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     *  |        |A|B|C|D|E|F|G|H|I|J|
     *  |        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     *  |         |        |a|b|c|d|e|f|g|h|i|j|
     *  |         |        +-+-+-+-+-+-+-+-+-+-+
     *  |         |         |
     *  +---> 90  +-->95    +---> 100
     */

    seg = tcp_segment_new(90, "0123456789", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 1) == 0);

    seg = tcp_segment_new(95, "ABCDEFGHIJ", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 2) == 1);

    seg = tcp_segment_new(100, "abcdefghij", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 3) == 1);

    seg = tcp_reassembly_expire(queue, 11);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "0123456789", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_expire(queue, 11);
    EXPECT_TRUE(seg == NULL);

    tcp_reassembly_inc_recv_next(queue, 10);
    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 5);
    EXPECT_TRUE(memcmp(seg->data, "FGHIJ", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 5);
    EXPECT_TRUE(memcmp(seg->data, "fghij", seg->len) == 0);
    tcp_segment_free(seg);

    EXPECT_TRUE(tcp_reassembly_pop(queue) == NULL);

    tcp_reassembly_free(queue);
}
#endif

#if 1
TEST(TCP_REASSEMBLY, MAX_PACKETS)
{
    struct tcp_segment *seg;
    struct tcp_reassembly *queue = tcp_reassembly_new(10, 2);
    EXPECT_TRUE(queue != NULL);

    tcp_reassembly_set_recv_next(queue, 90);

    /*
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     * |0|1|2|3|4|5|6|7|8|9|a|b|c|d|e|f|g|h|i|j|
     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
     *  |                  |A|B|C|D|E|F|G|H|I|J|
     *  |                  +-+-+-+-+-+-+-+-+-+-+
     *  |                   |                 |
     *  +---> 90            +---> 100         +---> 109
     */

    seg = tcp_segment_new(90, "0123456789", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 0) == 0);

    seg = tcp_segment_new(100, "abcdefghij", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 1) == 0);

    seg = tcp_segment_new(100, "ABCDEFGHIJ", 10);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(tcp_reassembly_push(queue, seg, 2) == -1);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "0123456789", seg->len) == 0);
    tcp_segment_free(seg);

    seg = tcp_reassembly_pop(queue);
    EXPECT_TRUE(seg != NULL);
    EXPECT_TRUE(seg->len == 10);
    EXPECT_TRUE(memcmp(seg->data, "abcdefghij", seg->len) == 0);
    tcp_segment_free(seg);

    EXPECT_TRUE(tcp_reassembly_pop(queue) == NULL);

    tcp_reassembly_free(queue);
}
#endif

int main(int argc, char **argv)
{
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}