stellar-stellar/test/lpi_plus/gtest_lpip_main.cpp

/*
 *  author:yangwei
 *  create time:2021-8-21
 *
 */

#pragma GCC diagnostic ignored "-Wunused-parameter"

#include <stdio.h>
#include <unistd.h>
#include <assert.h>
#include <gtest/gtest.h>
#include <stdlib.h>
#include <time.h>


#include "test_result_validator.h"

#include "stellar/stellar.h"
#include "stellar/module.h"
#include "stellar/session.h"
#include "stellar/utils.h"

#include "stellar/lpi_plus.h"


struct test_lpip_env
{
    struct module_manager *mod_mgr;
    struct session_manager *sess_mgr;
	struct test_result_validator *result;
    struct lpi_plus *lpip;
    int l7_exdata_idx;
    int session_num;
};

struct test_lpip_env env;

#ifndef MAX_APPID_NUM
#define MAX_APPID_NUM 8
#endif

struct test_lpip_exdata
{
    int appid[MAX_APPID_NUM];
    size_t appid_num;
    struct session *sess;
};


static void gtest_lpip_exdata_free(int idx __attribute__((unused)), void *ex_ptr, void *arg)
{
    struct test_lpip_exdata *test_appid_exdata=(struct test_lpip_exdata *)ex_ptr;

    if(test_appid_exdata ==NULL)return;

    cJSON *ctx = cJSON_CreateObject();
    cJSON_AddStringToObject(ctx, "Tuple4", session_get_readable_addr(test_appid_exdata->sess));
    enum session_type type = session_get_type(test_appid_exdata->sess);
    if (type == SESSION_TYPE_TCP)
    {
        cJSON_AddStringToObject(ctx, "STREAM_TYPE", "TCP");
    }
    if (type == SESSION_TYPE_UDP)
    {
        cJSON_AddStringToObject(ctx, "STREAM_TYPE", "UDP");
    }
    if (test_appid_exdata->appid_num > 0)
    {
        const char *proto_names[MAX_APPID_NUM] = {};
        for (unsigned int i = 0; i < test_appid_exdata->appid_num; i++)
        {
            proto_names[i] = lpi_plus_appid2name(env.lpip ,test_appid_exdata->appid[i]);
        }
        cJSON *label_ids = cJSON_CreateIntArray(test_appid_exdata->appid, test_appid_exdata->appid_num);
        cJSON_AddItemToObject(ctx, "l7_label_id", label_ids);
        cJSON *label_names = cJSON_CreateStringArray(proto_names, test_appid_exdata->appid_num);
        cJSON_AddItemToObject(ctx, "l7_label_name", label_names);
    }
    else
    {
        cJSON_AddStringToObject(ctx, "l7_label_id", "UNKNOWN");
    }
    unsigned char dir_flag;
    int is_symmetric = session_is_symmetric(test_appid_exdata->sess, &dir_flag);
    if (is_symmetric)
    {
        cJSON_AddStringToObject(ctx, "STREAM_DIR", "DOUBLE");
    }
    else if (dir_flag == SESSION_SEEN_C2S_FLOW)
    {
        cJSON_AddStringToObject(ctx, "STREAM_DIR", "C2S");
    }
    else if (dir_flag == SESSION_SEEN_S2C_FLOW)
    {
        cJSON_AddStringToObject(ctx, "STREAM_DIR", "S2C");
    }
    else
    {
        assert(0);
    }
    char result_name[128] = "";
    env.session_num++;
    sprintf(result_name, "APP_PROTO_IDENTIFY_RESULT_%d", env.session_num);
    cJSON_AddStringToObject(ctx, "name", result_name);

	test_result_validator_add_actual(env.result, ctx);
    free(test_appid_exdata);
    return;
}


void gtest_lpip_on_packet(struct packet *pkt, struct module *mod)
{
    struct session *sess=packet_exdata_to_session(env.sess_mgr, pkt);
    if(sess==NULL)return;
    struct test_lpip_exdata *test_appid_exdata=(struct test_lpip_exdata *)session_get_exdata(sess, env.l7_exdata_idx);
    if (test_appid_exdata==NULL)
    {
        test_appid_exdata = CALLOC(struct test_lpip_exdata, 1);
        test_appid_exdata->sess=sess;
        session_set_exdata(sess, env.l7_exdata_idx, test_appid_exdata);
    }

    size_t appid_num=0;
    int32_t *appid = packet_exdata_to_lpip_appid(env.lpip, pkt,  &appid_num);
    if(appid_num>0 && appid!=NULL)
    {
       test_appid_exdata->appid_num=appid_num;
       memcpy(test_appid_exdata->appid, appid, sizeof(*appid)*appid_num);
    }
    return;
}

/**********************************************
 * GTEST MAIN *
 **********************************************/

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

	printf("Usage: ./[gtest_main] [/path/to/expect_json]\n");

	env.result = test_result_validator_new(argv[1]);
	
	struct stellar *st=stellar_new("./conf/stellar.toml");

	env.mod_mgr=stellar_get_module_manager(st);

	// init gtest module
    struct module *lpip_mod = module_manager_get_module(env.mod_mgr, LPI_PLUS_MODULE_NAME);
    env.lpip=module_to_lpi_plus(lpip_mod);
    struct module *sess_mgr_mod=module_manager_get_module(env.mod_mgr, SESSION_MANAGER_MODULE_NAME);
    env.sess_mgr = module_to_session_manager(sess_mgr_mod);
    env.l7_exdata_idx = session_manager_new_session_exdata_index(env.sess_mgr, "EXDATA_L7", gtest_lpip_exdata_free, &env);


	struct module *pkt_mgr_mod=module_manager_get_module(env.mod_mgr, PACKET_MANAGER_MODULE_NAME);
	struct packet_manager *pkt_mgr=module_to_packet_manager(pkt_mgr_mod);

	packet_manager_register_node(pkt_mgr, "GTEST_LPIP", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP | PKT_TAG_VAL_IPPROTO_UDP, gtest_lpip_on_packet, NULL);

	stellar_run(st);
	stellar_free(st);

	EXPECT_EQ(test_result_validator_compare(env.result), 1);    
	test_result_validator_free(env.result);

	return ::testing::Test::HasFailure() ? 1 : 0;
}