// Copyright (c) 2019, Xiaomi, Inc.  All rights reserved.
// This source code is licensed under the Apache License Version 2.0, which
// can be found in the LICENSE file in the root directory of this source tree.

#include <gtest/gtest.h>
#include <dsn/utility/crc.h>
#include <dsn/utility/endians.h>
#include <dsn/cpp/serialization_helper/thrift_helper.h>
#include <dsn/cpp/serialization_helper/dsn.layer2_types.h>

#include "runtime/rpc/thrift_message_parser.h"

namespace dsn {

DEFINE_TASK_CODE_RPC(RPC_TEST_THRIFT_MESSAGE_PARSER, TASK_PRIORITY_COMMON, THREAD_POOL_DEFAULT)

class thrift_message_parser_test : public testing::Test
{
public:
    void mock_reader_read_data(message_reader &reader, const std::string &data)
    {
        char *buf = reader.read_buffer_ptr(data.length());
        memcpy(buf, data.c_str(), data.size());
        reader.mark_read(data.length());
    }

    void test_get_message_on_receive_v0_data(message_reader &reader,
                                             apache::thrift::protocol::TMessageType messageType,
                                             bool is_request)
    {
        /// write rpc message
        size_t body_length = 0;
        message_ptr msg =
            message_ex::create_request(RPC_TEST_THRIFT_MESSAGE_PARSER, 1000, 64, 5000000000);
        rpc_write_stream stream(msg);
        binary_writer_transport binary_transport(stream);
        boost::shared_ptr<binary_writer_transport> trans_ptr(&binary_transport,
                                                             [](binary_writer_transport *) {});
        ::apache::thrift::protocol::TBinaryProtocol oprot(trans_ptr);
        body_length += oprot.writeMessageBegin("RPC_TEST_THRIFT_MESSAGE_PARSER", messageType, 999);
        body_length += oprot.writeMessageEnd();
        stream.commit_buffer();

        thrift_message_parser parser;
        std::string data;
        int read_next = 0;
        data = std::string("THFT") + std::string(44 + body_length, '\0'); // header+body_length
        data_output out(&data[4], 44);
        out.write_u32(0);           // hdr_version
        out.write_u32(48);          // hdr_length
        out.write_u32(0);           // hdr_crc32
        out.write_u32(body_length); // body_length
        out.write_u32(0);           // body_crc32
        out.write_u32(1);           // app_id
        out.write_u32(28);          // partition_index
        out.write_u32(1000);        // client_timeout
        out.write_u32(64);          // client_thread_hash
        out.write_u64(5000000000);  // client_partition_hash
        ASSERT_EQ(stream.get_buffer().size(), body_length);
        memcpy(&data[48], stream.get_buffer().data(), stream.get_buffer().size());

        mock_reader_read_data(reader, data);
        msg = parser.get_message_on_receive(&reader, read_next);

        if (is_request) {
            ASSERT_NE(msg, nullptr);
            ASSERT_EQ(msg->hdr_format, NET_HDR_THRIFT);

            ASSERT_EQ(msg->header->body_length, body_length);
            ASSERT_EQ(msg->header->gpid, gpid(1, 28));
            ASSERT_EQ(msg->header->hdr_type, THRIFT_HDR_SIG);
            ASSERT_EQ(msg->header->hdr_length, sizeof(message_header));
            ASSERT_EQ(msg->header->hdr_crc32, CRC_INVALID);
            ASSERT_EQ(msg->header->body_crc32, CRC_INVALID);
            ASSERT_EQ(msg->header->id, 999);

            ASSERT_EQ(msg->header->client.timeout_ms, 1000);
            ASSERT_EQ(msg->header->client.thread_hash, 64);
            ASSERT_EQ(msg->header->client.partition_hash, 5000000000);

            ASSERT_EQ(msg->header->context.u.is_request, true);
            ASSERT_EQ(msg->header->context.u.serialize_format, DSF_THRIFT_BINARY);

            // v0 Thrift network format doesn't support message context.
            ASSERT_EQ(msg->header->context.u.is_backup_request, false);
            ASSERT_EQ(msg->header->context.u.is_forwarded, false);
            ASSERT_EQ(msg->header->context.u.is_forward_supported, false);

            ASSERT_EQ(reader.buffer().size(), 0);

            // must be reset
            ASSERT_EQ(parser._header_version, -1);
            ASSERT_EQ(parser._v1_specific_vars->_meta_parsed, false);
            ASSERT_EQ(parser._v1_specific_vars->_meta_length, 0);
            ASSERT_EQ(parser._v1_specific_vars->_body_length, 0);
        } else {
            ASSERT_EQ(msg, nullptr);
            ASSERT_EQ(read_next, -1);
        }
    }

    void test_get_message_on_receive_v1_data(message_reader &reader,
                                             apache::thrift::protocol::TMessageType messageType,
                                             bool is_request,
                                             bool is_backup_request)
    {
        /// write rpc message
        size_t body_length = 0;
        message_ptr msg =
            message_ex::create_request(RPC_TEST_THRIFT_MESSAGE_PARSER, 1000, 64, 5000000000);
        rpc_write_stream body_stream(msg);
        {
            binary_writer_transport transport(body_stream);
            boost::shared_ptr<binary_writer_transport> trans_ptr(&transport,
                                                                 [](binary_writer_transport *) {});
            ::apache::thrift::protocol::TBinaryProtocol oprot(trans_ptr);
            body_length +=
                oprot.writeMessageBegin("RPC_TEST_THRIFT_MESSAGE_PARSER", messageType, 999);
            body_length += oprot.writeMessageEnd();
            body_stream.commit_buffer();
            ASSERT_EQ(body_stream.get_buffer().size(), body_length);
        }

        // write rpc meta
        size_t meta_length = 0;
        thrift_request_meta_v1 meta;
        meta.__set_is_backup_request(is_backup_request);
        meta.__set_app_id(1);
        meta.__set_partition_index(28);
        meta.__set_client_timeout(1000);
        meta.__set_client_partition_hash(5000000000);

        binary_writer meta_writer(1024);
        ::dsn::binary_writer_transport trans(meta_writer);
        boost::shared_ptr<::dsn::binary_writer_transport> transport(
            &trans, [](::dsn::binary_writer_transport *) {});
        ::apache::thrift::protocol::TBinaryProtocol proto(transport);
        meta.write(&proto);

        meta_length = meta_writer.get_buffer().size();

        thrift_message_parser parser;
        std::string data;
        int read_next = 0;
        data = std::string("THFT") + std::string(12 + meta_length + body_length, '\0');
        data_output out(&data[4], 12);
        out.write_u32(1);
        out.write_u32(meta_length);
        out.write_u32(body_length);

        memcpy(&data[16], meta_writer.get_buffer().data(), meta_writer.get_buffer().size());
        memcpy(&data[16 + meta_length],
               body_stream.get_buffer().data(),
               body_stream.get_buffer().size());

        mock_reader_read_data(reader, data);
        ASSERT_EQ(reader.buffer().size(), data.size());
        ASSERT_EQ(reader.buffer().size(), 16 + meta_length + body_length);

        msg = parser.get_message_on_receive(&reader, read_next);

        if (is_request) {
            ASSERT_NE(msg, nullptr);
            ASSERT_EQ(msg->hdr_format, NET_HDR_THRIFT);

            ASSERT_EQ(msg->header->body_length, body_length);
            ASSERT_EQ(msg->header->gpid, gpid(1, 28));
            ASSERT_EQ(msg->header->hdr_type, THRIFT_HDR_SIG);
            ASSERT_EQ(msg->header->hdr_length, sizeof(message_header));
            ASSERT_EQ(msg->header->hdr_crc32, CRC_INVALID);
            ASSERT_EQ(msg->header->body_crc32, CRC_INVALID);
            ASSERT_EQ(msg->header->id, 999);

            ASSERT_EQ(msg->header->client.timeout_ms, 1000);
            ASSERT_EQ(msg->header->client.thread_hash, 7947);
            ASSERT_EQ(msg->header->client.partition_hash, 5000000000);

            ASSERT_EQ(msg->header->context.u.is_request, true);
            ASSERT_EQ(msg->header->context.u.serialize_format, DSF_THRIFT_BINARY);
            ASSERT_EQ(msg->header->context.u.is_backup_request, is_backup_request);
            ASSERT_EQ(msg->header->context.u.is_forwarded, false);
            ASSERT_EQ(msg->header->context.u.is_forward_supported, false);

            ASSERT_EQ(reader.buffer().size(), 0);

            // must be reset
            ASSERT_EQ(parser._header_version, -1);
            ASSERT_EQ(parser._v1_specific_vars->_meta_parsed, false);
        } else {
            ASSERT_EQ(msg, nullptr);
            ASSERT_EQ(read_next, -1);
        }
    }
};

TEST_F(thrift_message_parser_test, get_message_on_receive_incomplete_second_field)
{
    for (int i = 0; i < 4; i++) {
        thrift_message_parser parser;

        std::string data;
        int read_next = 0;
        message_reader reader(64);
        data = std::string("THFT") + std::string(i, ' ');
        mock_reader_read_data(reader, data);
        ASSERT_EQ(reader._buffer_occupied, 4 + i);
        ASSERT_EQ(reader.buffer().size(), 4 + i);

        message_ex *msg = parser.get_message_on_receive(&reader, read_next);
        ASSERT_EQ(msg, nullptr);
        ASSERT_EQ(read_next, 4 - i);
        ASSERT_EQ(parser._header_version, -1);
        ASSERT_EQ(parser._v1_specific_vars->_meta_parsed, false);
        ASSERT_EQ(parser._v1_specific_vars->_meta_length, 0);

        // not consumed
        ASSERT_EQ(reader._buffer_occupied, data.length());
        ASSERT_EQ(reader.buffer().size(), data.length());
    }
}

TEST_F(thrift_message_parser_test, get_message_on_receive_incomplete_v0_hdr_len)
{
    for (int i = 4; i < 44; i++) {
        thrift_message_parser parser;

        std::string data;
        int read_next = 0;
        message_reader reader(64);
        data = std::string("THFT") + std::string(i, ' ');

        data_output out(&data[4], 8);
        out.write_u32(0);
        out.write_u32(48);

        mock_reader_read_data(reader, data);
        ASSERT_EQ(reader.buffer().size(), data.length());

        message_ex *msg = parser.get_message_on_receive(&reader, read_next);
        ASSERT_EQ(msg, nullptr);
        ASSERT_EQ(read_next, 48 - data.length()); // read remaining fields
        ASSERT_EQ(parser._header_version, -1);
    }
}

TEST_F(thrift_message_parser_test, get_message_on_receive_invalid_v0_hdr_length)
{
    for (int i = 0; i < 48; i++) {
        thrift_message_parser parser;

        std::string data;
        int read_next = 0;
        message_reader reader(64);
        data = std::string("THFT") + std::string(44, '\0'); // full 48 bytes

        // hdr_version = 0
        data_output out(&data[4], 8);
        out.write_u32(0);
        // hdr_length = i
        out.write_u32(i);

        mock_reader_read_data(reader, data);
        message_ex *msg = parser.get_message_on_receive(&reader, read_next);
        ASSERT_EQ(msg, nullptr);
        ASSERT_EQ(read_next, -1);
        ASSERT_EQ(parser._header_version, -1);
    }
}

TEST_F(thrift_message_parser_test, get_message_on_receive_valid_v0_hdr)
{
    thrift_message_parser parser;
    std::string data;
    int read_next = 0;
    message_reader reader(64);
    data = std::string("THFT") + std::string(44, '\0'); // full 48 bytes
    data_output out(&data[4], 44);
    out.write_u32(0);          // hdr_version
    out.write_u32(48);         // hdr_length
    out.write_u32(0);          // hdr_crc32
    out.write_u32(100);        // body_length
    out.write_u32(0);          // body_crc32
    out.write_u32(1);          // app_id
    out.write_u32(28);         // partition_index
    out.write_u32(1000);       // client_timeout
    out.write_u32(64);         // client_thread_hash
    out.write_u64(5000000000); // client_partition_hash

    mock_reader_read_data(reader, data);

    message_ex *msg = parser.get_message_on_receive(&reader, read_next);
    ASSERT_EQ(msg, nullptr);
    ASSERT_EQ(read_next, 100); // required to read more
    ASSERT_EQ(parser._header_version, 0);
    ASSERT_EQ(reader.buffer().size(), 0);
    ASSERT_EQ(parser._meta_v0->hdr_crc32, 0);
    ASSERT_EQ(parser._meta_v0->body_length, 100);
    ASSERT_EQ(parser._meta_v0->body_crc32, 0);
    ASSERT_EQ(parser._meta_v0->app_id, 1);
    ASSERT_EQ(parser._meta_v0->partition_index, 28);
    ASSERT_EQ(parser._meta_v0->client_timeout, 1000);
    ASSERT_EQ(parser._meta_v0->client_thread_hash, 64);
    ASSERT_EQ(parser._meta_v0->client_partition_hash, 5000000000);
}

TEST_F(thrift_message_parser_test, get_message_on_receive_valid_v0_data)
{
    message_reader reader(64);

    ASSERT_NO_FATAL_FAILURE(
        test_get_message_on_receive_v0_data(reader, apache::thrift::protocol::T_CALL, true));
    ASSERT_NO_FATAL_FAILURE(
        test_get_message_on_receive_v0_data(reader, apache::thrift::protocol::T_ONEWAY, true));
}

TEST_F(thrift_message_parser_test, get_message_on_receive_v0_not_request)
{
    message_reader reader(64);

    // ensure server won't corrupt when it receives a non-request.
    ASSERT_NO_FATAL_FAILURE(
        test_get_message_on_receive_v0_data(reader, apache::thrift::protocol::T_REPLY, false));
    reader.truncate_read();
    ASSERT_NO_FATAL_FAILURE(test_get_message_on_receive_v0_data(
        reader, apache::thrift::protocol::TMessageType(65), false));
    reader.truncate_read();
}

TEST_F(thrift_message_parser_test, get_message_on_receive_incomplete_v1_hdr)
{
    for (int i = 4; i < 12; i++) {
        thrift_message_parser parser;

        std::string data;
        int read_next = 0;
        message_reader reader(64);
        data = std::string("THFT") + std::string(i, ' ');

        data_output out(&data[4], 8);
        out.write_u32(1);

        mock_reader_read_data(reader, data);
        ASSERT_EQ(reader.buffer().size(), data.length());

        message_ex *msg = parser.get_message_on_receive(&reader, read_next);
        ASSERT_EQ(msg, nullptr);
        ASSERT_EQ(read_next, 16 - data.length()); // read remaining fields
        ASSERT_EQ(parser._header_version, -1);
        ASSERT_EQ(parser._v1_specific_vars->_meta_length, 0);
        ASSERT_EQ(parser._v1_specific_vars->_body_length, 0);
    }
}

TEST_F(thrift_message_parser_test, get_message_on_receive_valid_v1_hdr)
{
    thrift_message_parser parser;
    std::string data;
    int read_next = 0;
    message_reader reader(64);
    data = std::string("THFT") + std::string(12, '\0'); // full 12 bytes
    data_output out(&data[4], 12);
    out.write_u32(1);   // header_version
    out.write_u32(100); // meta_length
    out.write_u32(200); // body_length

    mock_reader_read_data(reader, data);
    ASSERT_EQ(reader.buffer().size(), 16);

    message_ex *msg = parser.get_message_on_receive(&reader, read_next);
    ASSERT_EQ(msg, nullptr);
    ASSERT_EQ(read_next, 100); // required to read more
    ASSERT_EQ(parser._header_version, 1);
    ASSERT_EQ(parser._v1_specific_vars->_meta_length, 100);
    ASSERT_EQ(parser._v1_specific_vars->_body_length, 200);
    ASSERT_EQ(parser._v1_specific_vars->_meta_parsed, false);
    ASSERT_EQ(reader.buffer().size(), 0);
}

TEST_F(thrift_message_parser_test, get_message_on_receive_v1_data)
{
    message_reader reader(64);
    ASSERT_NO_FATAL_FAILURE(
        test_get_message_on_receive_v1_data(reader, apache::thrift::protocol::T_CALL, true, true));
    ASSERT_NO_FATAL_FAILURE(
        test_get_message_on_receive_v1_data(reader, apache::thrift::protocol::T_CALL, true, false));
    ASSERT_NO_FATAL_FAILURE(test_get_message_on_receive_v1_data(
        reader, apache::thrift::protocol::T_ONEWAY, true, true));
    ASSERT_NO_FATAL_FAILURE(test_get_message_on_receive_v1_data(
        reader, apache::thrift::protocol::T_ONEWAY, true, false));

    ASSERT_NO_FATAL_FAILURE(test_get_message_on_receive_v1_data(
        reader, apache::thrift::protocol::TMessageType(65), false, false));
    reader.truncate_read();
}

} // namespace dsn
