#pragma once

#include <rte_mbuf.h>
#include <rte_mempool.h>
#include <rte_spinlock.h>

/* 全局asan开关 */
extern unsigned int g_in_protect_mode;

/* asan_poison和asan_unpoison非线程安全，
 * 设置一个全局锁，保证不会同时执行asan_posion和asan_unpoison操作 */
extern rte_spinlock_t g_in_protect_lock;

/* TODO: 这些定义应通过包含GCC的头文件来实现，不应通过这种方式 */
#if defined(__SANITIZE_ADDRESS__)
void __asan_poison_memory_region(void const volatile * addr, size_t size);
void __asan_unpoison_memory_region(void const volatile * addr, size_t size);
#endif

#if defined(__SANITIZE_ADDRESS__)
#define MR_ASAN_POISON_MEMORY_REGION(addr, size)                                                                       \
    do                                                                                                                 \
    {                                                                                                                  \
        if (g_in_protect_mode)                                                                                         \
        {                                                                                                              \
            rte_spinlock_lock(&g_in_protect_lock);                                                                     \
            __asan_poison_memory_region((addr), (size));                                                               \
            rte_spinlock_unlock(&g_in_protect_lock);                                                                   \
        }                                                                                                              \
    } while (0)

#define MR_ASAN_UNPOISON_MEMORY_REGION(addr, size)                                                                     \
    do                                                                                                                 \
    {                                                                                                                  \
        if (g_in_protect_mode)                                                                                         \
        {                                                                                                              \
            rte_spinlock_lock(&g_in_protect_lock);                                                                     \
            __asan_unpoison_memory_region((addr), (size));                                                             \
            rte_spinlock_unlock(&g_in_protect_lock);                                                                   \
        }                                                                                                              \
    } while (0)

#else
#define MR_ASAN_POISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
#define MR_ASAN_UNPOISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
#endif

static inline void PROTECT_rte_mempool_unpoison(struct rte_mempool * pool)
{
    MR_ASAN_UNPOISON_MEMORY_REGION(pool, sizeof(struct rte_mempool));
    MR_ASAN_UNPOISON_MEMORY_REGION(pool, RTE_MEMPOOL_HEADER_SIZE(pool, pool->cache_size));
}

static inline void PROTECT_rte_mbuf_unpoison_meta(struct rte_mbuf * mbuf)
{
    MR_ASAN_UNPOISON_MEMORY_REGION(mbuf, sizeof(struct rte_mbuf));
    MR_ASAN_UNPOISON_MEMORY_REGION((const char *)mbuf + sizeof(struct rte_mbuf), mbuf->priv_size);
}

static inline void PROTECT_rte_mbuf_poison_meta(struct rte_mbuf * mbuf)
{
    MR_ASAN_POISON_MEMORY_REGION((const char *)mbuf + sizeof(struct rte_mbuf), mbuf->priv_size);
    MR_ASAN_POISON_MEMORY_REGION(mbuf, sizeof(struct rte_mbuf));
}

static inline void PROTECT_rte_mbuf_poison_buff(struct rte_mbuf * mbuf)
{
    MR_ASAN_POISON_MEMORY_REGION(mbuf->buf_addr, mbuf->buf_len);
}

static inline void PROTECT_rte_mbuf_unpoison_buff(struct rte_mbuf * mbuf)
{
    MR_ASAN_UNPOISON_MEMORY_REGION(mbuf->buf_addr, mbuf->buf_len);
}

static inline void PROTECT_rte_mbuf_poison(struct rte_mbuf * mbuf)
{
    PROTECT_rte_mbuf_poison_buff(mbuf);
    PROTECT_rte_mbuf_poison_meta(mbuf);
}

static inline void PROTECT_rte_mbuf_unpoison(struct rte_mbuf * mbuf)
{
    PROTECT_rte_mbuf_unpoison_meta(mbuf);
    PROTECT_rte_mbuf_unpoison_buff(mbuf);
}

static inline void PROTECT_rte_mbuf_poison_bulk(struct rte_mbuf * mbufs[], int nr_mbufs)
{
    for (unsigned int i = 0; i < nr_mbufs; i++)
        PROTECT_rte_mbuf_poison(mbufs[i]);
}

static inline void PROTECT_rte_mbuf_unpoison_bulk(struct rte_mbuf * mbufs[], int nr_mbufs)
{
    for (unsigned int i = 0; i < nr_mbufs; i++)
        PROTECT_rte_mbuf_unpoison(mbufs[i]);
}

static inline void PROTECT_rte_mbuf_unpoison_data(struct rte_mbuf * mbuf)
{
    /* Poison all buff area, because some area may be unpoison before. */
    PROTECT_rte_mbuf_poison_buff(mbuf);

    /* Unpoison the area we need to access */
    const char * data_ptr = rte_pktmbuf_mtod(mbuf, const char *);
    size_t data_len = rte_pktmbuf_data_len(mbuf);
    MR_ASAN_UNPOISON_MEMORY_REGION(data_ptr, data_len);
}

/* 从DPDK17.05中移植rte_pktmbuf_alloc
 * 该函数在申请mbuf的同时还要修改mbuf中的元数据，此时由于还没有拿到mbuf的指针，
 * 也就无法去保护，函数功能无法实现。现移植，在拿到mbuf的指针后立即去保护，继续后面的初始化流程 */
static inline struct rte_mbuf * __PROTECT_PORT_rte_mbuf_raw_alloc(struct rte_mempool * mp)
{
    struct rte_mbuf * m;
    void * mb = NULL;

    if (rte_mempool_get(mp, &mb) < 0)
        return NULL;

    m = (struct rte_mbuf *)mb;
    PROTECT_rte_mbuf_unpoison_meta(m);
    MBUF_RAW_ALLOC_CHECK(m);
    return m;
}

static inline struct rte_mbuf * __PROTECT_PORT_rte_pktmbuf_alloc(struct rte_mempool * mp)
{
    struct rte_mbuf * m;
    if ((m = __PROTECT_PORT_rte_mbuf_raw_alloc(mp)) != NULL)
        rte_pktmbuf_reset(m);
    return m;
}

/* 移植完毕，以后DPDK版本升级时，要格外注意以上的实现是否发生了变化 */

static inline struct rte_mbuf * PROTECT_rte_pktmbuf_alloc(struct rte_mempool * mp)
{
#ifndef __SANITIZE_ADDRESS__
    return rte_pktmbuf_alloc(mp);
#else
    struct rte_mbuf * m = __PROTECT_PORT_rte_pktmbuf_alloc(mp);
    PROTECT_rte_mbuf_poison_meta(m);
    return m;
#endif
}

#if 0
static inline void PROTECT_rte_pktmbuf_free(struct rte_mbuf * mbuf)
{
    PROTECT_rte_mbuf_unpoison_meta(mbuf);
    rte_pktmbuf_free(mbuf);
    /* TODO: 应重新对释放的mbuf区域进行保护，需要在释放前计算mbuf的指针，长度，再执行保护 */
}

static inline int PROTECT_rte_pktmbuf_alloc_bulk(struct rte_mempool * pool, struct rte_mbuf ** mbufs, unsigned count)
{
#ifndef __SANITIZE_ADDRESS__
    return rte_pktmbuf_alloc_bulk(pool, mbufs, count);
#else

    unsigned int alloc_count;
    for (alloc_count = 0; alloc_count < count; alloc_count++)
    {
        struct rte_mbuf * m = PROTECT_rte_pktmbuf_alloc(pool);
        if (unlikely(m == NULL))
            goto errout;
        mbufs[alloc_count] = m;
    }

    return 0;

errout:
    for (unsigned i = 0; i < alloc_count; i++)
    {
        PROTECT_rte_pktmbuf_free(mbufs[i]);
        mbufs[i] = NULL;
    }

    return -1;
#endif
}
#endif

static inline char * PROTECT_rte_pktmbuf_append(struct rte_mbuf * m, uint16_t len)
{
    PROTECT_rte_mbuf_unpoison_meta(m);
    char * ret = rte_pktmbuf_append(m, len);
    PROTECT_rte_mbuf_unpoison_data(m);
    return ret;
}

static inline char * PROTECT_rte_pktmbuf_prepend(struct rte_mbuf * m, uint16_t len)
{
    PROTECT_rte_mbuf_unpoison_meta(m);
    char * ret = rte_pktmbuf_prepend(m, len);
    PROTECT_rte_mbuf_unpoison_data(m);
    return ret;
}

static inline char * PROTECT_rte_pktmbuf_adj(struct rte_mbuf * m, uint16_t len)
{
    PROTECT_rte_mbuf_unpoison_meta(m);
    char * ret = rte_pktmbuf_adj(m, len);
    PROTECT_rte_mbuf_unpoison_data(m);
    return ret;
}

static inline int PROTECT_rte_pktmbuf_trim(struct rte_mbuf * m, uint16_t len)
{
    PROTECT_rte_mbuf_unpoison_meta(m);
    int ret = rte_pktmbuf_trim(m, len);
    return ret;
}

static inline uint16_t PROTECT_rte_pktmbuf_data_len(struct rte_mbuf * m)
{
    PROTECT_rte_mbuf_unpoison_meta(m);
    uint16_t ret = rte_pktmbuf_data_len(m);
    return ret;
}

static inline uint32_t PROTECT_rte_pktmbuf_pkt_len(struct rte_mbuf * m)
{
    PROTECT_rte_mbuf_unpoison_meta(m);
    uint32_t ret = rte_pktmbuf_pkt_len(m);
    return ret;
}

static inline uint16_t PROTECT_rte_mbuf_refcnt_update(struct rte_mbuf * m, int16_t value)
{
    PROTECT_rte_mbuf_unpoison_meta(m);
    uint16_t ret = rte_mbuf_refcnt_update(m, value);
    return ret;
}

static inline uint16_t PROTECT_rte_mbuf_refcnt_read(const struct rte_mbuf * m)
{
    PROTECT_rte_mbuf_unpoison_meta((struct rte_mbuf *)m);
    uint16_t ret = rte_mbuf_refcnt_read(m);
    return ret;
}