#include <tfe_proxy.h>
#include <tfe_utils.h>
#include <tfe_types.h>
#include <traffic_mirror.h>
#include <marsio.h>
#include <netinet/ip.h>
#include <netinet/ip6.h>

struct traffic_mirror_rebuild
{
    struct tfe_stream_addr * c_s_addr;
    struct tfe_stream_addr * s_c_addr;
    struct ether_addr c_ether_addr;
    struct ether_addr s_ether_addr;

    struct traffic_mirror_rebuild_target * 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 & 0xffff0000U) >> 16U) + (sum & 0xffffU);
    sum = ((sum & 0xffff0000U) >> 16U) + (sum & 0xffffU);
    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 __tcp_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 & 0xffff0000U) >> 16U) + (cksum & 0xffffU);
    cksum = (~cksum) & 0xffffU;
    if (cksum == 0) cksum = 0xffff;
    return cksum;
}

static inline uint16_t __tcp_cksum_by_stream_addr_v6(const struct tfe_stream_addr * addr,
    const void * l4_hdr, unsigned int l4_len)
{
    uint32_t sum;
    struct
    {
        uint32_t len;   /* L4 length. */
        uint32_t proto; /* L4 protocol - top 3 bytes must be zero */
    } psd_hdr
    {
        .len = l4_len,
        .proto =(uint32_t)(IPPROTO_TCP << 24U)
    };

    sum = __rte_raw_cksum(&addr->tuple4_v6->saddr, sizeof(addr->tuple4_v6->saddr), 0);
    sum = __rte_raw_cksum(&addr->tuple4_v6->daddr, sizeof(addr->tuple4_v6->daddr), sum);
    sum = __rte_raw_cksum(&psd_hdr, sizeof(psd_hdr), sum);
    return __rte_raw_cksum_reduce(sum);
}

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 char * payload, unsigned int payload_len)
{
    unsigned short sport=0;
    unsigned short dport=0;
    unsigned short cksum=0;

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

    unsigned int tcphdr_len = tcp_header_construct(buf, sport, dport, seq, ack, flags, win, urg);
    memcpy(buf + tcphdr_len, payload, payload_len);

    if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V4)
    {
        cksum = __tcp_cksum_by_stream_addr_v4(addr, (void *) buf, tcphdr_len + payload_len);
    }
    else if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V6)
    {
        cksum = __tcp_cksum_by_stream_addr_v6(addr, (void *)buf, tcphdr_len + payload_len);
    }

    struct tcp_hdr * tcp_hdr = (struct tcp_hdr *)buf;
    tcp_hdr->cksum = cksum;
    return tcphdr_len + payload_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 ipv6_header_construct(unsigned char *buf, unsigned short carry_layer_len,
    unsigned char ttl, unsigned char protocol, struct in6_addr * src, struct in6_addr * dst)
{
    struct ip6_hdr * ip6_hdr = (struct ip6_hdr *) buf;

    ip6_hdr->ip6_flow = 0;
    ip6_hdr->ip6_plen = htons(carry_layer_len);
    ip6_hdr->ip6_nxt = protocol;
    ip6_hdr->ip6_hlim = 128;
    ip6_hdr->ip6_vfc &= 0x0F;
    ip6_hdr->ip6_vfc |= (6U << 4U);

    ip6_hdr->ip6_src = *src;
    ip6_hdr->ip6_dst = *dst;
    return sizeof(struct ip6_hdr);
}

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)
    {
        return ipv6_header_construct(buf, carry_layer_len, ttl, protocol,
            &addr->tuple4_v6->saddr, &addr->tuple4_v6->daddr);
    }

    assert(0);
    return -1;
}

static void ether_header_construct_helper(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 int ether_header_construct(struct traffic_mirror_ethdev * ethdev, char * buffer,
	struct ether_addr * source, struct ether_addr * dest, unsigned int vlan_tci, unsigned l3_protocol)
{
	unsigned int header_len = 0;
	unsigned int eth_protocol = vlan_tci > 0 ? ETH_P_8021Q : l3_protocol;

	ether_header_construct_helper((unsigned char *) buffer + header_len,
		(unsigned char *) dest, (unsigned char *) source, eth_protocol);

	header_len += sizeof(struct ethhdr);

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

	return header_len;
}

static void tcp_segment_send_to_target_group(struct tfe_stream_addr * addr, struct ether_addr * ether_addr_src,
	struct ether_addr * ether_addr_dst, struct traffic_mirror_ethdev * ethdev,
	struct traffic_mirror_rebuild_target * target, unsigned int tid, const char * payload, unsigned int payload_len,
	unsigned int seq, unsigned int ack, char flags)
{
    char * pkt_buffer = ethdev->fn_send_prepare(ethdev, tid);
    assert(pkt_buffer != NULL);

    unsigned int l3_protocol = 0;
    switch(addr->addrtype)
    {
        case TFE_ADDR_STREAM_TUPLE4_V4: { l3_protocol = ETH_P_IP;  break;}
        case TFE_ADDR_STREAM_TUPLE4_V6: { l3_protocol = ETH_P_IPV6; break;}
        default: assert(0);
    }

    unsigned int pkt_len = 0;

    /* Ethernet and VLAN header */
    if ((target->rewrite_as_target_mac &&  target->rewrite_as_target_vlan) ||
        (target->rewrite_as_target_mac && !target->rewrite_as_target_vlan))
    {
        pkt_len += ether_header_construct(ethdev, pkt_buffer, ether_addr_src,
			&target->ether_addr, target->vlan_tci, l3_protocol);
    }
    // !target->rewrite_as_target_mac &&  target->rewrite_as_target_vlan
    // !target->rewrite_as_target_mac && !target->rewrite_as_target_vlan
    else
    {
        pkt_len += ether_header_construct(ethdev, pkt_buffer, ether_addr_src,
			ether_addr_dst, target->vlan_tci, l3_protocol);
    }

    /* IPv4/IPv6 Header */
    pkt_len += ip_header_construct_by_stream_addr(addr,
        (unsigned char *)pkt_buffer + pkt_len,
        payload_len + sizeof(struct tcphdr), 0, 0x1000, 0, 128, IPPROTO_TCP);

    /* TCP header and payload */
    pkt_len += tcp_header_construct_by_stream_addr(addr,
        (unsigned char *)pkt_buffer + pkt_len, seq, ack, flags, 0xffff, 0,
        (unsigned char *)payload, payload_len);

    int ret = traffic_mirror_ethdev_finish(ethdev, tid, pkt_len);
    if (unlikely(ret < 0))
    {
        //TODO: 统计计数
        return;
    }
}

static void tcp_send_to_target(struct tfe_stream_addr * addr, struct ether_addr * ether_addr_src,
	struct ether_addr * ether_addr_dst, struct traffic_mirror_ethdev * ethdev,
	struct traffic_mirror_rebuild_target * target, unsigned int tid, const char * payload, unsigned int payload_len,
	unsigned int seq, unsigned int ack, char flags)
{
	unsigned int payload_offset = 0;
	unsigned mss = ethdev->mtu - (MAX(sizeof(struct iphdr), sizeof(struct ip6_hdr)) + sizeof(struct tcphdr));

	/* handshake or farewell */
	if (payload == NULL || payload_len == 0)
	{
		tcp_segment_send_to_target_group(addr, ether_addr_src, ether_addr_dst,
			ethdev, target, tid, 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, ether_addr_src, ether_addr_dst,
			ethdev, target, tid, 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 ether_addr * c_ether_addr, struct ether_addr * s_ether_addr,
    struct traffic_mirror_rebuild_target * 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->c_ether_addr = *c_ether_addr;
    instance->s_ether_addr = *s_ether_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, unsigned int tid)
{
	tcp_send_to_target(instance->c_s_addr, &instance->c_ether_addr, &instance->s_ether_addr,
		instance->ethdev, instance->target, tid,NULL, 0, instance->c_seq, 0, TCP_SYN_FLAG);

	tcp_send_to_target(instance->s_c_addr, &instance->s_ether_addr, &instance->c_ether_addr,
		instance->ethdev, instance->target, tid,NULL, 0, instance->s_seq,
		instance->c_seq + 1, TCP_SYN_FLAG | TCP_ACK_FLAG);

	tcp_send_to_target(instance->c_s_addr, &instance->c_ether_addr, &instance->s_ether_addr,
		instance->ethdev, instance->target, tid,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, unsigned int tid,
    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(instance->c_s_addr, &instance->c_ether_addr, &instance->s_ether_addr,
			instance->ethdev, instance->target, tid, data, datalen, instance->c_seq, instance->s_seq, TCP_ACK_FLAG);
		instance->c_seq += datalen;
	}
	else
	{
		tcp_send_to_target(instance->s_c_addr, &instance->s_ether_addr, &instance->c_ether_addr, instance->ethdev,
			instance->target, tid, 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, unsigned int tid)
{
    /* C->S FIN */
	tcp_send_to_target(instance->c_s_addr, &instance->c_ether_addr, &instance->s_ether_addr, instance->ethdev, instance->target, tid,
		NULL, 0, instance->c_seq, instance->s_seq, TCP_FIN_FLAG | TCP_ACK_FLAG);

    /* C->S FIN, ACK */
	tcp_send_to_target(instance->s_c_addr, &instance->s_ether_addr, &instance->c_ether_addr, instance->ethdev, instance->target, tid,
		NULL, 0, instance->s_seq, instance->c_seq + 1, TCP_ACK_FLAG);

    instance->c_seq += 1;

    /* S->C FIN */
	tcp_send_to_target(instance->s_c_addr, &instance->s_ether_addr, &instance->c_ether_addr, instance->ethdev, instance->target, tid,
		NULL, 0, instance->s_seq, instance->c_seq, TCP_FIN_FLAG | TCP_ACK_FLAG);

    /* C->S FIN, ACK */
	tcp_send_to_target(instance->c_s_addr, &instance->c_ether_addr, &instance->s_ether_addr, instance->ethdev, instance->target, tid,
		NULL, 0, instance->c_seq, instance->s_seq + 1, TCP_ACK_FLAG);

    instance->s_seq += 1;
}