tango-tfe/common/test/test_mpack.cpp

#include <gtest/gtest.h>
#include <arpa/inet.h>

#include "mpack.h"
#include "tfe_cmsg.h"
#include "tfe_utils.h"
#include "tfe_ctrl_packet.h"

bool g_print_to_stderr = true;
void *g_packet_io_logger = NULL;

void build_mpack_data(char **data, size_t *size)
{
    mpack_writer_t writer;

	uint32_t tcp_seq = 3750033295;
	uint32_t tcp_ack = 739256665;
	uint16_t mss_client = 1460;
	uint16_t mss_server = 1460;
	uint8_t wsacle_client_flag = 1;
	uint8_t wsacle_server_flag = 1;
	uint8_t wsacle_client = 10;
	uint8_t wsacle_server = 10;
	uint8_t sack_client = 1;
	uint8_t sack_server = 1;
	uint8_t ts_client = 1;
	uint8_t ts_server = 1;
	uint8_t protocol = 0;
	uint16_t window_client = 64240;
	uint16_t window_server = 65160;
	uint32_t ts_client_val = 4213358146;
	uint32_t ts_server_val = 1218524397;
	uint8_t packet_cur_dir = 1;
	uint8_t src_sub_id[256] = {0};
	uint8_t dst_sub_id[256] = {0};
	uint64_t src_asn_val = 7676;
	uint64_t dst_asn_val = 7878;
	uint64_t src_asn_id = 1234;
	uint64_t dst_asn_id = 4321;
	// ip location array
	const char *src_ip_location[4] = {
		"src_region",
		"src_province",
		"src_city",
		"src_subdivision"
	};
	const char *dst_ip_location[4] = {
		"dst_region",
		"dst_province",
		"dst_city",
		"dst_subdivision"
	};

	uint64_t src_ip_location_obj_id[4] = {1, 2, 3, 4};
	uint64_t dst_ip_location_obj_id[4] = {1, 2, 3, 4};
	uint8_t ja3_fingerprint[32] = {0};
	uint16_t seq_sid[4] = {1000, 1001, 1002, 1003};
	uint16_t ack_sid[4] = {1003, 1002, 1001, 1000};
	uint8_t seq_route_ctx[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07};
	uint8_t ack_route_ctx[] = {0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F};
	uint8_t seq_header[] = {0x00, 0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70};
	uint8_t ack_header[] = {0x80, 0x90, 0xA0, 0xB0, 0xC0, 0xD0, 0xE0, 0xF0};
	uint8_t tfe_flag = 0;

	mpack_writer_init_growable(&writer, data, size);
	mpack_build_map(&writer);

    mpack_write_cstr(&writer, "tsync");
    mpack_write_cstr(&writer, "2.0");

    mpack_write_cstr(&writer, "session_id");
    mpack_write_u64(&writer, 290591070686249153);

    mpack_write_cstr(&writer, "state");
    mpack_write_cstr(&writer, "active");

    mpack_write_cstr(&writer, "method");
    mpack_write_cstr(&writer, "policy_update");

    mpack_write_cstr(&writer, "params");
    // params map
    mpack_build_map(&writer);

	mpack_write_cstr(&writer, "sce");
    mpack_build_map(&writer);
    mpack_write_cstr(&writer, "rule_ids");
    mpack_build_array(&writer);
    mpack_write_u64(&writer, 1);
	mpack_complete_array(&writer);
	mpack_complete_map(&writer);

    mpack_write_cstr(&writer, "proxy");
    mpack_build_map(&writer);
    mpack_write_cstr(&writer, "rule_ids");
    mpack_build_array(&writer);
    mpack_write_u64(&writer, 3);
	mpack_complete_array(&writer);

	mpack_write_cstr(&writer, "tcp_handshake");
	mpack_build_array(&writer);
	mpack_write_u32(&writer, tcp_seq);
	mpack_write_u32(&writer, tcp_ack);
	mpack_write_u16(&writer, mss_client);
	mpack_write_u16(&writer, mss_server);
	mpack_write_u8(&writer, wsacle_client_flag);
	mpack_write_u8(&writer, wsacle_server_flag);
	mpack_write_u8(&writer, wsacle_client);
	mpack_write_u8(&writer, wsacle_server);
	mpack_write_u8(&writer, sack_client);
	mpack_write_u8(&writer, sack_server);
	mpack_write_u8(&writer, ts_client);
	mpack_write_u8(&writer, ts_server);
	mpack_write_u8(&writer, protocol);
	mpack_write_u16(&writer, window_client);
	mpack_write_u16(&writer, window_server);
	mpack_write_u32(&writer, ts_client_val);
	mpack_write_u32(&writer, ts_server_val);
	mpack_write_u8(&writer, packet_cur_dir);
	mpack_write_str(&writer, (const char*)src_sub_id, 0);
	mpack_write_str(&writer, (const char*)dst_sub_id, 0);
	mpack_write_u64(&writer, src_asn_val);
	mpack_write_u64(&writer, dst_asn_val);
	mpack_write_u64(&writer, src_asn_id);
	mpack_write_u64(&writer, dst_asn_id);

	mpack_build_array(&writer);
	mpack_write_str(&writer, src_ip_location[0], strlen(src_ip_location[0]));
	mpack_write_str(&writer, src_ip_location[1], strlen(src_ip_location[1]));
	mpack_write_str(&writer, src_ip_location[2], strlen(src_ip_location[2]));
	mpack_write_str(&writer, src_ip_location[3], strlen(src_ip_location[3]));
	mpack_complete_array(&writer);

	mpack_build_array(&writer);
	mpack_write_str(&writer, dst_ip_location[0], strlen(dst_ip_location[0]));
	mpack_write_str(&writer, dst_ip_location[1], strlen(dst_ip_location[1]));
	mpack_write_str(&writer, dst_ip_location[2], strlen(dst_ip_location[2]));
	mpack_write_str(&writer, dst_ip_location[3], strlen(dst_ip_location[3]));
	mpack_complete_array(&writer);

	mpack_build_array(&writer);
	mpack_write_u64(&writer, src_ip_location_obj_id[0]);
	mpack_write_u64(&writer, src_ip_location_obj_id[1]);
	mpack_write_u64(&writer, src_ip_location_obj_id[2]);
	mpack_write_u64(&writer, src_ip_location_obj_id[3]);
	mpack_complete_array(&writer);

	mpack_build_array(&writer);
	mpack_write_u64(&writer, dst_ip_location_obj_id[0]);
	mpack_write_u64(&writer, dst_ip_location_obj_id[1]);
	mpack_write_u64(&writer, dst_ip_location_obj_id[2]);
	mpack_write_u64(&writer, dst_ip_location_obj_id[3]);
	mpack_complete_array(&writer);

	mpack_write_str(&writer, (const char*)ja3_fingerprint, 0);

	mpack_build_array(&writer);
	int cnt = sizeof(seq_sid)/sizeof(uint16_t);
	for (int i = 0; i < cnt; i++) {
		mpack_write_u16(&writer, seq_sid[i]);
	}
	mpack_complete_array(&writer);

	mpack_build_array(&writer);
	cnt = sizeof(ack_sid)/sizeof(uint16_t);
	for (int i = 0; i < cnt; i++) {
		mpack_write_u16(&writer, ack_sid[i]);
	}
	mpack_complete_array(&writer);

	mpack_write_bin(&writer, (const char*)seq_route_ctx, sizeof(seq_route_ctx));
	mpack_write_bin(&writer, (const char*)ack_route_ctx, sizeof(ack_route_ctx));
	mpack_write_bin(&writer, (const char*)seq_header, sizeof(seq_header));
	mpack_write_bin(&writer, (const char*)ack_header, sizeof(ack_header));
	mpack_write_u8(&writer, tfe_flag);

	mpack_complete_array(&writer);
	mpack_complete_map(&writer);
	mpack_complete_map(&writer);
	mpack_complete_map(&writer);

	// finish writing
    if (mpack_writer_destroy(&writer) != mpack_ok) {
        assert(0);
        if (*data) {
            free(*data);
            *data = NULL;
        }
        return;
    }
}

TEST(CONTROL_PACKET, PARSER)
{
	int ret = 0;
	char *data = NULL;
	size_t size = 0;

	build_mpack_data(&data, &size);
	if (data == NULL) {
		printf("build mpack data error");
		return;
	}

	struct ctrl_pkt_parser ctrl_parser;
	memset(&ctrl_parser, 0, sizeof(struct ctrl_pkt_parser));
	ret = ctrl_packet_parser_parse((void *)&ctrl_parser, (const char*)data, size, g_packet_io_logger, 1);
	EXPECT_TRUE(ret == 0);
}

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