////////////////////////////////////////
///@file rx_sample.c
///
///@brief pag接口收包示例代码，收包线程可以指定收包个数和打印报文内容
///
///修改记录:
///
///@author lifengw@sugon.com
///@date 2013-10-13
///
///修改内容：建立文档
////////////////////////////////////////

//系统路径下头文件
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>  //getopt
#include <string.h>
#include <pthread.h>
#include <signal.h>
#include <sys/time.h>
#include <linux/types.h>
#include <netinet/ip.h> //struct ip
#include <netinet/ip6.h>
#include <linux/tcp.h>
#include <net/ethernet.h>  //struct ether_header

#include <libpag.h>

///// 宏定义 ///////

#define DEF_THREADS_NUM 16   ///<默认线程数
#define MAX_THREADS_NUM 64  ///<系统支持的最大线程数 

#define TCPDUMP_MAGIC 0xa1b2c3d4
#define SNAP_LEN_DEF 96
#define PCAP_VERSION_MAJOR 2
#define PCAP_VERSION_MINOR 4
#define DLT_EN10MB 1

#define rtol(r, n) ((r) << (n) | ((r) >> ((8 * sizeof(r)) - (n))))

///// 数据结构定义 ///////

struct rand_func {
        uint32_t x, y, z;
};

struct pcap_file_header {
	uint32_t magic;
	uint16_t version_major;
	uint16_t version_minor;
	uint32_t thiszone;
	uint32_t sigfigs; 
	uint32_t snaplen; 
	uint32_t linktype;
};

struct pcap_timeval {
	uint32_t tv_sec;		
	uint32_t tv_usec;
};

struct pcap_pkthdr {
	struct pcap_timeval ts;	
	uint32_t caplen;		    
	uint32_t len;		   
};

////////// 数据 /////////

static int glob_loop_flag_main;
static int glob_loop_flag_thread;
static int glob_copy_time;
static int glob_print_pkt_flag;
static int glob_pkts_num;
static int glob_threads_num;
static int glob_pktsend_flag;
static int glob_get_frame;
static int glob_rand_test;
static char *glob_write_pcap = NULL;
static FILE *glob_filep = NULL;

static struct rand_func randf[MAX_THREADS_NUM];

///计数器，每个线程分别计数
ulong glob_pkts[MAX_THREADS_NUM][16];
ulong glob_bytes[MAX_THREADS_NUM][16];


pthread_mutex_t print_mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t pcap_mutex = PTHREAD_MUTEX_INITIALIZER;

///////   函数列表  ///////////////

static inline void my_srand(struct rand_func *f, uint32_t seed)
{
        unsigned n;

        f->x = 6247;
        f->y = 3848;
        f->z = 0;

        for (n = ((seed >> 22) & 0x3ff) + 20; n > 0; n--)
                f->x -= rtol(f->x, 21);
        for (n = ((seed >> 11) & 0x7ff) + 20; n > 0; n--)
                f->y = rtol(f->y, 11) - f->y;
        for (n = ((seed      ) & 0x7ff) + 20; n > 0; n--)
                f->z = 3286325185u - rtol(f->z, 19);
}

static inline uint32_t my_rand(struct rand_func *f)
{
//	f->x -= rtol(f->x, 21);
//	return f->x;

	f->x -= rtol(f->x, 21);
	f->y = rtol(f->y, 11) - f->y;
	f->z = 3286325185u - rtol(f->z, 19);

	return f->x ^ f->y ^ f->z;

//	f->x = ~(2911329625u*f->x); 
//	f->x = rtol(f->x, 17);
//
//	f->y = 4031235431u * f->y; 
//	f->y = rtol(f->y, 15);
//
//	f->z = 3286325185u - rtol(f->z, 19);

//	return (f->x + f->y) ^ f->z;
}

///@brief 输出程序帮助信息
///@param cmd 程序命令(输入)
void usage(char* cmd) 
{
	printf("Usage: %s [-m threads_num][-n pkt_num][-c copy_time][-?]\n", cmd);
	printf("  -m threads number, default = %d\n", DEF_THREADS_NUM);
	printf("  -n pkts num to rx, default = continously\n");
	printf("  -c copy time of each rx pkt, default = 0\n");
	printf("  -s send rcv pkt\n");
	printf("  -? print usage message\n");
	printf("\n");
}

///@brief 打印收包信息的线程，每隔1秒显示一次
///@param ptr 线程建立时的指针，这里没有使用
void print_speed_thread(void *ptr)
{
	int i;
	ulong old_pkts[MAX_THREADS_NUM];
	ulong old_bytes[MAX_THREADS_NUM];

	for (i = 0; i < glob_threads_num; i++)
	{
		old_pkts[i] = 0;
		old_bytes[i] = 0;
	}

	while (glob_loop_flag_thread)
	{
		sleep(1); 
		printf("\nCurrent speed----- ");
		printf("\n");        
		for (i = 0; i < glob_threads_num; i++) 
		{
			printf("\t Thread[%d]:\t %8lu pps, %8lu M bps\n", i, 
					glob_pkts[i][0] - old_pkts[i], 
					(glob_bytes[i][0] - old_bytes[i]) >> 17);
			old_pkts[i] = glob_pkts[i][0];
			old_bytes[i] = glob_bytes[i][0];
		}
	}
}

///@brief 拷贝报文（模拟应用处理包的负载）
///@param pkt 指向报文的指针
///@param len 报文长度
/* Make sure the gcc not optimization out this function */
static void copy_pkt(void *pkt, int len) __attribute__((optimize("O0")));

static void copy_pkt(void *pkt, int len)
{
	int i;
	volatile uint8_t p8_tmp[2000];
	uint8_t *p8 = (uint8_t *)pkt;

	len = (len < 2000) ? len : 2000;
	
	///拷贝报文
	for (i = 0; i < len; i++) 
	{
		p8_tmp[i] = p8[i];
	}

    (void)p8_tmp;
    (void)p8;
}

///@brief 转发接收到的数据报文
///@param pkt 指向报文帧头的指针
static void send_rcvframe(void *pkt, int len, int sid)
{
	void *buff;
	int type;

	type = 0;
	if (glob_get_frame) {
		buff = pag_getsendbuf_eth(sid, 0);
		while (!buff) {
			usleep(50);
			buff = pag_getsendbuf_eth(sid, 0);
		}

		memcpy(buff, pkt, len);
		pag_send_eth(sid, len, 0);
	} else {
		buff = pag_getsendbuf(sid);
		while (!buff) {
			usleep(50);
			buff = pag_getsendbuf(sid);
		}

		memcpy(buff, pkt, len);
		type = (type == 0) ? 1 : 0;
		pag_send(type, sid, len);
	}
}

///@brief 打印报文
///@param pkt 指向报文的指针
///@param len 报文长度
static void print_pkt(char *pkt_s, int len)
{
	int i, rc;
	int pos = 0; 
	unsigned char *pkt = (unsigned char *) pkt_s;

	rc = pthread_mutex_lock(&print_mutex);
	if (rc) {
		fprintf(stderr, "pthread_mutex_lock error\n");
		exit(-1);
	}
	for (i = 0; i < len;i++) {
		if (i % 16 == 0) {
			fprintf(stdout, "\n");
			fprintf(stdout, "%08x:  ", pos);
			fprintf(stdout, "%02x  ", pkt[i]);
			pos += 16;
		}
		else
			fprintf(stdout, "%02x  ", pkt[i]);
	}
	fprintf(stdout, "\n");
	rc = pthread_mutex_unlock(&print_mutex);
	if (rc) {
		fprintf(stderr, "pthread_mutex_unlock error\n");
		exit(-1);
	}
}

int get_ip_len(void *pkt)
{
	int len;
	unsigned char *p8;
	struct iphdr* iph;     

	len = 0;
	p8 = (unsigned char*)(pkt);

	iph = (struct iphdr *) p8; 
	if (iph->version == 4) {
		len = ntohs((unsigned short)iph->tot_len);
	} else if (iph->version == 6) {
		struct ip6_hdr *ip6h;
		ip6h = (struct ip6_hdr *) p8;
		len = ntohs((unsigned short)ip6h->ip6_plen);
	}

	return len;
}

int write_pcap_file(FILE *fp, void *pkt, int len)
{
	struct timeval tv;
	static int first = 1;
	struct pcap_file_header hdr;
	struct pcap_pkthdr pcap_hdr;
	int ret;

	gettimeofday(&tv, NULL);
	ret = pthread_mutex_lock(&pcap_mutex);
	if (ret) {
		fprintf(stderr, "pthread_mutex_lock error\n");
		exit(-1);
	}

	if (first == 1) {
		hdr.magic = TCPDUMP_MAGIC;
		hdr.version_major = PCAP_VERSION_MAJOR;
		hdr.version_minor = PCAP_VERSION_MINOR;
		hdr.thiszone = 0;
		hdr.snaplen = 65535;
		hdr.sigfigs = 0;
		hdr.linktype = DLT_EN10MB;

		ret = fwrite((char *)&hdr, sizeof(struct pcap_file_header), 1, fp);
		if (ret != 1) {
			fprintf(stderr, "### Error : fwrite pcap_file_header error!\n");
			return -1;
		}
		first = 0;
	}
	
	pcap_hdr.ts.tv_sec = tv.tv_sec;
	pcap_hdr.ts.tv_usec = tv.tv_usec;
	pcap_hdr.len = len;
	pcap_hdr.caplen = len;

	ret = fwrite((char *)&pcap_hdr, sizeof(struct pcap_pkthdr), 1, fp);
	if (ret != 1) {
		fprintf(stderr, "### Error : fwrite pcap_pkthdr error!\n");
		return -1;
	}
	ret = fwrite((char *)pkt, len, 1, fp);
	if (ret != 1) {
		fprintf(stderr, "### Error : fwrite pcap_pkt error!\n");
		return -1;
	}

	ret = pthread_mutex_unlock(&pcap_mutex);
	if (ret) {
		fprintf(stderr, "pthread_mutex_unlock error\n");
		exit(-1);
	}
	return 0;
}

///@brief 接收报文的线程
///@param ptr 线程建立时的指针，这里指向线程的数据流号
void rx_thread(void *ptr)
{
	int ret;
	int sid, i;
	int len;
	void *pkt;

	///获取缓冲区id
	sid = *((int *)ptr);

	my_srand(&randf[sid], time(NULL));

	///循环，直到线程循环标志清零   
	while (glob_loop_flag_thread)
	{ 
		///--获取一个报文，如果没有报文，则循环等待
		if (glob_get_frame || (glob_filep != NULL)) {
			pkt = pag_get_frame(sid);
			if (!pkt) {
				//usleep(500);
				continue;
			}
			len = pag_get_frame_length(sid);
		} else {
			pkt = pag_get(sid);
			if (!pkt) {
				//usleep(500);
				continue;
			}
			len = get_ip_len(pkt);
		}

		if (glob_filep != NULL) {
			ret = write_pcap_file(glob_filep, pkt, len);
			if (ret != 0)
				return ;
			fflush(glob_filep);
		}

		///--增加报文计数
		glob_pkts[sid][0]++;
		glob_bytes[sid][0] += len;

		///--根据参数对报文做相应处理
		if (glob_rand_test) {
			int n;
			n = my_rand(&randf[sid]) & 0xF;
			for (i = 0; i < n; i++) 
				copy_pkt(pkt, len);
			if (n & 0x1)
				send_rcvframe(pkt, len, sid);
			continue;
		} else {
			for (i = 0; i < glob_copy_time; i++) {
				copy_pkt(pkt, len);
			}
		}

		if (glob_pktsend_flag == 1) {
			send_rcvframe(pkt, len, sid);
		}

		if (glob_print_pkt_flag) {
			print_pkt(pkt, len);
		}
	}
}

///@brief 信号响应函数，收到信号后把全局循环标志无效掉
void stop_loop(int sig)
{
	glob_loop_flag_main = 0;
}

///@brief 主函数
///@return 成功返回0，失败返回-1
int main(int argc, char *argv[])
{
	int c, i;
	struct timeval ts, te;
	double sec; 
	ulong all_pkts, all_bytes;
	sigset_t set;
	pthread_t rx_tids[MAX_THREADS_NUM], print_tid;
	int sids[MAX_THREADS_NUM];

	///打印使用信息
	usage(argv[0]);

	///读取参数
	glob_copy_time = 0;
	glob_print_pkt_flag = 0;
	glob_pkts_num = 0;
	glob_threads_num = DEF_THREADS_NUM;
	glob_pktsend_flag = 0;
	glob_rand_test = 0;
	glob_get_frame = 0;

	while ((c = getopt(argc, argv, "c:n:m:spfrw:?")) != -1) 
	{
		switch (c)
		{
			case 'c':
				glob_copy_time = atoi(optarg);
				break;
			case 'n':
				glob_pkts_num = atoi(optarg);
				break;
			case 'm':
				glob_threads_num = atoi(optarg);
				if (glob_threads_num > MAX_THREADS_NUM)
				{
					printf("\nError: threads num invalid\n");
					return -1;
				}
				break;
			case 's':
				glob_pktsend_flag = 1;
				break;
			case 'p':
				glob_print_pkt_flag = 1;
				break;
			case 'f':
				glob_get_frame = 1;
				break;
			case 'r':
				glob_rand_test = 1;
				break;
			case 'w':
				glob_write_pcap = (char *)optarg;
				break;
			case '?':
			default:
				usage(argv[0]);
				return 0;
		}
	}

	if (glob_write_pcap != NULL) {
		if ((glob_filep = fopen(glob_write_pcap, "w")) == NULL) {
			printf("### Error : can not create file %s\n", glob_write_pcap);
			return -1;
		}
	}

	///打开设备
	if (pag_open() < 0)
	{
		printf("### Error : can not open device! ###\n");
		return -1;
	}

	///创建线程
	glob_loop_flag_thread = 1;
	for (i = 0; i < glob_threads_num; i++)
	{
		sids[i] = i;
		if (pthread_create(&rx_tids[i], NULL, (void *)rx_thread, (void *)&sids[i]) != 0)
		{
			printf("### Error: Can not creat threads!\n");    
			pag_close();
			return -1;
		}
	}

	///如果不打印报文内容，则打印报文速度
	if (!glob_print_pkt_flag)
	{
		if (pthread_create(&print_tid, NULL, (void *)print_speed_thread, NULL) != 0)   
		{
			printf("### Error: Can not creat threads!\n");    
			pag_close();
			return -1;
		}
	}

	///等待结束信号，或着收到期望的包数 
	sigemptyset(&set);
	sigprocmask(SIG_SETMASK, &set, NULL);
	signal(SIGINT, stop_loop);

	///记录开始时间
	gettimeofday(&ts, NULL);   

	glob_loop_flag_main = 1;
	while (glob_loop_flag_main) 
	{
		usleep(10);  

		all_pkts = 0;
		all_bytes = 0;
		for (i = 0; i < glob_threads_num; i++) 
		{
			all_pkts += glob_pkts[i][0];
		}

		if (glob_pkts_num != 0 && glob_pkts_num <= all_pkts)
		{
			break;
		}
	}

	///记录结束时间
	gettimeofday(&te, NULL);  

	///停止线程
	glob_loop_flag_thread = 0;
	for (i = 0; i < glob_threads_num; i++)
	{
		pthread_join(rx_tids[i], NULL);
	} 
	if (!glob_print_pkt_flag)
	{   
		pthread_join(print_tid, NULL);
	}

	///关闭设备
	pag_close();    

	if (glob_filep != NULL)
		fclose(glob_filep);

	///打印总结信息
	printf("\n");        
	printf("Pkts recieved in threads: \n");

	sec = (float)(te.tv_sec - ts.tv_sec) + ((float)(te.tv_usec - ts.tv_usec)) / 1000000.f;
	all_pkts = 0;
	all_bytes = 0;
	for(i = 0; i < glob_threads_num; i++) 
	{
		printf("\t Thread[%d]:\t <%8lu> pkts, <%8lu M> bytes\n", i,
				glob_pkts[i][0], glob_bytes[i][0] >> 20);
		all_pkts += glob_pkts[i][0];
		all_bytes += glob_bytes[i][0];
	}

	printf("Total : %lu pkts,  %lu bytes\n", all_pkts, all_bytes);
	printf("Speed : < %u > pps, < %u M > bps\n", 
			(int)(all_pkts / sec), (int)((all_bytes >> 17) / sec));
	printf("\n");        

	return 0;
}