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_reassembly.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_reassembly *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_options *dev_opts = &stellar_context.config.dev_opts;
struct packet_io_options *pkt_io_opts = &stellar_context.config.pkt_io_opts;
struct ip_reassembly_options *ip_reass_opts = &stellar_context.config.ip_reass_opts;
struct session_manager_options *sess_mgr_opts = &stellar_context.config.sess_mgr_opts;

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

// TODO
void *plugin_manager_new_ctx()
{
    return NULL;
}

void plugin_manager_free_ctx(void *ctx)
{
    return;
}

void plugin_manager_dispatch(void *plugin_mgr, struct session *sess, const struct packet *pkt)
{
    if (sess == NULL)
    {
        return;
    }

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

/******************************************************************************
 * 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 session *sess;
    struct session *evicted_sess;
    struct session *expired_sess;
    struct packet *pkt;
    struct packet packets[RX_BURST_MAX];
    struct thread_context *threads_ctx = (struct thread_context *)arg;
    struct packet_io *packet_io = stellar_ctx->packet_io;
    struct session_manager *sess_mgr = threads_ctx->sess_mgr;
    struct ip_reassembly *ip_mgr = threads_ctx->ip_mgr;
    void *plug_mgr = NULL;
    void *plug_mgr_ctx = NULL;

    int nr_recv;
    uint16_t thd_idx = threads_ctx->index;
    uint64_t now_msec = 0;

    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_STATE("worker thread %d runing", thd_idx);

    while (ATOMIC_READ(&threads_ctx->need_exit) == 0)
    {
        now_msec = timestamp_get_msec();
        nr_recv = packet_io_ingress(packet_io, thd_idx, packets, RX_BURST_MAX);
        if (nr_recv == 0)
        {
            goto idle_tasks;
        }

        for (int i = 0; i < nr_recv; i++)
        {
            pkt = &packets[i];

            // TODO
            // call packet plugin

            // ip fragment reassemble
            if (pkt->frag_layer)
            {
                struct packet *temp = ip_reassembly_packet(ip_mgr, pkt);
                packet_io_egress(packet_io, thd_idx, pkt, 1); // forward the original fragment packet
                if (temp == NULL)
                {
                    continue;
                }
                else
                {
                    pkt = temp;
                }
            }

            sess = session_manager_lookup_session(sess_mgr, pkt);
            if (sess == NULL)
            {
                sess = session_manager_new_session(sess_mgr, pkt);
                if (sess == NULL)
                {
                    // 1.Not TCP or UDP
                    // 2.UDP evict packet
                    // 3.UDP overloading and config to bypass new session
                    // 4.TCP no SYN flag
                    // 5.UDP overloading and config to bypass new session
                    goto fast_forward;
                }
                plug_mgr_ctx = plugin_manager_new_ctx();
                session_set_user_data(sess, plug_mgr_ctx);
            }
            else
            {
                if (session_manager_update_session(sess_mgr, sess, pkt) == -1)
                {
                    // TCP duplicate packet
                    goto fast_forward;
                }
            }
            plugin_manager_dispatch(plug_mgr, sess, pkt);

        fast_forward:
            packet_io_egress(packet_io, thd_idx, pkt, 1);

            evicted_sess = session_manager_get_evicted_session(sess_mgr);
            if (evicted_sess)
            {
                plug_mgr_ctx = session_get_user_data(evicted_sess);
                plugin_manager_free_ctx(plug_mgr_ctx);
                session_manager_free_session(sess_mgr, evicted_sess);
            }
        }

    idle_tasks:
        expired_sess = session_manager_get_expired_session(sess_mgr);
        if (expired_sess)
        {
            plug_mgr_ctx = session_get_user_data(expired_sess);
            plugin_manager_free_ctx(plug_mgr_ctx);
            session_manager_free_session(sess_mgr, expired_sess);
        }

        // TODO
        // plugin_manager_cron();
        // poll_non_packet_events();
        // packet_io_yield();
    }

    ATOMIC_SET(&threads_ctx->is_runing, 0);
    STELLAR_LOG_STATE("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_new(sess_mgr_opts);
        if (threads_ctx->sess_mgr == NULL)
        {
            STELLAR_LOG_ERROR("unable to create session manager");
            return -1;
        }

        threads_ctx->ip_mgr = ip_reassembly_new(ip_reass_opts);
        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_free(threads_ctx->sess_mgr);
            ip_reassembly_free(threads_ctx->ip_mgr);
        }
    }
}

static int thread_new(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_free(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_opts->device_base, dev_opts->device_offset) != 0)
    {
        STELLAR_LOG_ERROR("unable to init id generator");
        return -1;
    }

    // TODO load plugin

    uint8_t nr_threads = pkt_io_opts->nr_threads;
    stellar_ctx->packet_io = packet_io_new(pkt_io_opts);
    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_new(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_free(stellar_ctx->threads_ctx, nr_threads);
    thread_context_free(stellar_ctx, nr_threads);
    packet_io_free(stellar_ctx->packet_io);

    // TODO free plugin

    id_generator_free();
    log_free();

    return 0;
}