#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 IP_FRAG_HASH_FNUM 2
#define PRIME_VALUE 0xeaad8405

#define IP_FRAG_TBL_POS(mgr, sig) ((mgr)->table + ((sig) & (mgr)->entry_mask))

#define IPV4_KEYLEN 1
#define IPV6_KEYLEN 4
#define IPv6_KEY_BYTES(key) (key)[0], (key)[1], (key)[2], (key)[3]

#define IP_REASSEMBLE_DEBUG_ARG1(desc, key, ...)                                                                             \
    do                                                                                                                       \
    {                                                                                                                        \
        if ((key)->src_dst_len == IPV4_KEYLEN)                                                                               \
        {                                                                                                                    \
            IP_REASSEMBLE_DEBUG(desc, "key <%lu, %#x>", (key)->src_dst_addr[0], (key)->ip_id);                               \
        }                                                                                                                    \
        else                                                                                                                 \
        {                                                                                                                    \
            IP_REASSEMBLE_DEBUG(desc, "key <%08lu%08lu%08lu%08lu, %#x>", IPv6_KEY_BYTES((key)->src_dst_addr), (key)->ip_id); \
        }                                                                                                                    \
    } while (0)

#define IP_REASSEMBLE_ERROR_ARG1(desc, key, ...)                                                                             \
    do                                                                                                                       \
    {                                                                                                                        \
        if ((key)->src_dst_len == IPV4_KEYLEN)                                                                               \
        {                                                                                                                    \
            IP_REASSEMBLE_ERROR(desc, "key <%lu, %#x>", (key)->src_dst_addr[0], (key)->ip_id);                               \
        }                                                                                                                    \
        else                                                                                                                 \
        {                                                                                                                    \
            IP_REASSEMBLE_ERROR(desc, "key <%08lu%08lu%08lu%08lu, %#x>", IPv6_KEY_BYTES((key)->src_dst_addr), (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
{
    uint16_t l3_offset;
    uint16_t l3_len;
    uint16_t hdr_len;
    uint8_t next_proto;
    void *hdr_data; // need be freed
};

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

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_stat
{
    uint64_t find_num;
    uint64_t add_num;
    uint64_t del_num;
    uint64_t fail_total;
    uint64_t fail_nospace;
};

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

/******************************************************************************
 * Private API
 ******************************************************************************/

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

/******************************************************************************
 * struct 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;
}

/******************************************************************************
 * struct ip_frag_hdr
 ******************************************************************************/

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

/******************************************************************************
 * struct 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 = {
        .offset = 0,
        .len = 0,
        .data = NULL,
    };

    flow->lru.tqe_next = NULL;
    flow->lru.tqe_prev = NULL;
    flow->key = *key;
    flow->create_time = timestamp_get_msec();
    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];
        frag->offset = 0;
        frag->len = 0;
        if (frag->data != NULL)
        {
            free(frag->data);
            frag->data = NULL;
        }
    }
    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_flow *flow, const struct packet *pkt,
                                 char *data, uint16_t len, uint16_t offset, bool more_frags)
{
    uint32_t idx;

    if (offset == 0)
    {
        if (flow->frags[IP_FIRST_FRAG_IDX].data != NULL)
        {
            IP_REASSEMBLE_ERROR_ARG1("duplicate first fragment encountered: ", &flow->key);
            return -1;
        }
        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_ERROR_ARG1("duplicate last fragment encountered: ", &flow->key);
            return -1;
        }
        idx = IP_LAST_FRAG_IDX;
        flow->expected_total_size = offset + len;
    }
    else
    {
        if (flow->next_fill_idx >= IP_MAX_FRAG_NUM)
        {
            IP_REASSEMBLE_ERROR_ARG1("max number of fragment exceeded: ", &flow->key);
            return -1;
        }
        idx = flow->next_fill_idx;
        flow->next_fill_idx++;
    }

    flow->received_frag_size += len;
    flow->frags[idx].data = memdup(data, len);
    flow->frags[idx].offset = offset;
    flow->frags[idx].len = len;

    return 0;
}

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

static inline void ip_flow_table_add(struct ip_reassemble_manager *mgr, struct ip_flow *flow)
{
    TAILQ_INSERT_TAIL(&mgr->lru, flow, lru);
    mgr->entry_used++;
    mgr->stat.add_num++;
}

static inline void ip_flow_table_del(struct ip_reassemble_manager *mgr, struct ip_flow *flow)
{
    TAILQ_REMOVE(&mgr->lru, flow, lru);
    mgr->entry_used--;
    mgr->stat.del_num++;
}

static inline void ip_flow_table_reuse(struct ip_reassemble_manager *mgr, struct ip_flow *flow)
{
    ip_flow_table_del(mgr, flow);
    ip_flow_table_add(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_flow_table_find(struct ip_reassemble_manager *mgr, const struct ip_flow_key *key, struct ip_flow **free, struct ip_flow **expired)
{
    mgr->stat.find_num++;

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

    uint32_t sig1 = 0;
    uint32_t sig2 = 0;
    uint64_t timeout = mgr->timeout;
    uint32_t assoc = mgr->bucket_entries;
    uint64_t tms = timestamp_get_msec();

    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;
    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_flow_table_update(struct ip_reassemble_manager *mgr, const struct ip_flow_key *key)
{
    struct ip_flow *free = NULL;
    struct ip_flow *expired = NULL;
    struct ip_flow *flow = NULL;
    uint64_t tms = timestamp_get_msec();

    flow = ip_flow_table_find(mgr, key, &free, &expired);
    if (flow == NULL)
    {
        if (expired)
        {
            IP_REASSEMBLE_DEBUG_ARG1("ip flow new (use expired iterm): ", key);
            ip_flow_free(expired);
            ip_flow_init(expired, key);
            ip_flow_table_reuse(mgr, expired);

            mgr->last = expired;
            return expired;
        }

        if (free)
        {
            IP_REASSEMBLE_DEBUG_ARG1("ip flow new (use free iterm): ", key);
            ip_flow_init(free, key);
            ip_flow_table_add(mgr, free);

            mgr->last = free;
            return free;
        }

        // no space
        IP_REASSEMBLE_DEBUG_ARG1("bucket full discarding new fragmented packets: ", key);
        mgr->stat.fail_nospace++;
        return NULL;
    }
    else
    {
        // expired
        if (mgr->timeout + flow->create_time < tms)
        {
            IP_REASSEMBLE_DEBUG_ARG1("ip flow find, but expired: ", key);
            ip_flow_free(flow);
            ip_flow_init(flow, key);
            ip_flow_table_reuse(mgr, flow);

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

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

    TAILQ_FOREACH(flow, &mgr->lru, lru)
    if (timeout + flow->create_time < curr_ts)
    {
        IP_REASSEMBLE_DEBUG_ARG1("time expires discarding old fragmented packets: ", &flow->key);
        ip_flow_free(flow);
        ip_flow_table_del(mgr, flow);
    }
}

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

static struct packet *ip_frag_reassemble(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;

    // calculate the length of the reassembled packet
    uint32_t buff_len = flow->expected_total_size + flow->hdr.hdr_len;
    char *buff = (char *)calloc(1, buff_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 + buff_len;

    // copy last frag
    if (last->len > end - ptr)
    {
        IP_REASSEMBLE_ERROR_ARG1("last packet length is greater than the expected reassembled length: ", &flow->key);
        free(buff);
        return NULL;
    }
    end -= last->len;
    memcpy(end, last->data, last->len);

    uint32_t loop = 0;
    uint16_t last_offset = last->offset;
    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--)
        {
            if (i == IP_FIRST_FRAG_IDX || i == IP_LAST_FRAG_IDX)
            {
                continue;
            }

            temp = &flow->frags[i];
            if (temp->offset + temp->len == last_offset)
            {
                if (temp->len > end - ptr)
                {
                    IP_REASSEMBLE_ERROR_ARG1("middle fragment packet length doesn't match last fragment packet offset: ", &flow->key);
                    free(buff);
                    return NULL;
                }

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

        if (loop > flow->next_fill_idx)
        {
            IP_REASSEMBLE_ERROR_ARG1("holes appear during frag reassemble: ", &flow->key);
            free(buff);
            return NULL;
        }

        loop++;
    }

    // copy fist fragment data
    if (first->len > end - ptr)
    {
        IP_REASSEMBLE_ERROR_ARG1("last packet length is greater than the expected reassembled length: ", &flow->key);
        free(buff);
        return NULL;
    }
    end -= first->len;
    memcpy(end, first->data, first->len);

    // copy frag hdr
    if (flow->hdr.hdr_len > end - ptr)
    {
        IP_REASSEMBLE_ERROR_ARG1("hdr length is greater than the expected reassembled length: ", &flow->key);
        free(buff);
        return NULL;
    }
    end -= flow->hdr.hdr_len;
    memcpy(end, flow->hdr.hdr_data, flow->hdr.hdr_len);
    assert(ptr == end);

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

    // create a new packet
    struct packet *new_pkt = (struct packet *)buff;
    packet_parse(new_pkt, buff + sizeof(struct packet), buff_len);

    return new_pkt;
}

/******************************************************************************
 * 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_stat(struct ip_reassemble_manager *mgr)
{
    // TODO
}

/*
 * 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, pkt);
            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, pkt);
            packet_free(pkt1);
            return pkt2;
        }

        return pkt1;
    }
    else
    {
        return NULL;
    }
}

struct packet *ipv4_reassemble_packet(struct ip_reassemble_manager *mgr, const struct packet *pkt)
{
    struct layer_record *layer = pkt->frag_layer;
    struct ip *hdr = (struct ip *)layer->hdr_ptr;
    char *data = (char *)layer->pld_ptr;
    uint16_t len = ipv4_hdr_get_total_len(hdr) - ipv4_hdr_get_hdr_len(hdr);
    if (len > layer->pld_len)
    {
        IP_REASSEMBLE_ERROR("unexpected header length on fragmented packet, 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_flow_table_update(mgr, &key);
    if (flow == NULL)
    {
        return NULL;
    }

    bool more_frags = ipv4_hdr_get_mf_flag(hdr);
    uint16_t offset = ipv4_hdr_get_frag_offset(hdr);
    if (ip_flow_update(flow, pkt, data, len, offset, more_frags) != 0)
    {
        ip_flow_free(flow);
        ip_flow_table_del(mgr, flow);
        return NULL;
    }

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

    struct packet *new_pkt = ip_frag_reassemble(flow);
    ip_flow_free(flow);
    ip_flow_table_del(mgr, 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)
{
    struct layer_record *layer = pkt->frag_layer;
    const struct ip6_hdr *hdr = (const struct ip6_hdr *)layer->hdr_ptr;
    struct ip6_frag *frag_ext = ipv6_hdr_get_frag_ext(hdr);
    if (frag_ext == NULL)
    {
        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 = frag_ext->ip6f_ident;
    key.proto = 0; // only first fragment has the upper layer protocol

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

    char *data = (char *)layer->hdr_ptr + sizeof(struct ip6_hdr) + sizeof(struct ip6_frag);
    uint16_t len = ipv6_hdr_get_payload_len(hdr) - sizeof(struct ip6_frag);

    if (data + len > pkt->data_ptr + pkt->data_len)
    {
        IP_REASSEMBLE_ERROR("unexpected header length on fragmented packet, frag_id: %lu", frag_ext->ip6f_ident);
        return NULL;
    }

    bool more_frags = (frag_ext->ip6f_offlg & IP6F_MORE_FRAG);
    uint16_t offset = ntohs(frag_ext->ip6f_offlg & IP6F_OFF_MASK);
    if (ip_flow_update(flow, pkt, data, len, offset, more_frags) != 0)
    {
        ip_flow_free(flow);
        ip_flow_table_del(mgr, flow);
        return NULL;
    }

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

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

    return new_pkt;
}