#include <tfe_proxy.h>
#include <tfe_utils.h>
#include <tfe_types.h>
#include <traffic_mirror.h>

struct traffic_mirror_rebuild
{
    struct tfe_stream_addr * c_s_addr;
    struct tfe_stream_addr * s_c_addr;

    struct profile_table_ex_data * target;
    struct traffic_mirror_ethdev * ethdev;

    uint32_t c_seq;
    uint32_t s_seq;
};

/* The definition of vlan_hdr and tcp_hdr is from DPDK 17.05 */
struct vlan_hdr
{
    uint16_t vlan_tci; /**< Priority (3) + CFI (1) + Identifier Code (12) */
    uint16_t eth_proto;/**< Ethernet type of encapsulated frame. */
} __attribute__((__packed__));

struct tcp_hdr {
    uint16_t src_port;  /**< TCP source port. */
    uint16_t dst_port;  /**< TCP destination port. */
    uint32_t sent_seq;  /**< TX data sequence number. */
    uint32_t recv_ack;  /**< RX data acknowledgement sequence number. */
    uint8_t  data_off;  /**< Data offset. */
    uint8_t  tcp_flags; /**< TCP flags */
    uint16_t rx_win;    /**< RX flow control window. */
    uint16_t cksum;     /**< TCP checksum. */
    uint16_t tcp_urp;   /**< TCP urgent pointer, if any. */
} __attribute__((__packed__));

#define TCP_URG_FLAG 0x20
#define TCP_ACK_FLAG 0x10
#define TCP_PSH_FLAG 0x08
#define TCP_RST_FLAG 0x04
#define TCP_SYN_FLAG 0x02
#define TCP_FIN_FLAG 0x01
#define TCP_FLAG_ALL 0x3F

static inline uint32_t __rte_raw_cksum(const void *buf, size_t len, uint32_t sum)
{
    /* workaround gcc strict-aliasing warning */
    uintptr_t ptr = (uintptr_t) buf;
    typedef uint16_t __attribute__((__may_alias__)) u16_p;
    const u16_p * u16 = (const u16_p *) ptr;

    while (len >= (sizeof(*u16) * 4))
    {
        sum += u16[0];
        sum += u16[1];
        sum += u16[2];
        sum += u16[3];
        len -= sizeof(*u16) * 4;
        u16 += 4;
    }
    while (len >= sizeof(*u16))
    {
        sum += *u16;
        len -= sizeof(*u16);
        u16 += 1;
    }

    /* if length is in odd bytes */
    if (len == 1)
        sum += *((const uint8_t *) u16);

    return sum;
}

static inline uint16_t __rte_raw_cksum_reduce(uint32_t sum)
{
    sum = ((sum & 0xffff0000) >> 16) + (sum & 0xffff);
    sum = ((sum & 0xffff0000) >> 16) + (sum & 0xffff);
    return (uint16_t) sum;
}

static inline uint16_t rte_raw_cksum(const void *buf, size_t len)
{
    uint32_t sum = __rte_raw_cksum(buf, len, 0);
    return __rte_raw_cksum_reduce(sum);
}

static inline uint16_t __ipv4_cksum(const struct iphdr * ipv4_hdr)
{
    uint16_t cksum = rte_raw_cksum(ipv4_hdr, sizeof(struct iphdr));
    return (cksum == 0xffff) ? cksum : ~cksum;
}

static inline uint16_t __tcpudp_cksum_by_stream_addr_v4(const struct tfe_stream_addr * addr,
    const void *l4_hdr, unsigned int l4_len)
{
    struct ipv4_psd_header
    {
        uint32_t src_addr; /* IP address of source host. */
        uint32_t dst_addr; /* IP address of destination host. */
        uint8_t zero;     /* zero. */
        uint8_t proto;    /* L4 protocol type. */
        uint16_t len;      /* L4 length. */
    } psd_hdr;

    psd_hdr.src_addr = addr->tuple4_v4->saddr.s_addr;
    psd_hdr.dst_addr = addr->tuple4_v4->daddr.s_addr;
    psd_hdr.zero = 0;
    psd_hdr.proto = IPPROTO_TCP;
    psd_hdr.len = htons(l4_len);

    uint32_t cksum = rte_raw_cksum(l4_hdr, l4_len);
    cksum += rte_raw_cksum(&psd_hdr, sizeof(psd_hdr));

    cksum = ((cksum & 0xffff0000) >> 16) + (cksum & 0xffff);
    cksum = (~cksum) & 0xffff;
    if (cksum == 0)
        cksum = 0xffff;

    return cksum;
}

static int tcp_header_construct(unsigned char *buf, unsigned short sp,
    unsigned short dp, unsigned int seq, unsigned int ack,
    unsigned char flags, unsigned short win, unsigned short urg)
{
    struct tcp_hdr * tcp_hdr = (struct tcp_hdr *) buf;
    assert(tcp_hdr != NULL);

    tcp_hdr->src_port = sp;
    tcp_hdr->dst_port = dp;
    tcp_hdr->sent_seq = htonl(seq);
    tcp_hdr->recv_ack = htonl(ack);
    tcp_hdr->data_off = 0x50;
    tcp_hdr->tcp_flags = flags;
    tcp_hdr->rx_win = htons(win);
    tcp_hdr->cksum = 0;
    tcp_hdr->tcp_urp = 0;
    return sizeof(struct tcp_hdr);
}

static int tcp_header_construct_by_stream_addr(struct tfe_stream_addr * addr, unsigned char *buf,
    unsigned int seq, unsigned int ack, unsigned char flags, unsigned short win, unsigned short urg,
    unsigned int payload_len)
{
    unsigned short sport;
    unsigned short dport;
    unsigned short cksum;
    unsigned int tcphdr_len;

    if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V4)
    {
        sport = addr->tuple4_v4->source;
        dport = addr->tuple4_v4->dest;
    }
    else if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V6)
    {
        sport = addr->tuple4_v6->source;
        dport = addr->tuple4_v6->dest;
    }
    else
    {
        assert(0);
    }

    tcphdr_len = tcp_header_construct(buf, sport, dport, seq, ack, flags, win, urg);
    if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V4)
    {
        cksum = __tcpudp_cksum_by_stream_addr_v4(addr, (void *)buf, tcphdr_len + payload_len);
    }
    else if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V6)
    {
        assert(0);
    }
    else
    {
        assert(0);
    }

    struct tcp_hdr * tcp_hdr = (struct tcp_hdr *)buf;
    tcp_hdr->cksum = cksum;
    return tcphdr_len;
}

static int ipv4_header_construct(unsigned char *buf, unsigned short carry_layer_len,
    unsigned char tos, unsigned short id, unsigned short frag, unsigned char ttl,
    unsigned char protocol, unsigned int src, unsigned int dst)
{
    struct iphdr * ip_hdr = (struct iphdr *) buf;
    ip_hdr->version = 4;
    ip_hdr->ihl = 5;
    ip_hdr->tos = tos;
    ip_hdr->tot_len = htons(sizeof(struct iphdr) + carry_layer_len);
    ip_hdr->id = htons(id);
    ip_hdr->frag_off = 0;
    ip_hdr->ttl = ttl;
    ip_hdr->protocol = protocol;
    ip_hdr->check = 0;
    ip_hdr->saddr = src;
    ip_hdr->daddr = dst;

    /* checksum */
    ip_hdr->check = __ipv4_cksum(ip_hdr);
    return sizeof(struct iphdr);
}

static int ip_header_construct_by_stream_addr(struct tfe_stream_addr * addr,
    unsigned char *buf, unsigned short carry_layer_len, unsigned char tos,
    unsigned short id, unsigned short frag, unsigned char ttl, unsigned char protocol)
{
    if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V4)
    {
        return ipv4_header_construct(buf, carry_layer_len, tos, id,
            frag, ttl, protocol, addr->tuple4_v4->saddr.s_addr, addr->tuple4_v4->daddr.s_addr);
    }
    else if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V6)
    {
        /* TODO: IPv6 */
        assert(0);
    }
    else
    {
        assert(0);
    }

    assert(0);
    return -1;
}

static void ether_header_construct(unsigned char *buf, unsigned char *dst,
    unsigned char *src, unsigned short type)
{
    struct ethhdr * eth_hdr = (struct ethhdr *) buf;
    memcpy(eth_hdr->h_dest, dst, ETHER_ADDR_LEN);
    memcpy(eth_hdr->h_source, src, ETHER_ADDR_LEN);
    eth_hdr->h_proto = htons(type);
}

static void vlan_tag_construct(unsigned char *buf, unsigned short tci, unsigned short type)
{
    struct vlan_hdr * vlan_hdr = (struct vlan_hdr *)buf;
    vlan_hdr->vlan_tci = htons(tci);
    vlan_hdr->eth_proto = htons(type);
}

static void l2_send_to_target(struct traffic_mirror_ethdev * ethdev,
    unsigned char * snd_buffer, unsigned int l3_data_offset, unsigned int l3_data_len,
    struct ether_addr * target_addr, unsigned int vlan_tci, unsigned l3_protocol)
{
    assert(l3_data_offset >= (sizeof(struct ethhdr) + sizeof(struct vlan_hdr)));
    unsigned int header_offset = l3_data_offset;
    unsigned int header_len = 0;

    /* need to construct vlan header */
    if (vlan_tci > 0)
    {
        header_offset -= sizeof(struct vlan_hdr);
        header_len += sizeof(struct vlan_hdr);
        vlan_tag_construct(snd_buffer + header_offset, vlan_tci, l3_protocol);
    }

    unsigned int eth_protocol = vlan_tci > 0 ? ETH_P_8021Q : l3_protocol;
    header_offset -= sizeof(struct ethhdr);
    header_len += sizeof(struct ethhdr);

    ether_header_construct(snd_buffer + header_offset, (unsigned char *)target_addr->ether_addr_octet,
        (unsigned char *)ethdev->local_ether_addr, eth_protocol);

    traffic_mirror_ethdev_inject(ethdev, (char *)snd_buffer + header_offset, header_len + l3_data_len);
}

static void l2_send_to_target_group(struct traffic_mirror_ethdev * ethdev, struct profile_table_ex_data * t_group,
    unsigned char * snd_buffer, unsigned int l3_data_offset, unsigned int l3_data_len, unsigned l3_protocol)
{
    for(unsigned int i = 0; i < t_group->nr_targets; i++)
    {
        l2_send_to_target(ethdev, snd_buffer, l3_data_offset, l3_data_len,
            &t_group->ether_addrs[i], t_group->vlans[i], l3_protocol);
    }
}

static void tcp_segment_send_to_target_group(struct tfe_stream_addr * addr, struct traffic_mirror_ethdev * ethdev,
    struct profile_table_ex_data * t_group, const char * payload, unsigned int payload_len,
        unsigned int seq, unsigned int ack, char flags)
{
    char pkt[ETHER_MAX_LEN];

    unsigned sz_pkt_prepend = sizeof(struct ethhdr) + sizeof(struct vlan_hdr) +
        sizeof(struct iphdr) + sizeof(struct tcphdr);

    unsigned header_len = 0;

    assert(sizeof(pkt) - sz_pkt_prepend >= payload_len);
    memcpy(pkt + sz_pkt_prepend, payload, payload_len);

    sz_pkt_prepend -= sizeof(struct tcp_hdr);
    header_len += tcp_header_construct_by_stream_addr(addr, (unsigned char *)pkt + sz_pkt_prepend,
        seq, ack, flags, 0xffff, 0, payload_len);

    sz_pkt_prepend -= sizeof(struct iphdr);
    header_len += ip_header_construct_by_stream_addr(addr, (unsigned char *)pkt + sz_pkt_prepend,
        header_len + payload_len, 0, 0x1000, 0, 128, IPPROTO_TCP);

    l2_send_to_target_group(ethdev, t_group, (unsigned char *)pkt, sz_pkt_prepend,
        header_len + payload_len, ETHERTYPE_IP);
}

static void tcp_send_to_target_group(struct tfe_stream_addr * addr, struct traffic_mirror_ethdev * ethdev,
    struct profile_table_ex_data * t_group, const char * payload,
        unsigned int payload_len, unsigned int seq, unsigned int ack, char flags)
{
    unsigned int payload_offset = 0;
    unsigned mss = ethdev->mtu - (sizeof(struct iphdr) + sizeof(struct tcphdr));

    /* Handshake or farewall */
    if (payload == NULL || payload_len == 0)
    {
        tcp_segment_send_to_target_group(addr, ethdev, t_group, NULL, 0, seq, ack, flags);
        return;
    }

    while(payload_offset < payload_len)
    {
        unsigned int payload_sz_seg = MIN(payload_len - payload_offset, mss);
        const char * payload_ptr_seg = payload + payload_offset;

        tcp_segment_send_to_target_group(addr, ethdev, t_group, payload_ptr_seg, payload_sz_seg, seq, ack, flags);
        seq += payload_sz_seg;
        payload_offset += payload_sz_seg;
    }
}

struct traffic_mirror_rebuild * traffic_mirror_rebuild_create(struct tfe_stream_addr * addr,
    struct profile_table_ex_data * target, struct traffic_mirror_ethdev * ethdev)
{
    struct traffic_mirror_rebuild * instance = ALLOC(struct traffic_mirror_rebuild, 1);
    instance->c_s_addr = addr;
    instance->s_c_addr = tfe_stream_addr_reverse(addr);
    instance->target = target;
    instance->ethdev = ethdev;

    /* the c_seq, s_seq, c_ipid, s_ipid is random
     * TODO: use the fast random algorithm like Linux TCP/IP stack */
    instance->c_seq = random();
    instance->s_seq = random();
    return instance;
}

void traffic_mirror_rebuild_destroy(struct traffic_mirror_rebuild * instance)
{
    tfe_stream_addr_free(instance->s_c_addr);
    free(instance);
}

void traffic_mirror_rebuild_handshake(struct traffic_mirror_rebuild * instance)
{
    tcp_send_to_target_group(instance->c_s_addr, instance->ethdev, instance->target,
        NULL, 0, instance->c_seq, 0, TCP_SYN_FLAG);

    tcp_send_to_target_group(instance->s_c_addr, instance->ethdev, instance->target,
        NULL, 0, instance->s_seq, instance->c_seq + 1, TCP_SYN_FLAG | TCP_ACK_FLAG);

    tcp_send_to_target_group(instance->c_s_addr, instance->ethdev, instance->target,
        NULL, 0, instance->c_seq + 1, instance->s_seq + 1, TCP_ACK_FLAG);

    instance->s_seq++;
    instance->c_seq++;
}

void traffic_mirror_rebuild_data(struct traffic_mirror_rebuild * instance,
    const char * data, unsigned int datalen, enum tfe_conn_dir dir)
{
    if (data == NULL || datalen == 0)
    {
        return;
    }

    if (dir == CONN_DIR_DOWNSTREAM)
    {
        tcp_send_to_target_group(instance->c_s_addr, instance->ethdev, instance->target,
            data, datalen, instance->c_seq, instance->s_seq, TCP_ACK_FLAG);
        instance->c_seq += datalen;
    }
    else
    {
        tcp_send_to_target_group(instance->s_c_addr, instance->ethdev, instance->target,
            data, datalen, instance->s_seq, instance->c_seq, TCP_ACK_FLAG);
        instance->s_seq += datalen;
    }
}

void traffic_mirror_rebuild_farewell(struct traffic_mirror_rebuild * instance)
{
    /* C->S FIN */
    tcp_send_to_target_group(instance->c_s_addr, instance->ethdev, instance->target,
        NULL, 0, instance->c_seq, instance->s_seq, TCP_FIN_FLAG | TCP_ACK_FLAG);

    /* C->S FIN, ACK */
    tcp_send_to_target_group(instance->s_c_addr, instance->ethdev, instance->target,
        NULL, 0, instance->s_seq, instance->c_seq + 1, TCP_ACK_FLAG);

    instance->c_seq += 1;

    /* S->C FIN */
    tcp_send_to_target_group(instance->s_c_addr, instance->ethdev, instance->target,
        NULL, 0, instance->s_seq, instance->c_seq, TCP_FIN_FLAG | TCP_ACK_FLAG);

    /* C->S FIN, ACK */
    tcp_send_to_target_group(instance->c_s_addr, instance->ethdev, instance->target,
        NULL, 0, instance->c_seq, instance->s_seq + 1, TCP_ACK_FLAG);

    instance->s_seq += 1;
}