stellar-stellar/src/stellar/stellar.cpp

#include <stdio.h>
#include <assert.h>
#include <errno.h>
#include <pthread.h>
#include <unistd.h>
#include <string.h>
#include <signal.h>
#include <stdlib.h>
#include <sys/prctl.h>

#include "logo.h"
#include "stellar.h"
#include "config.h"
#include "packet.h"
#include "packet_io.h"
#include "timestamp.h"
#include "id_generator.h"
#include "ip_reassemble.h"
#include "session_manager.h"

#define STELLAR_LOG_STATE(format, ...) LOG_STATE("stellar", format, ##__VA_ARGS__)
#define STELLAR_LOG_ERROR(format, ...) LOG_ERROR("stellar", format, ##__VA_ARGS__)
#define STELLAR_LOG_DEBUG(format, ...) LOG_DEBUG("stellar", format, ##__VA_ARGS__)

#define ATOMIC_SET(x, y) __atomic_store_n(x, y, __ATOMIC_RELAXED)
#define ATOMIC_READ(x) __atomic_load_n(x, __ATOMIC_RELAXED)

struct thread_context
{
    pthread_t tid;
    uint16_t index;
    uint64_t need_exit;
    uint64_t is_runing;
    struct session_manager *sess_mgr;
    struct ip_reassemble_manager *ip_mgr;
};

struct stellar_context
{
    uint64_t need_exit;
    struct config config;

    struct packet_io *packet_io;
    struct thread_context threads_ctx[MAX_THREAD_NUM];
};

struct stellar_context stellar_context;
struct stellar_context *stellar_ctx = &stellar_context;

// config
struct device_config *dev_cfg = &stellar_context.config.dev_cfg;
struct packet_io_config *pkt_io_cfg = &stellar_context.config.pkt_io_cfg;
struct ip_reassemble_config *ip_reass_cfg = &stellar_context.config.ip_reass_cfg;
struct session_manager_config *sess_mgr_cfg = &stellar_context.config.sess_mgr_cfg;

static const char *log_config_file = "./conf/log.toml";
static const char *stellar_config_file = "./conf/stellar.toml";

/******************************************************************************
 * example
 ******************************************************************************/

static void __packet_plugin(const struct packet *pkt)
{
    if (pkt == NULL)
    {
        return;
    }

    printf("=> packet dispatch: %p\n", pkt);
    printf("<= packet dispatch\n");
}

static void __session_plugin(struct session *sess)
{
    if (sess == NULL)
    {
        return;
    }

    printf("=> session dispatch: %p\n", sess);
    session_dump(sess);
    printf("<= session dispatch\n");

    // after session dispatch, we should reset session current packet and direction
    session_set0_cur_pkt(sess, NULL);
    session_set_cur_dir(sess, SESSION_DIR_NONE);
}

/******************************************************************************
 * util
 ******************************************************************************/

static void signal_handler(int signo)
{
    if (signo == SIGINT)
    {
        STELLAR_LOG_STATE("SIGINT received, exit !!!");
        ATOMIC_SET(&stellar_ctx->need_exit, 1);
    }

    if (signo == SIGQUIT)
    {
        STELLAR_LOG_STATE("SIGQUIT received, exit !!!");
        ATOMIC_SET(&stellar_ctx->need_exit, 1);
    }

    if (signo == SIGTERM)
    {
        STELLAR_LOG_STATE("SIGTERM received, exit !!!");
        ATOMIC_SET(&stellar_ctx->need_exit, 1);
    }

    if (signo == SIGHUP)
    {
        STELLAR_LOG_STATE("SIGHUP received, reload log level !!!");
        log_reload_level(log_config_file);
    }
}

/******************************************************************************
 * thread
 ******************************************************************************/

static inline void thread_set_name(const char *thd_symbol, uint16_t thd_idx)
{
    char thd_name[16];
    snprintf(thd_name, sizeof(thd_name), "%s:%d", thd_symbol, thd_idx);
    prctl(PR_SET_NAME, (unsigned long long)thd_name, NULL, NULL, NULL);
}

static void *main_loop(void *arg)
{
    struct packet *pkt = NULL;
    struct session *sess;
    struct thread_context *threads_ctx = (struct thread_context *)arg;
    uint16_t thd_idx = threads_ctx->index;
    struct packet_io *packet_io = stellar_ctx->packet_io;
    struct session_manager *sess_mgr = threads_ctx->sess_mgr;
    struct ip_reassemble_manager *ip_mgr = threads_ctx->ip_mgr;

    if (packet_io_init(packet_io, thd_idx) != 0)
    {
        STELLAR_LOG_ERROR("unable to init marsio thread");
        return NULL;
    }

    ATOMIC_SET(&threads_ctx->is_runing, 1);
    thread_set_name("stellar", thd_idx);
    STELLAR_LOG_DEBUG("worker thread %d runing", thd_idx);

    while (ATOMIC_READ(&threads_ctx->need_exit) == 0)
    {
        // recv packet
        if (packet_io_recv(packet_io, thd_idx, &pkt) != 0)
        {
            goto poll_wait;
        }

        // call packet plugin
        __packet_plugin(pkt);

        // ip fragment reassemble
        if (pkt->frag_layer)
        {
            struct packet *temp = ip_reassemble_packet(ip_mgr, pkt);
            // forward the original fragment packet
            packet_io_send(packet_io, thd_idx, pkt);
            if (temp == NULL)
            {
                goto poll_wait;
            }
            else
            {
                pkt = temp;
            }
        }

        // update session
        sess = session_manager_update_session(sess_mgr, pkt);
        __session_plugin(sess);

        // get evicted session
        sess = session_manager_get_evicted_session(sess_mgr);
        __session_plugin(sess);

        packet_io_send(packet_io, thd_idx, pkt);

    poll_wait:
        // get expired session
        sess = session_manager_get_expired_session(sess_mgr);
        __session_plugin(sess);
    }

    ATOMIC_SET(&threads_ctx->is_runing, 0);
    STELLAR_LOG_DEBUG("worker thread %d stop", thd_idx);

    return NULL;
}

static int thread_context_init(struct stellar_context *ctx, uint8_t nr_threads)
{
    for (uint8_t i = 0; i < nr_threads; i++)
    {
        struct thread_context *threads_ctx = &ctx->threads_ctx[i];
        threads_ctx->index = i;
        threads_ctx->need_exit = 0;
        threads_ctx->is_runing = 0;

        threads_ctx->sess_mgr = session_manager_create(sess_mgr_cfg);
        if (threads_ctx->sess_mgr == NULL)
        {
            STELLAR_LOG_ERROR("unable to create session manager");
            return -1;
        }

        threads_ctx->ip_mgr = ip_reassemble_manager_create(ip_reass_cfg);
        if (threads_ctx->ip_mgr == NULL)
        {
            STELLAR_LOG_ERROR("unable to create ip reassemble manager");
            return -1;
        }
    }

    return 0;
}

static void thread_context_free(struct stellar_context *ctx, uint8_t nr_threads)
{
    for (uint8_t i = 0; i < nr_threads; i++)
    {
        struct thread_context *threads_ctx = &ctx->threads_ctx[i];
        if (ATOMIC_READ(&threads_ctx->is_runing) == 0)
        {
            STELLAR_LOG_STATE("wait worker thread %d free context", i);
            session_manager_destroy(threads_ctx->sess_mgr);
            ip_reassemble_manager_destory(threads_ctx->ip_mgr);
        }
    }
}

static int thread_create(struct thread_context threads_ctx[], uint8_t nr_threads)
{
    for (uint8_t i = 0; i < nr_threads; i++)
    {
        struct thread_context *ctx = &threads_ctx[i];
        if (pthread_create(&ctx->tid, NULL, main_loop, (void *)ctx) < 0)
        {
            STELLAR_LOG_ERROR("unable to create worker thread, error %d: %s", errno, strerror(errno));
            return -1;
        }
    }

    return 0;
}

static void thread_destroy(struct thread_context threads_ctx[], uint8_t nr_threads)
{
    for (uint8_t i = 0; i < nr_threads; i++)
    {
        struct thread_context *ctx = &threads_ctx[i];
        ATOMIC_SET(&ctx->need_exit, 1);
        while (ATOMIC_READ(&ctx->is_runing) == 1)
        {
            STELLAR_LOG_STATE("wait worker thread %d stop", i);
            sleep(1);
        }
    }
}

/******************************************************************************
 * main
 ******************************************************************************/

int main(int argc, char **argv)
{
    memset(stellar_ctx, 0, sizeof(struct stellar_context));
    timestamp_update();

    signal(SIGINT, signal_handler);
    signal(SIGQUIT, signal_handler);
    signal(SIGTERM, signal_handler);
    signal(SIGHUP, signal_handler);

    if (log_init(log_config_file) != 0)
    {
        STELLAR_LOG_ERROR("unable to init log");
        return -1;
    }

    STELLAR_LOG_STATE("Start Stellar (version: %s)\n %s", __stellar_version, logo_str);

    if (config_load(&stellar_ctx->config, stellar_config_file) != 0)
    {
        STELLAR_LOG_ERROR("unable to load config file");
        return -1;
    }

    config_dump(&stellar_ctx->config);

    if (id_generator_init(dev_cfg->device_base, dev_cfg->device_offset) != 0)
    {
        STELLAR_LOG_ERROR("unable to init id generator");
        return -1;
    }

    // TODO load plugin

    uint8_t nr_threads = pkt_io_cfg->nr_threads;
    stellar_ctx->packet_io = packet_io_create(pkt_io_cfg);
    if (stellar_ctx->packet_io == NULL)
    {
        STELLAR_LOG_ERROR("unable to create packet io");
        goto error_out;
    }

    if (thread_context_init(stellar_ctx, nr_threads) != 0)
    {
        STELLAR_LOG_ERROR("unable to init thread context");
        goto error_out;
    }

    if (thread_create(stellar_ctx->threads_ctx, nr_threads) != 0)
    {
        STELLAR_LOG_ERROR("unable to create worker thread");
        goto error_out;
    }

    while (!ATOMIC_READ(&stellar_ctx->need_exit))
    {
        timestamp_update();
        sleep(1);
    }

error_out:
    thread_destroy(stellar_ctx->threads_ctx, nr_threads);
    thread_context_free(stellar_ctx, nr_threads);
    packet_io_destroy(stellar_ctx->packet_io);

    // TODO free plugin

    id_generator_free();
    log_free();

    return 0;
}