stellar-stellar/test/lpi_plus/gtest_lpip_module.c

#include <stdlib.h>
#pragma GCC diagnostic ignored "-Wunused-parameter"

#include <stdio.h>
#include <time.h>
#include <unistd.h>
#include <assert.h>

#include "stellar/appid.h"
#include "stellar/module_manager.h"
#include "stellar/session.h"
#include "stellar/utils.h"

#include "cjson/cJSON.h"

#include "lpi_plus/lpi_plus.h"
#include "appid/appid_internal.h"

struct test_lpip_env
{
    struct stellar_module_manager *mod_mgr;
    struct lpi_plus_mapper *lpi_mapper;
    int l7_exdata_idx;
    int session_num;
};

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

cJSON *g_result_json=NULL;

void stellar_test_result_setup()
{
    if(g_result_json!=NULL)return;
    g_result_json=cJSON_CreateArray();
}

char *stellar_test_result_json_export()
{
    if(g_result_json==NULL)return NULL;
    return cJSON_Print(g_result_json);
}

void stellar_test_result_cleanup()
{
    if(g_result_json)cJSON_Delete(g_result_json);
}

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

    if(env==NULL || 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->lpi_mapper ,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);

    if(g_result_json)cJSON_AddItemToArray(g_result_json, ctx);
    free(test_appid_exdata);
    return;
}


static void gtest_lpip_on_appid_msg(struct session *sess, enum APPID_ORIGIN origin, int appid[], size_t appid_num, void *args)
{
    if(sess==NULL || appid==NULL || args==NULL)return;
    struct test_lpip_env *env = (struct test_lpip_env *)args;
    struct test_lpip_exdata *test_appid_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;

    }
    memcpy(test_appid_exdata->appid, appid, appid_num*sizeof(appid[0]));
    test_appid_exdata->appid_num=appid_num;
    session_set_exdata(sess, env->l7_exdata_idx, test_appid_exdata);
    return;
}

static void on_session(struct session *sess, enum session_state state, struct packet *pkt, void *args)
{
    if (state == SESSION_STATE_CLOSED)
    {
        assert(pkt == NULL);
        return;
    }

    struct test_lpip_env *env = (struct test_lpip_env *)args;
    if (session_get_current_state(sess) == SESSION_STATE_OPENING)
    {
        struct test_lpip_exdata *test_appid_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);
        }
    }
    return;
}

struct stellar_module *gtest_lpip_module_init(struct stellar_module_manager *mod_mgr)
{
    struct test_lpip_env *env = (struct test_lpip_env *)calloc(1, sizeof(struct test_lpip_env));

    env->lpi_mapper=stellar_module_get_lpip(mod_mgr);
    struct session_manager *sess_mgr = stellar_module_get_session_manager(mod_mgr);
    if(sess_mgr == NULL)
    {
        perror("gtest_lpi_plugin_load:stellar_module_get_session_manager failed !!!\n");
        exit(-1);
    }

    session_manager_subscribe_udp(sess_mgr, on_session, env);
    session_manager_subscribe_tcp(sess_mgr, on_session, env);

    env->l7_exdata_idx = session_manager_new_session_exdata_index(sess_mgr, "EXDATA_L7", gtest_lpip_exdata_free, env);

    stellar_appid_subscribe(mod_mgr, gtest_lpip_on_appid_msg, env);
    printf("gtest_lpip_module_init OK!\n");	

    return stellar_module_new("TEST_LPIP", env);
}

void gtest_lpip_module_exit(struct stellar_module_manager *mod_mgr, struct stellar_module *mod)
{
    assert(mod_mgr!=NULL);
    struct test_lpip_env *env = (struct test_lpip_env *)stellar_module_get_ctx(mod);
	free(env);
    printf("gtest_lpip_module_exit OK!\n");	
    stellar_module_free(mod);
    return ;
}