stellar-stellar/src/ip_reassemble/ip_reassemble.cpp

#include <stdlib.h>
#include <string.h>
#include <sys/queue.h>
#include <assert.h>

#include "packet.h"
#include "timestamp.h"
#include "crc32_hash.h"
#include "checksum.h"
#include "ipv4_utils.h"
#include "ipv6_utils.h"
#include "packet_helpers.h"
#include "ip_reassemble.h"

#define IPV4_KEYLEN 1
#define IPV6_KEYLEN 4
#define PRIME_VALUE 0xeaad8405
#define IP_FRAG_HASH_FNUM 2
#define IP_FRAG_TBL_POS(mgr, sig) ((mgr)->table + ((sig) & (mgr)->entry_mask))

#define KEY_TO_STR(key, str_str, dst_str)                                            \
    do                                                                               \
    {                                                                                \
        if ((key)->src_dst_len == IPV4_KEYLEN)                                       \
        {                                                                            \
            uint32_t src_addr = (key)->src_dst_addr[0] >> 32;                        \
            uint32_t dst_addr = (key)->src_dst_addr[0] & 0xffffffff;                 \
            inet_ntop(AF_INET, &src_addr, src_str, INET6_ADDRSTRLEN);                \
            inet_ntop(AF_INET, &dst_addr, dst_str, INET6_ADDRSTRLEN);                \
        }                                                                            \
        else                                                                         \
        {                                                                            \
            inet_ntop(AF_INET6, &(key)->src_dst_addr[0], src_str, INET6_ADDRSTRLEN); \
            inet_ntop(AF_INET6, &(key)->src_dst_addr[2], dst_str, INET6_ADDRSTRLEN); \
        }                                                                            \
    } while (0)

#define IP_REASSEMBLE_DEBUG1(desc, key, ...)                                            \
    do                                                                                  \
    {                                                                                   \
        char src_str[INET6_ADDRSTRLEN] = {0};                                           \
        char dst_str[INET6_ADDRSTRLEN] = {0};                                           \
        KEY_TO_STR(key, src_str, dst_str);                                              \
        IP_REASSEMBLE_DEBUG("%s (%s->%s 0x%0x)", desc, src_str, dst_str, (key)->ip_id); \
    } while (0)

#define IP_REASSEMBLE_ERROR1(desc, key, ...)                                            \
    do                                                                                  \
    {                                                                                   \
        char src_str[INET6_ADDRSTRLEN] = {0};                                           \
        char dst_str[INET6_ADDRSTRLEN] = {0};                                           \
        KEY_TO_STR(key, src_str, dst_str);                                              \
        IP_REASSEMBLE_ERROR("%s (%s->%s 0x%0x)", desc, src_str, dst_str, (key)->ip_id); \
    } while (0)

/******************************************************************************
 * Structs
 ******************************************************************************/

enum
{
    IP_LAST_FRAG_IDX,
    IP_FIRST_FRAG_IDX,
    IP_MIN_FRAG_NUM,
    IP_MAX_FRAG_NUM = 8,
};

struct ip_frag_hdr
{
    void *hdr_data; // need be freed
    uint16_t hdr_len;

    uint16_t l3_offset;
    uint16_t l3_len;

    uint8_t next_proto;
};

struct ip_frag_pkt
{
    void *data; // need be freed
    uint16_t len;

    uint16_t offset;
};

struct ip_flow_key
{
    uint64_t src_dst_addr[4]; // src and dst address (only first 8 bytes used for IPv4)
    uint32_t src_dst_len;
    uint32_t ip_id; // ipv4: identification is uint16_t; ipv6: identification is uint32_t
    uint8_t proto;
};

struct ip_flow
{
    struct
    {
        struct ip_flow *tqe_next;
        struct ip_flow **tqe_prev;
    } lru;
    struct ip_flow_key key;
    struct ip_frag_hdr hdr;
    uint64_t create_time;
    uint32_t expected_total_size;
    uint32_t received_frag_size;
    uint32_t next_fill_idx;
    struct ip_frag_pkt frags[IP_MAX_FRAG_NUM]; // first two entries in the frags[] array are for the last and first fragments.
};

struct ip_reassemble_manager
{
    // config
    bool enable;
    uint32_t timeout;
    uint32_t bucket_entries;
    uint32_t bucket_num;

    // runtime
    uint32_t entry_used;
    uint32_t entry_total;
    uint32_t entry_mask;

    // stats
    struct ip_reassemble_stat stat;

    // hash table
    struct
    {
        struct ip_flow *tqh_first;
        struct ip_flow **tqh_last;
    } lru;
    struct ip_flow *last;
    struct ip_flow *table; // array of ip_flow
};

/******************************************************************************
 * utils
 ******************************************************************************/

#define ip_reassemble_stat_inc(mgr, filed, key) \
    {                                           \
        if ((key)->src_dst_len == IPV4_KEYLEN)  \
        {                                       \
            (mgr)->stat.ip4_flow_##filed++;     \
        }                                       \
        else                                    \
        {                                       \
            (mgr)->stat.ip6_flow_##filed++;     \
        }                                       \
    }

static inline void *memdup(const void *src, size_t len)
{
    if (src == NULL || len == 0)
    {
        return NULL;
    }

    void *dst = malloc(len);
    if (dst == NULL)
    {
        return NULL;
    }
    return memcpy(dst, src, len);
}

static inline uint32_t combine32ms1b(uint32_t x)
{
    x |= x >> 1;
    x |= x >> 2;
    x |= x >> 4;
    x |= x >> 8;
    x |= x >> 16;

    return x;
}

static inline uint32_t align32pow2(uint32_t x)
{
    x--;
    x = combine32ms1b(x);

    return x + 1;
}

static inline int is_power_of_2(uint32_t n)
{
    return n && !(n & (n - 1));
}

static inline int ip_reassemble_check_config(const struct ip_reassemble_config *config)
{
    if (config == NULL)
    {
        IP_REASSEMBLE_DEBUG("invalid config");
        return -1;
    }

    if (config->enable)
    {
        if (config->timeout == 0)
        {
            IP_REASSEMBLE_DEBUG("invalid timeout");
            return -1;
        }

        if (config->bucket_entries == 0 || is_power_of_2(config->bucket_entries) == 0)
        {
            IP_REASSEMBLE_DEBUG("invalid bucket entries, must be power of 2");
            return -1;
        }

        if (config->bucket_num == 0)
        {
            IP_REASSEMBLE_DEBUG("invalid bucket num");
            return -1;
        }
    }

    return 0;
}

/******************************************************************************
 * ip flow key
 ******************************************************************************/

static inline void ipv4_flow_key_hash(const struct ip_flow_key *key, uint32_t *value1, uint32_t *value2)
{
    uint32_t v = 0;
    const uint32_t *p = (const uint32_t *)&key->src_dst_addr;

    v = crc32_hash_4byte(p[0], PRIME_VALUE);
    v = crc32_hash_4byte(p[1], v);
    v = crc32_hash_4byte(key->ip_id, v);

    *value1 = v;
    *value2 = (v << 7) + (v >> 14);
}

static inline void ipv6_flow_key_hash(const struct ip_flow_key *key, uint32_t *value1, uint32_t *value2)
{
    uint32_t v = 0;
    const uint32_t *p = (const uint32_t *)&key->src_dst_addr;

    v = crc32_hash_4byte(p[0], PRIME_VALUE);
    v = crc32_hash_4byte(p[1], v);
    v = crc32_hash_4byte(p[2], v);
    v = crc32_hash_4byte(p[3], v);
    v = crc32_hash_4byte(p[4], v);
    v = crc32_hash_4byte(p[5], v);
    v = crc32_hash_4byte(p[6], v);
    v = crc32_hash_4byte(p[7], v);
    v = crc32_hash_4byte(key->ip_id, v);

    *value1 = v;
    *value2 = (v << 7) + (v >> 14);
}

static inline uint64_t ip_flow_key_cmp(const struct ip_flow_key *key1, const struct ip_flow_key *key2)
{
    if (key1->ip_id != key2->ip_id)
    {
        return 1;
    }

    if (key1->src_dst_len != key2->src_dst_len)
    {
        return 1;
    }

    for (uint32_t i = 0; i < key1->src_dst_len; i++)
    {
        if (key1->src_dst_addr[i] != key2->src_dst_addr[i])
        {
            return 1;
        }
    }

    return 0;
}

static inline int ip_flow_key_is_empty(const struct ip_flow_key *key)
{
    return (key->src_dst_len == 0);
}

static inline void ip_flow_key_zero(struct ip_flow_key *key)
{
    key->src_dst_addr[0] = 0;
    key->src_dst_addr[1] = 0;
    key->src_dst_addr[2] = 0;
    key->src_dst_addr[3] = 0;

    key->src_dst_len = 0;
    key->ip_id = 0;
    key->proto = 0;
}

/******************************************************************************
 * ip frag hdr
 ******************************************************************************/

static inline void ip_frag_hdr_init(struct ip_frag_hdr *hdr, const struct packet *pkt)
{
    struct layer_record *layer = pkt->frag_layer;

    if (layer->type == LAYER_TYPE_IPV6)
    {
        struct ip6_frag *frag_ext = ipv6_hdr_get_frag_ext((const struct ip6_hdr *)layer->hdr_ptr);
        hdr->next_proto = frag_ext->ip6f_nxt;
    }
    else
    {
        hdr->next_proto = ipv4_hdr_get_proto((const struct ip *)layer->hdr_ptr);
    }

    hdr->l3_offset = layer->hdr_offset;
    hdr->l3_len = layer->hdr_len;

    hdr->hdr_len = layer->hdr_offset + layer->hdr_len;
    hdr->hdr_data = memdup(pkt->data_ptr, hdr->hdr_len);
}

static inline void ip_frag_hdr_free(struct ip_frag_hdr *hdr)
{
    hdr->next_proto = 0;

    hdr->l3_offset = 0;
    hdr->l3_len = 0;

    hdr->hdr_len = 0;
    if (hdr->hdr_data != NULL)
    {
        free(hdr->hdr_data);
        hdr->hdr_data = NULL;
    }
}

/******************************************************************************
 * ip frag pkt
 ******************************************************************************/

static inline void ip_frag_pkt_init(struct ip_frag_pkt *frag, void *data, uint16_t len, uint16_t offset)
{
    frag->data = memdup(data, len);
    frag->len = len;
    frag->offset = offset;
}

static inline void ip_frag_pkt_free(struct ip_frag_pkt *frag)
{
    if (frag)
    {
        if (frag->data)
        {
            free(frag->data);
            frag->data = NULL;
        }

        frag->len = 0;
        frag->offset = 0;
    }
}

/******************************************************************************
 * ip flow
 ******************************************************************************/

static inline void ip_flow_init(struct ip_flow *flow, const struct ip_flow_key *key)
{
    static const struct ip_frag_pkt zero_frag = {
        .data = NULL,
        .len = 0,
        .offset = 0,
    };

    flow->lru.tqe_next = NULL;
    flow->lru.tqe_prev = NULL;
    flow->key = *key;
    flow->create_time = timestamp_get_sec();
    flow->expected_total_size = UINT32_MAX;
    flow->received_frag_size = 0;
    flow->next_fill_idx = IP_MIN_FRAG_NUM;
    flow->frags[IP_LAST_FRAG_IDX] = zero_frag;
    flow->frags[IP_FIRST_FRAG_IDX] = zero_frag;
}

static inline void ip_flow_free(struct ip_flow *flow)
{
    for (uint32_t i = 0; i < IP_MAX_FRAG_NUM; i++)
    {
        struct ip_frag_pkt *frag = &flow->frags[i];
        ip_frag_pkt_free(frag);
    }
    ip_flow_key_zero(&flow->key);
    ip_frag_hdr_free(&flow->hdr);
}

static inline int ip_flow_is_ready(struct ip_flow *flow)
{
    return (flow->received_frag_size == flow->expected_total_size && flow->frags[IP_FIRST_FRAG_IDX].data != NULL);
}

// return  0 : success
// return -1 : failed
static inline int ip_flow_update(struct ip_reassemble_manager *mgr,
                                 struct ip_flow *flow, const struct packet *pkt,
                                 char *frag_data, uint16_t frag_len, uint16_t frag_offset, bool more_frags)
{
    uint32_t idx;
    struct ip_frag_pkt *frag_pkt;
    /*
     * Internet Protocol, Version 6 (IPv6) Specification
     *
     * https://datatracker.ietf.org/doc/html/rfc8200#section-4.5
     *
     * It should be noted that fragments may be duplicated in the
     * network.  Instead of treating these exact duplicate fragments
     * as overlapping fragments, an implementation may choose to
     * detect this case and drop exact duplicate fragments while
     * keeping the other fragments belonging to the same packet.
     */
    if (frag_offset == 0)
    {
        if (flow->frags[IP_FIRST_FRAG_IDX].data != NULL)
        {
            IP_REASSEMBLE_DEBUG1("duplicate first fragment bypass", &flow->key);
            ip_reassemble_stat_inc(mgr, bypass_dup_fist_frag, &flow->key);
            return 0;
        }
        idx = IP_FIRST_FRAG_IDX;
        ip_frag_hdr_init(&flow->hdr, pkt);
    }
    else if (more_frags == 0)
    {
        if (flow->frags[IP_LAST_FRAG_IDX].data != NULL)
        {
            IP_REASSEMBLE_DEBUG1("duplicate last fragment bypass", &flow->key);
            ip_reassemble_stat_inc(mgr, bypass_dup_last_frag, &flow->key);
            return 0;
        }
        idx = IP_LAST_FRAG_IDX;
        flow->expected_total_size = frag_offset + frag_len;
    }
    else
    {
        if (flow->next_fill_idx >= IP_MAX_FRAG_NUM)
        {
            IP_REASSEMBLE_ERROR1("max number of fragment exceeded", &flow->key);
            ip_reassemble_stat_inc(mgr, fail_many_frag, &flow->key);
            return -1;
        }
        idx = flow->next_fill_idx;
        flow->next_fill_idx++;
    }

    flow->received_frag_size += frag_len;
    frag_pkt = &flow->frags[idx];
    ip_frag_pkt_init(frag_pkt, frag_data, frag_len, frag_offset);

    return 0;
}

/******************************************************************************
 * ip reassemble manager add/del/reuse/find/update flow
 ******************************************************************************/

static inline void ip_reassemble_manager_add_flow(struct ip_reassemble_manager *mgr, struct ip_flow *flow)
{
    ip_reassemble_stat_inc(mgr, add, &flow->key);
    TAILQ_INSERT_TAIL(&mgr->lru, flow, lru);
    mgr->entry_used++;
}

static inline void ip_reassemble_manager_del_flow(struct ip_reassemble_manager *mgr, struct ip_flow *flow)
{
    ip_reassemble_stat_inc(mgr, del, &flow->key);
    TAILQ_REMOVE(&mgr->lru, flow, lru);
    mgr->entry_used--;
}

static inline void ip_reassemble_manager_reuse_flow(struct ip_reassemble_manager *mgr, struct ip_flow *flow, const struct ip_flow_key *key)
{
    ip_reassemble_manager_del_flow(mgr, flow);
    ip_flow_free(flow);
    ip_flow_init(flow, key);
    ip_reassemble_manager_add_flow(mgr, flow);
}

/*
 * if return NULL, then *free and *expired are valid
 *      free   : the first empty entry in the bucket
 *      expired: the first timed-out entry in the bucket
 */
static struct ip_flow *ip_reassemble_manager_find_flow(struct ip_reassemble_manager *mgr, const struct ip_flow_key *key, struct ip_flow **free, struct ip_flow **expired)
{
    ip_reassemble_stat_inc(mgr, find, key);

    if (mgr->last != NULL && ip_flow_key_cmp(key, &mgr->last->key) == 0)
    {
        return mgr->last;
    }

    uint32_t sig1 = 0;
    uint32_t sig2 = 0;
    if (key->src_dst_len == IPV4_KEYLEN)
    {
        ipv4_flow_key_hash(key, &sig1, &sig2);
    }
    else
    {
        ipv6_flow_key_hash(key, &sig1, &sig2);
    }

    // get the bucket by hash
    struct ip_flow *p1 = IP_FRAG_TBL_POS(mgr, sig1);
    struct ip_flow *p2 = IP_FRAG_TBL_POS(mgr, sig2);

    // search in the bucket
    struct ip_flow *old = NULL;
    struct ip_flow *empty = NULL;
    uint64_t timeout = mgr->timeout;
    uint32_t assoc = mgr->bucket_entries;
    uint64_t tms = timestamp_get_sec();
    for (uint32_t i = 0; i != assoc; i++)
    {
        if (ip_flow_key_cmp(key, &p1[i].key) == 0)
        {
            *free = NULL;
            *expired = NULL;
            return p1 + i;
        }
        else if (ip_flow_key_is_empty(&p1[i].key))
        {
            empty = (empty == NULL) ? (p1 + i) : empty;
        }
        else if (timeout + p1[i].create_time <= tms)
        {
            old = (old == NULL) ? (p1 + i) : old;
        }

        if (ip_flow_key_cmp(key, &p2[i].key) == 0)
        {
            *free = NULL;
            *expired = NULL;
            return p2 + i;
        }
        else if (ip_flow_key_is_empty(&p2[i].key))
        {
            empty = (empty == NULL) ? (p2 + i) : empty;
        }
        else if (timeout + p2[i].create_time <= tms)
        {
            old = (old == NULL) ? (p2 + i) : old;
        }
    }

    *free = empty;
    *expired = old;
    return NULL;
}

static struct ip_flow *ip_reassemble_manager_update_flow(struct ip_reassemble_manager *mgr, const struct ip_flow_key *key)
{
    struct ip_flow *flow = NULL;
    struct ip_flow *free = NULL;
    struct ip_flow *expired = NULL;

    flow = ip_reassemble_manager_find_flow(mgr, key, &free, &expired);
    if (flow == NULL)
    {
        if (expired)
        {
            IP_REASSEMBLE_DEBUG1("add ip flow success: reuse expired entry", key);
            ip_reassemble_manager_reuse_flow(mgr, expired, key);
            ip_reassemble_stat_inc(mgr, timeout, key);

            mgr->last = expired;
            return expired;
        }

        if (free)
        {
            IP_REASSEMBLE_DEBUG1("add ip flow success: use free entry", key);
            ip_flow_init(free, key);
            ip_reassemble_manager_add_flow(mgr, free);

            mgr->last = free;
            return free;
        }

        // no space
        IP_REASSEMBLE_ERROR1("add ip flow failed: bucket full", key);
        ip_reassemble_stat_inc(mgr, fail_no_space, key);
        return NULL;
    }
    else
    {
        // expired
        if (mgr->timeout + flow->create_time <= timestamp_get_sec())
        {
            IP_REASSEMBLE_DEBUG1("add ip flow success: reuse expired entry", key);
            ip_reassemble_manager_reuse_flow(mgr, flow, key);
            ip_reassemble_stat_inc(mgr, timeout, key);

            mgr->last = flow;
            return flow;
        }
        // not expired
        else
        {
            IP_REASSEMBLE_DEBUG1("find ip flow success: not expire", key);

            mgr->last = flow;
            return flow;
        }
    }
}

/******************************************************************************
 * frag reassemble
 ******************************************************************************/

static struct packet *ip_frag_reassemble(struct ip_reassemble_manager *mgr, struct ip_flow *flow)
{
    struct ip_frag_pkt *first = &flow->frags[IP_FIRST_FRAG_IDX];
    struct ip_frag_pkt *last = &flow->frags[IP_LAST_FRAG_IDX];
    struct ip_frag_pkt *temp = NULL;

    uint32_t loop = 0;
    uint16_t last_offset = last->offset;

    struct ip *ip4_hdr = NULL;
    struct ip6_hdr *ip6_hdr = NULL;

    // calculate the length of the reassembled packet
    uint32_t packet_len = flow->expected_total_size + flow->hdr.hdr_len;
    char *buff = (char *)calloc(1, packet_len + sizeof(struct packet));
    if (buff == NULL)
    {
        IP_REASSEMBLE_ERROR("unable to allocate memory");
        return NULL;
    }

    char *ptr = buff + sizeof(struct packet);
    char *end = ptr + packet_len;

    // copy last frag
    if (last->len > end - ptr)
    {
        IP_REASSEMBLE_ERROR1("last frag length not match expected reassembled length", &flow->key);
        goto error_out_invalid_length;
    }
    end -= last->len;
    memcpy(end, last->data, last->len);

    while (first->len != last_offset)
    {
        /*
         * https://datatracker.ietf.org/doc/html/rfc791
         *
         * In the case that two or more fragments contain the same data
         * either identically or through a partial overlap, this procedure
         * will use the more recently arrived copy in the data buffer and
         * datagram delivered.
         */
        for (uint32_t i = flow->next_fill_idx - 1; i >= IP_MIN_FRAG_NUM; i--)
        {
            temp = &flow->frags[i];
            if (temp->offset + temp->len == last_offset)
            {
                if (temp->len > end - ptr)
                {
                    IP_REASSEMBLE_ERROR1("middle frag length not match expected reassembled length", &flow->key);
                    goto error_out_invalid_length;
                }

                end -= temp->len;
                memcpy(end, temp->data, temp->len);
                last_offset = temp->offset;
                break;
            }
        }

        if (loop > flow->next_fill_idx - IP_MIN_FRAG_NUM)
        {
            IP_REASSEMBLE_ERROR1("overlap appear during frag reassemble", &flow->key);
            goto error_out_overlap;
        }

        loop++;
    }

    // copy fist fragment data
    if (first->len > end - ptr)
    {
        IP_REASSEMBLE_ERROR1("first frag length not match expected reassembled length", &flow->key);
        goto error_out_invalid_length;
    }
    end -= first->len;
    memcpy(end, first->data, first->len);

    // copy frag hdr
    if (flow->hdr.hdr_len > end - ptr)
    {
        IP_REASSEMBLE_ERROR1("packet header length not match expected reassembled length", &flow->key);
        goto error_out_invalid_length;
    }
    end -= flow->hdr.hdr_len;
    memcpy(end, flow->hdr.hdr_data, flow->hdr.hdr_len);

    // assert
    assert(ptr == end);

    if (flow->key.src_dst_len == IPV4_KEYLEN)
    {
        // update ip total length & ip checksum
        ip4_hdr = (struct ip *)(ptr + flow->hdr.l3_offset);
        ipv4_hdr_set_total_len(ip4_hdr, packet_len - flow->hdr.l3_offset); // update total length
        ipv4_hdr_set_mf_flag(ip4_hdr, false);                              // update more fragment flag
        ipv4_hdr_set_frag_offset(ip4_hdr, 0);                              // update fragment offset
        ip4_hdr->ip_sum = 0;                                               // update checksum
        ip4_hdr->ip_sum = checksum((char *)ip4_hdr, flow->hdr.l3_len);
    }
    else
    {
        // update ipv6 payload length & next header
        ip6_hdr = (struct ip6_hdr *)(ptr + flow->hdr.l3_offset);
        ipv6_hdr_set_payload_len(ip6_hdr, flow->expected_total_size); // update payload length
        ipv6_hdr_set_next_header(ip6_hdr, flow->hdr.next_proto);      // update next header
    }

    // create a new packet
    packet_parse((struct packet *)buff, buff + sizeof(struct packet), packet_len);
    return (struct packet *)buff;

error_out_invalid_length:
    ip_reassemble_stat_inc(mgr, fail_invalid_length, &flow->key);
    free(buff);
    return NULL;

error_out_overlap:
    ip_reassemble_stat_inc(mgr, fail_overlap, &flow->key);
    free(buff);
    return NULL;
}

/******************************************************************************
 * Public API
 ******************************************************************************/

struct ip_reassemble_manager *ip_reassemble_manager_create(const struct ip_reassemble_config *config)
{
    if (ip_reassemble_check_config(config) != 0)
    {
        return NULL;
    }

    struct ip_reassemble_manager *mgr = (struct ip_reassemble_manager *)calloc(1, sizeof(struct ip_reassemble_manager));
    if (mgr == NULL)
    {
        IP_REASSEMBLE_ERROR("unable to allocate memory");
        return NULL;
    }
    mgr->enable = config->enable;
    mgr->timeout = config->timeout;
    mgr->bucket_entries = config->bucket_entries;
    mgr->bucket_num = config->bucket_num;

    if (!mgr->enable)
    {
        return mgr;
    }

    uint64_t entry_total = align32pow2(mgr->bucket_num) * mgr->bucket_entries * IP_FRAG_HASH_FNUM;
    if (entry_total > UINT32_MAX)
    {
        IP_REASSEMBLE_ERROR("bucket_num * bucket_entries is too large");
        free(mgr);
        return NULL;
    }

    mgr->entry_total = (uint32_t)entry_total;
    mgr->entry_mask = (mgr->entry_total - 1) & ~(mgr->bucket_entries - 1);
    mgr->table = (struct ip_flow *)calloc(mgr->entry_total, sizeof(struct ip_flow));
    if (mgr->table == NULL)
    {
        IP_REASSEMBLE_ERROR("unable to allocate memory");
        free(mgr);
        return NULL;
    }

    TAILQ_INIT(&(mgr->lru));

    return mgr;
}

void ip_reassemble_manager_destory(struct ip_reassemble_manager *mgr)
{
    if (mgr)
    {
        if (mgr->table)
        {
            for (uint32_t i = 0; i < mgr->entry_total; i++)
            {
                ip_flow_free(mgr->table + i);
            }

            free(mgr->table);
            mgr->table = NULL;
        }
        free(mgr);
        mgr = NULL;
    }
}

void ip_reassemble_manager_expire(struct ip_reassemble_manager *mgr)
{
    struct ip_flow *flow = NULL;
    uint64_t curr_ts = timestamp_get_sec();
    uint64_t timeout = mgr->timeout;

    TAILQ_FOREACH(flow, &mgr->lru, lru)
    if (timeout + flow->create_time <= curr_ts)
    {
        IP_REASSEMBLE_DEBUG1("expire ip flow: discarding old fragmented packets", &flow->key);
        ip_reassemble_manager_del_flow(mgr, flow);
        ip_reassemble_stat_inc(mgr, timeout, &flow->key);
        ip_flow_free(flow);
    }
}

void ip_reassemble_manager_print_stat(struct ip_reassemble_manager *mgr)
{
    if (mgr)
    {
        IP_REASSEMBLE_DEBUG("ip4_flow_find                 : %lu", mgr->stat.ip4_flow_find);
        IP_REASSEMBLE_DEBUG("ip4_flow_add                  : %lu", mgr->stat.ip4_flow_add);
        IP_REASSEMBLE_DEBUG("ip4_flow_del                  : %lu", mgr->stat.ip4_flow_del);
        IP_REASSEMBLE_DEBUG("ip4_flow_timeout              : %lu", mgr->stat.ip4_flow_timeout);
        IP_REASSEMBLE_DEBUG("ip4_flow_fail_no_space        : %lu", mgr->stat.ip4_flow_fail_no_space);
        IP_REASSEMBLE_DEBUG("ip4_flow_fail_overlap         : %lu", mgr->stat.ip4_flow_fail_overlap);
        IP_REASSEMBLE_DEBUG("ip4_flow_fail_many_frag       : %lu", mgr->stat.ip4_flow_fail_many_frag);
        IP_REASSEMBLE_DEBUG("ip4_flow_fail_invalid_length  : %lu", mgr->stat.ip4_flow_fail_invalid_length);
        IP_REASSEMBLE_DEBUG("ip4_flow_bypass_dup_fist_frag : %lu", mgr->stat.ip4_flow_bypass_dup_fist_frag);
        IP_REASSEMBLE_DEBUG("ip4_flow_bypass_dup_last_frag : %lu", mgr->stat.ip4_flow_bypass_dup_last_frag);

        IP_REASSEMBLE_DEBUG("ip6_flow_find                 : %lu", mgr->stat.ip6_flow_find);
        IP_REASSEMBLE_DEBUG("ip6_flow_add                  : %lu", mgr->stat.ip6_flow_add);
        IP_REASSEMBLE_DEBUG("ip6_flow_del                  : %lu", mgr->stat.ip6_flow_del);
        IP_REASSEMBLE_DEBUG("ip6_flow_timeout              : %lu", mgr->stat.ip6_flow_timeout);
        IP_REASSEMBLE_DEBUG("ip6_flow_fail_no_space        : %lu", mgr->stat.ip6_flow_fail_no_space);
        IP_REASSEMBLE_DEBUG("ip6_flow_fail_overlap         : %lu", mgr->stat.ip6_flow_fail_overlap);
        IP_REASSEMBLE_DEBUG("ip6_flow_fail_many_frag       : %lu", mgr->stat.ip6_flow_fail_many_frag);
        IP_REASSEMBLE_DEBUG("ip6_flow_fail_invalid_length  : %lu", mgr->stat.ip6_flow_fail_invalid_length);
        IP_REASSEMBLE_DEBUG("ip6_flow_bypass_dup_fist_frag : %lu", mgr->stat.ip6_flow_bypass_dup_fist_frag);
        IP_REASSEMBLE_DEBUG("ip6_flow_bypass_dup_last_frag : %lu", mgr->stat.ip6_flow_bypass_dup_last_frag);
    }
}

struct ip_reassemble_stat *ip_reassemble_manager_get_stat(struct ip_reassemble_manager *mgr)
{
    if (mgr)
    {
        return &(mgr->stat);
    }
    else
    {
        return NULL;
    }
}

/*
 * Returns the reassembled packet, or NULL if the packet is not reassembled
 * The returned packet should be freed by calling the packet_free() function
 */

struct packet *ip_reassemble_packet(struct ip_reassemble_manager *mgr, const struct packet *pkt)
{
    struct packet *pkt1;
    struct packet *pkt2;

    if (!mgr->enable)
    {
        return NULL;
    }

    const struct layer_record *layer = pkt->frag_layer;
    if (layer == NULL)
    {
        return NULL;
    }

    if (layer->type == LAYER_TYPE_IPV4)
    {
        pkt1 = ipv4_reassemble_packet(mgr, pkt);
        if (pkt1 && pkt1->frag_layer)
        {
            pkt2 = ip_reassemble_packet(mgr, pkt1);
            packet_free(pkt1);
            return pkt2;
        }

        return pkt1;
    }
    else if (layer->type == LAYER_TYPE_IPV6)
    {
        pkt1 = ipv6_reassemble_packet(mgr, pkt);
        if (pkt1 && pkt1->frag_layer)
        {
            pkt2 = ip_reassemble_packet(mgr, pkt1);
            packet_free(pkt1);
            return pkt2;
        }

        return pkt1;
    }
    else
    {
        return NULL;
    }
}

struct packet *ipv4_reassemble_packet(struct ip_reassemble_manager *mgr, const struct packet *pkt)
{
    const struct layer_record *layer = pkt->frag_layer;
    const struct ip *hdr = (const struct ip *)layer->hdr_ptr;
    uint16_t frag_len = ipv4_hdr_get_total_len(hdr) - ipv4_hdr_get_hdr_len(hdr);
    if (frag_len > layer->pld_len)
    {
        IP_REASSEMBLE_ERROR("unexpected header length, ip id: %lu", ipv4_hdr_get_ipid(hdr));
        return NULL;
    }

    struct ip_flow_key key = {0};
    uint64_t src_addr = hdr->ip_src.s_addr;
    uint64_t dst_addr = hdr->ip_dst.s_addr;
    key.src_dst_addr[0] = src_addr << 32 | dst_addr;
    key.src_dst_len = IPV4_KEYLEN;
    key.ip_id = ipv4_hdr_get_ipid(hdr);
    key.proto = ipv4_hdr_get_proto(hdr);

    struct ip_flow *flow = ip_reassemble_manager_update_flow(mgr, &key);
    if (flow == NULL)
    {
        return NULL;
    }

    char *frag_data = (char *)layer->pld_ptr;
    bool more_frags = ipv4_hdr_get_mf_flag(hdr);
    uint16_t frag_offset = ipv4_hdr_get_frag_offset(hdr);
    if (ip_flow_update(mgr, flow, pkt, frag_data, frag_len, frag_offset, more_frags) != 0)
    {
        ip_reassemble_manager_del_flow(mgr, flow);
        ip_flow_free(flow);
        return NULL;
    }

    if (!ip_flow_is_ready(flow))
    {
        return NULL;
    }

    struct packet *new_pkt = ip_frag_reassemble(mgr, flow);
    ip_reassemble_manager_del_flow(mgr, flow);
    ip_flow_free(flow);

    return new_pkt;
}

/*
 * https://datatracker.ietf.org/doc/html/rfc8200#section-4.5
 *
 * Note: unlike IPv4, fragmentation in IPv6 is performed only by source nodes,
 * not by routers along a packet's delivery path
 */

/*
 * original packet:
 * +-----------------+-----------------+--------+--------+-//-+--------+
 * |  Per-Fragment   |Ext & Upper-Layer|  first | second |    |  last  |
 * |    Headers      |    Headers      |fragment|fragment|....|fragment|
 * +-----------------+-----------------+--------+--------+-//-+--------+
 *
 * fragment packets:
 * +-----------------+--------+-------------------+----------+
 * |  Per-Fragment   |Fragment| Ext & Upper-Layer |  first   |
 * |    Headers      | Header |   Headers         | fragment |
 * +-----------------+--------+-------------------+----------+
 *
 * +-----------------+--------+----------+
 * |  Per-Fragment   |Fragment|  second  |
 * |    Headers      | Header | fragment |
 * +-----------------+--------+----------+
 *                          o
 *                          o
 *                          o
 * +-----------------+--------+----------+
 * |  Per-Fragment   |Fragment|   last   |
 * |    Headers      | Header | fragment |
 * +-----------------+--------+----------+
 *
 * reassembled packet:
 * +-----------------+-----------------+--------+--------+-//-+--------+
 * |  Per-Fragment   |Ext & Upper-Layer|  first | second |    |  last  |
 * |    Headers      |    Headers      |fragment|fragment|....|fragment|
 * +-----------------+-----------------+--------+--------+-//-+--------+
 */

struct packet *ipv6_reassemble_packet(struct ip_reassemble_manager *mgr, const struct packet *pkt)
{
    const struct layer_record *layer = pkt->frag_layer;
    const struct ip6_hdr *hdr = (const struct ip6_hdr *)layer->hdr_ptr;
    const struct ip6_frag *frag_hdr = ipv6_hdr_get_frag_ext(hdr);
    if (frag_hdr == NULL)
    {
        return NULL;
    }

    char *frag_data = (char *)frag_hdr + sizeof(struct ip6_frag);
    uint16_t frag_len = ipv6_hdr_get_payload_len(hdr) - sizeof(struct ip6_frag);
    if (frag_data + frag_len > pkt->data_ptr + pkt->data_len)
    {
        IP_REASSEMBLE_ERROR("unexpected header length, frag id: %lu", ipv6_frag_get_ident(frag_hdr));
        return NULL;
    }

    struct ip_flow_key key = {0};
    memcpy(&key.src_dst_addr[0], hdr->ip6_src.s6_addr, 16);
    memcpy(&key.src_dst_addr[2], hdr->ip6_dst.s6_addr, 16);
    key.src_dst_len = IPV6_KEYLEN;
    key.ip_id = ipv6_frag_get_ident(frag_hdr);
    key.proto = 0; // only first fragment has the upper layer protocol

    struct ip_flow *flow = ip_reassemble_manager_update_flow(mgr, &key);
    if (flow == NULL)
    {
        return NULL;
    }

    bool more_frags = ipv6_frag_get_more(frag_hdr);
    uint16_t frag_offset = ipv6_frag_get_offset(frag_hdr);
    if (ip_flow_update(mgr, flow, pkt, frag_data, frag_len, frag_offset, more_frags) != 0)
    {
        ip_reassemble_manager_del_flow(mgr, flow);
        ip_flow_free(flow);
        return NULL;
    }

    if (!ip_flow_is_ready(flow))
    {
        return NULL;
    }

    struct packet *new_pkt = ip_frag_reassemble(mgr, flow);
    ip_reassemble_manager_del_flow(mgr, flow);
    ip_flow_free(flow);

    return new_pkt;
}