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 *device_opts = &stellar_context.config.device_opts;
struct packet_io_options *packet_io_opts = &stellar_context.config.packet_io_opts;
struct ip_reassembly_options *ip_reassembly_opts = &stellar_context.config.ip_reassembly_opts;
struct session_manager_options *session_manager_opts = &stellar_context.config.session_manager_opts;

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

static void hex_dump(const char *payload, uint32_t len)
{
    printf("Payload Length: %d\n", len);
    for (uint32_t i = 0; i < len; i++)
    {
        if (i > 0 && i % 16 == 0)
        {
            printf("\n");
        }
        printf("%02x ", (uint8_t)payload[i]);
    }
    printf("\n");
}

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;
    }

    uint32_t len = 0;
    const char *payload = NULL;

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

    if (session_get_type(sess) == SESSION_TYPE_TCP)
    {
        do
        {
            payload = tcp_data_peek(sess, &len);
            if (payload && len > 0)
            {
                hex_dump(payload, len);
            }
            tcp_data_dequeue(sess, len);
        } while (payload && len > 0);
    }
    session_clean_packet(sess, SESSION_PACKET_CURRENT);
    session_set_cur_dir(sess, SESSION_DIR_NONE);
    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_reass = threads_ctx->ip_mgr;
    void *plug_mgr = NULL;
    void *plug_mgr_ctx = NULL;

    int nr_recv;
    int need_drop_pkt = 1; // TODO
    uint64_t now = 0;
    uint16_t thr_idx = threads_ctx->index;

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

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

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

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

            // call plugin_manager_dispatch_raw_pkt();
            if (packet_is_fragment(pkt))
            {
                struct packet *defraged_pkt = ip_reassembly_packet(ip_reass, pkt, now);
                if (defraged_pkt == NULL)
                {
                    continue;
                }
                else
                {
                    // call plugin_manager_dispatch_defrag_pkt();
                    pkt = defraged_pkt;
                }
            }

            // TODO filter protocol before session lookup

            sess = session_manager_lookup_session(sess_mgr, pkt);
            if (sess == NULL)
            {
                sess = session_manager_new_session(sess_mgr, pkt, now);
                if (sess == NULL)
                {
                    continue;
                }
                plug_mgr_ctx = plugin_manager_new_ctx();
                session_set_user_data(sess, plug_mgr_ctx);
            }
            else
            {
                session_manager_update_session(sess_mgr, sess, pkt, now);
            }
            plugin_manager_dispatch(plug_mgr, sess, pkt);
        }

        if (unlikely(need_drop_pkt))
        {
            for (int i = 0; i < nr_recv; i++)
            {
                if (packet_get_action(&packets[i]) == PACKET_ACTION_DROP)
                {
                    packet_io_drop(packet_io, thr_idx, &packets[i], 1);
                }
                else
                {
                    packet_io_egress(packet_io, thr_idx, &packets[i], 1);
                }
            }
        }
        else
        {
            packet_io_egress(packet_io, thr_idx, packets, nr_recv);
        }

    idle_tasks:
        // nr_recv packet atmost trigger nr_recv session evict
        for (int i = 0; i < nr_recv; i++)
        {
            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);
            }
        }

        while ((expired_sess = session_manager_get_expired_session(sess_mgr, now)))
        {
            plug_mgr_ctx = session_get_user_data(expired_sess);
            plugin_manager_free_ctx(plug_mgr_ctx);
            session_manager_free_session(sess_mgr, expired_sess);
        }
        ip_reassembly_expire(ip_reass, now);

        // 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", thr_idx);

    return NULL;
}

static int thread_context_init(struct stellar_context *ctx, uint8_t nr_threads)
{
    uint64_t now = timestamp_get_msec();
    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(session_manager_opts, now);
        if (threads_ctx->sess_mgr == NULL)
        {
            STELLAR_LOG_ERROR("unable to create session manager");
            return -1;
        }

        threads_ctx->ip_mgr = ip_reassembly_new(ip_reassembly_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 (parse_config_file(stellar_config_file, &stellar_ctx->config) != 0)
    {
        STELLAR_LOG_ERROR("unable to load config file");
        return -1;
    }

    print_config_options(&stellar_ctx->config);

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

    // TODO load plugin

    uint8_t nr_threads = packet_io_opts->nr_threads;
    stellar_ctx->packet_io = packet_io_new(packet_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

    log_free();

    return 0;
}