#include <balancer/kernel/http/parser/pico_parser_helpers/request_head_ending_scanner.h>

#include <library/cpp/testing/unittest/registar.h>

using namespace NSrvKernel;

Y_UNIT_TEST_SUITE(TRequestHeadEngingScannerTest) {

    Y_UNIT_TEST(TestCrlfsInDifferentChunks) {
        TString buffer;
        TMessageHeadEndingScanner scanner(&buffer);

        TChunkList body;

        bool requestFinished;
        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
            TChunkList("GET /?test HTTP/1.1\r\nContent-Length"), body).AssignTo(requestFinished));
        UNIT_ASSERT(!requestFinished);
        UNIT_ASSERT(body.Empty());

        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
            TChunkList(": 3\r\nX-Forwarded-For: host\r\n\r"), body).AssignTo(requestFinished));
        UNIT_ASSERT(!requestFinished);
        UNIT_ASSERT(body.Empty());

        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
            TChunkList("\nbody"), body).AssignTo(requestFinished));
        UNIT_ASSERT(requestFinished);

        UNIT_ASSERT_VALUES_EQUAL(StrInplace(body), "body");
        UNIT_ASSERT_VALUES_EQUAL(buffer,
                                 "GET /?test HTTP/1.1\r\nContent-Length: 3\r\nX-Forwarded-For: host\r\n\r\n");
    }

    Y_UNIT_TEST(TestCrlfsInOneChunk) {
        TString buffer;
        TMessageHeadEndingScanner scanner(&buffer);

        TChunkList body;
        bool requestFinished;

        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
            TChunkList("GET /?test HTTP/1.1\r\nContent-Length"), body).AssignTo(requestFinished));
        UNIT_ASSERT(!requestFinished);
        UNIT_ASSERT(body.Empty());

        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
                TChunkList(": 3\nX-Forwarded-For: host\r\n\r\nbody"), body).AssignTo(requestFinished));
        UNIT_ASSERT(requestFinished);

        UNIT_ASSERT_VALUES_EQUAL(StrInplace(body), "body");
        UNIT_ASSERT_VALUES_EQUAL(buffer,
                                 "GET /?test HTTP/1.1\r\nContent-Length: 3\nX-Forwarded-For: host\r\n\r\n");
    }

    Y_UNIT_TEST(TestCrlfsInDifferentChunksOnlyNewlines) {
        TString buffer;
        TMessageHeadEndingScanner scanner(&buffer);

        TChunkList body;

        bool requestFinished;
        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
                TChunkList("GET /?test HTTP/1.1\nContent-Length"), body).AssignTo(requestFinished));
        UNIT_ASSERT(!requestFinished);
        UNIT_ASSERT(body.Empty());

        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
                TChunkList(": 3\nX-Forwarded-For: host\n"), body).AssignTo(requestFinished));
        UNIT_ASSERT(!requestFinished);
        UNIT_ASSERT(body.Empty());

        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
                TChunkList("\nbody"), body).AssignTo(requestFinished));
        UNIT_ASSERT(requestFinished);

        UNIT_ASSERT_VALUES_EQUAL(StrInplace(body), "body");
        UNIT_ASSERT_VALUES_EQUAL(buffer,
                                 "GET /?test HTTP/1.1\nContent-Length: 3\nX-Forwarded-For: host\n\n");
    }

    Y_UNIT_TEST(TestMakeSureToZeroCounters) {
        TString buffer;
        TMessageHeadEndingScanner scanner(&buffer);

        TChunkList body;

        bool requestFinished;
        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
                TChunkList("GET /?test HTTP/1.1\r\nContent-Length: 3"), body).AssignTo(requestFinished));
        UNIT_ASSERT(!requestFinished);
        UNIT_ASSERT(body.Empty());

        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
                TChunkList("\r\n"), body).AssignTo(requestFinished));
        UNIT_ASSERT(!requestFinished);
        UNIT_ASSERT(body.Empty());

        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
                TChunkList("\r\n"), body).AssignTo(requestFinished));
        UNIT_ASSERT(requestFinished);
        UNIT_ASSERT(body.Empty());
    }

    Y_UNIT_TEST(TestScanHeaders) {
        TString buffer;
        TMessageHeadEndingScanner scanner(&buffer, /*scanHeaders =*/ true);

        TChunkList body;

        bool requestFinished;
        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
            TChunkList("\r"), body).AssignTo(requestFinished));
        UNIT_ASSERT(!requestFinished);
        UNIT_ASSERT(body.Empty());

        UNIT_ASSERT_NO_EXCEPTION(scanner.ProcessNewDataChunkAndSeparateBody(
            TChunkList("\naba"), body).AssignTo(requestFinished));
        UNIT_ASSERT(requestFinished);
        UNIT_ASSERT_VALUES_EQUAL(body, "aba");
    }

};
