/* Copyright (C) 2007-2023 Open Information Security Foundation
 *
 * You can copy, redistribute or modify this Program under the terms of
 * the GNU General Public License version 2 as published by the Free
 * Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * version 2 along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301, USA.
 */

/**
 * \file
 *
 * \author Victor Julien <victor@inliniac.net>
 *
 * Flow queue handler functions
 */

#include "suricata-common.h"
#include "threads.h"
#include "flow-private.h"
#include "flow-queue.h"
#include "flow-util.h"
#include "flow-spare-pool.h"
#include "util-error.h"
#include "util-debug.h"
#include "util-print.h"
#include "util-validate.h"

typedef struct FlowSparePool {
    FlowQueuePrivate queue;
    struct FlowSparePool *next;
} FlowSparePool;

static uint32_t flow_spare_pool_flow_cnt = 0;
uint32_t flow_spare_pool_block_size = 100;
static FlowSparePool *flow_spare_pool = NULL;
static SCMutex flow_spare_pool_m = SCMUTEX_INITIALIZER;

uint32_t FlowSpareGetPoolSize(void)
{
    uint32_t size;
    SCMutexLock(&flow_spare_pool_m);
    size = flow_spare_pool_flow_cnt;
    SCMutexUnlock(&flow_spare_pool_m);
    return size;
}

static FlowSparePool *FlowSpareGetPool(void)
{
    FlowSparePool *p = SCCalloc(1, sizeof(*p));
    if (p == NULL)
        return NULL;
    return p;
}

static bool FlowSparePoolUpdateBlock(FlowSparePool *p)
{
    DEBUG_VALIDATE_BUG_ON(p == NULL);

    for (uint32_t i = p->queue.len; i < flow_spare_pool_block_size; i++)
    {
        Flow *f = FlowAlloc();
        if (f == NULL)
            return false;
        FlowQueuePrivateAppendFlow(&p->queue, f);
    }
    return true;
}

#ifdef FSP_VALIDATE
static void Validate(FlowSparePool *top, const uint32_t target)
{
    if (top == NULL) {
        assert(target == 0);
        return;
    }

    assert(top->queue.len >= 1);
    //if (top->next != NULL)
    //    assert(top->next->queue.len == flow_spare_pool_block_size);

    uint32_t cnt = 0;
    for (FlowSparePool *p = top; p != NULL; p = p->next)
    {
        assert(p->queue.len);
        cnt += p->queue.len;
    }
    assert(cnt == target);
}
#endif

void FlowSparePoolReturnFlow(Flow *f)
{
    SCMutexLock(&flow_spare_pool_m);
    if (flow_spare_pool == NULL) {
        flow_spare_pool = FlowSpareGetPool();
    }
    DEBUG_VALIDATE_BUG_ON(flow_spare_pool == NULL);

    /* if the top is full, get a new block */
    if (flow_spare_pool->queue.len >= flow_spare_pool_block_size) {
        FlowSparePool *p = FlowSpareGetPool();
        DEBUG_VALIDATE_BUG_ON(p == NULL);
        p->next = flow_spare_pool;
        flow_spare_pool = p;
    }
    /* add to the (possibly new) top */
    FlowQueuePrivateAppendFlow(&flow_spare_pool->queue, f);
    flow_spare_pool_flow_cnt++;

    SCMutexUnlock(&flow_spare_pool_m);
}

void FlowSparePoolReturnFlows(FlowQueuePrivate *fqp)
{
    FlowSparePool *p = FlowSpareGetPool();
    DEBUG_VALIDATE_BUG_ON(p == NULL);
    p->queue = *fqp;

    SCMutexLock(&flow_spare_pool_m);
    flow_spare_pool_flow_cnt += fqp->len;
    if (flow_spare_pool != NULL) {
        if (p->queue.len == flow_spare_pool_block_size) {
            /* full block insert */

            if (flow_spare_pool->queue.len < flow_spare_pool_block_size) {
                p->next = flow_spare_pool->next;
                flow_spare_pool->next = p;
                p = NULL;
            } else {
                p->next = flow_spare_pool;
                flow_spare_pool = p;
                p = NULL;
            }
        } else {
            /* incomplete block insert */

            if (p->queue.len + flow_spare_pool->queue.len <= flow_spare_pool_block_size) {
                FlowQueuePrivateAppendPrivate(&flow_spare_pool->queue, &p->queue);
                /* free 'p' outside of lock below */
            } else {
                // put smallest first
                if (p->queue.len < flow_spare_pool->queue.len) {
                    p->next = flow_spare_pool;
                    flow_spare_pool = p;
                } else {
                    p->next = flow_spare_pool->next;
                    flow_spare_pool->next = p;
                }
                p = NULL;
            }
        }
    } else {
        p->next = flow_spare_pool;
        flow_spare_pool = p;
        p = NULL;
    }
    SCMutexUnlock(&flow_spare_pool_m);

    FlowQueuePrivate empty = { NULL, NULL, 0 };
    *fqp = empty;

    if (p != NULL)
        SCFree(p);
}

FlowQueuePrivate FlowSpareGetFromPool(void)
{
    SCMutexLock(&flow_spare_pool_m);
    if (flow_spare_pool == NULL || flow_spare_pool_flow_cnt == 0) {
        SCMutexUnlock(&flow_spare_pool_m);
        FlowQueuePrivate empty = { NULL, NULL, 0 };
        return empty;
    }

    /* top if full or its the only block we have */
    if (flow_spare_pool->queue.len >= flow_spare_pool_block_size || flow_spare_pool->next == NULL) {
        FlowSparePool *p = flow_spare_pool;
        flow_spare_pool = p->next;
        DEBUG_VALIDATE_BUG_ON(flow_spare_pool_flow_cnt < p->queue.len);
        flow_spare_pool_flow_cnt -= p->queue.len;
#ifdef FSP_VALIDATE
        Validate(flow_spare_pool, flow_spare_pool_flow_cnt);
#endif
        SCMutexUnlock(&flow_spare_pool_m);

        FlowQueuePrivate ret = p->queue;
        SCFree(p);
        return ret;
    /* next should always be full if it exists */
    } else if (flow_spare_pool->next != NULL) {
        FlowSparePool *p = flow_spare_pool->next;
        flow_spare_pool->next = p->next;
        DEBUG_VALIDATE_BUG_ON(flow_spare_pool_flow_cnt < p->queue.len);
        flow_spare_pool_flow_cnt -= p->queue.len;
#ifdef FSP_VALIDATE
        Validate(flow_spare_pool, flow_spare_pool_flow_cnt);
#endif
        SCMutexUnlock(&flow_spare_pool_m);

        FlowQueuePrivate ret = p->queue;
        SCFree(p);
        return ret;
    }

    SCMutexUnlock(&flow_spare_pool_m);
    FlowQueuePrivate empty = { NULL, NULL, 0 };
    return empty;
}

void FlowSparePoolUpdate(uint32_t size)
{
    const int64_t todo = (int64_t)flow_config.prealloc - (int64_t)size;
    if (todo < 0) {
        uint32_t to_remove = (uint32_t)(todo * -1) / 10;
        while (to_remove) {
            if (to_remove < flow_spare_pool_block_size)
                return;

            FlowSparePool *p = NULL;
            SCMutexLock(&flow_spare_pool_m);
            p = flow_spare_pool;
            if (p != NULL) {
                flow_spare_pool = p->next;
                flow_spare_pool_flow_cnt -= p->queue.len;
                to_remove -= p->queue.len;
            }
            SCMutexUnlock(&flow_spare_pool_m);

            if (p != NULL) {
                Flow *f;
                while ((f = FlowQueuePrivateGetFromTop(&p->queue))) {
                    FlowFree(f);
                }
                SCFree(p);
            }
        }
    } else if (todo > 0) {
        FlowSparePool *head = NULL, *tail = NULL;

        uint32_t blocks = ((uint32_t)todo / flow_spare_pool_block_size) + 1;

        uint32_t flow_cnt = 0;
        for (uint32_t cnt = 0; cnt < blocks; cnt++) {
            FlowSparePool *p = FlowSpareGetPool();
            if (p == NULL) {
                break;
            }
            const bool ok = FlowSparePoolUpdateBlock(p);
            if (p->queue.len == 0) {
                SCFree(p);
                break;
            }
            flow_cnt += p->queue.len;

            /* prepend to list */
            p->next = head;
            head = p;
            if (tail == NULL)
                tail = p;
            if (!ok)
                break;
        }
        if (head) {
            SCMutexLock(&flow_spare_pool_m);
            if (flow_spare_pool == NULL) {
                flow_spare_pool = head;
            } else if (tail != NULL) {
                /* since these are 'full' buckets we don't put them
                 * at the top but right after as the top is likely not
                 * full. */
                tail->next = flow_spare_pool->next;
                flow_spare_pool->next = head;
            }

            flow_spare_pool_flow_cnt += flow_cnt;
#ifdef FSP_VALIDATE
            Validate(flow_spare_pool, flow_spare_pool_flow_cnt);
#endif
            SCMutexUnlock(&flow_spare_pool_m);
        }
    }
}

void FlowSparePoolInit(void)
{
    SCMutexLock(&flow_spare_pool_m);
    for (uint32_t cnt = 0; cnt < flow_config.prealloc; ) {
        FlowSparePool *p = FlowSpareGetPool();
        if (p == NULL) {
            FatalError("failed to initialize flow pool");
        }
        FlowSparePoolUpdateBlock(p);
        cnt += p->queue.len;

        /* prepend to list */
        p->next = flow_spare_pool;
        flow_spare_pool = p;
        flow_spare_pool_flow_cnt = cnt;
    }
    SCMutexUnlock(&flow_spare_pool_m);
}

void FlowSparePoolDestroy(void)
{
    SCMutexLock(&flow_spare_pool_m);
    for (FlowSparePool *p = flow_spare_pool; p != NULL; ) {
        uint32_t cnt = 0;
        Flow *f;
        while ((f = FlowQueuePrivateGetFromTop(&p->queue))) {
            FlowFree(f);
            cnt++;
        }
        flow_spare_pool_flow_cnt -= cnt;
        FlowSparePool *next = p->next;
        SCFree(p);
        p = next;
    }
    flow_spare_pool = NULL;
    SCMutexUnlock(&flow_spare_pool_m);
}