#include "test_framework.h"

void usage(char *progname) {
        fprintf(stderr,
		"%s [-h] [-n #] [-s #] [-d destination host] [-p port]\n",
		progname);
}

void help(char *progname) {
	usage(progname);
	fprintf(stderr,
		"  -h  .............. help (this message)\n"
		"  -d <dest> ........ send to dest, default localhost\n"
		"  -p <port> ........ remote port to contact, default %d\n"
		"  -n <max packets> . Max # of packets to send before exit\n"
		"  -s <datasize> .... Amt of extra data to include (bytes)\n",
		DEFAULT_SERVER_PORT);
}

int main(int argc, char **argv)
{
	extern char *optarg;
	extern int optind;
	int ch;

	struct hostent *hostptr;
	struct sockaddr_in dst, src;
	int sock;
	int port = DEFAULT_SERVER_PORT;
	int maxpackets = DEFAULT_MAXPACKETS;
	int extra_data = 0;
#ifdef SO_TIMESTAMP
	int on = 0;
#endif

	char *dest_host = "127.0.0.1";
	char *progname = argv[0];

	while ((ch = getopt(argc, argv, "hd:p:n:s:")) != -1)
		switch (ch) {
		case 's':
			extra_data = atoi(optarg);
			break;
		case 'n':
			maxpackets = atoi(optarg);
			break;
		case 'd':
			dest_host = optarg;
			break;
		case 'p':
			port = atoi(optarg);
			break;
		case 'h':
			help(progname);
			exit(0);
		default:
			usage(progname);
			exit(-1);
		}
	argc -= optind;
	argv += optind;

	/* Resolve the destination host */

	bzero(&src, sizeof(src));
	bzero(&dst, sizeof(dst));

	if (inet_aton(dest_host, &(dst.sin_addr))) {
		dst.sin_family = AF_INET;
	} else {
		if ( (hostptr = gethostbyname( dest_host )) == NULL) 
		{
			herror(dest_host);
			exit(-1);
		}
		
		dst.sin_family = hostptr->h_addrtype;
		bcopy(hostptr->h_addr, (caddr_t)&(dst.sin_addr),
		      hostptr->h_length);
	}

	dst.sin_port = htons(port);

	/* Set up the socket */
	sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (sock == -1) {
		perror("socket");
		exit(-1);
	}

#ifdef SO_TIMESTAMP
	on = 1;
	if (setsockopt(sock, SOL_SOCKET, SO_TIMESTAMP,
		       (const char *)&on, sizeof(on)) < 0) {
	        perror("setsockopt for SO_TIMESTAMP");
		exit(-1);
	}
#endif

	src.sin_family = AF_INET;
	src.sin_port = htons(DEFAULT_CLIENT_PORT);
	src.sin_addr.s_addr = INADDR_ANY;

	if (bind(sock, &src, sizeof(src)) == -1) {
		perror("could not bind udp socket");
		exit(-1);
	}

#ifndef NO_CONNECT
	if (connect(sock, &dst, sizeof(dst)) == -1) {
		perror("could not specify server");
		exit(-1);
	}
#endif

	/*
	 * The application provides this function;  it determines
	 * the functionality of the program.  Different clients
	 * (e.g. test_rate and test_alf) will provide different
	 * versions of the test_run function
	 *
	 * We'll time it and see how long things take to run,
	 * to make performance measurements easy.
	 */

	test_run(sock, maxpackets, &dst, dest_host, &src, extra_data);

	close(sock);
	
	exit(0);
}

long timediff(const struct timeval *time_now,
              const struct timeval *time_prev)
{
        long usecs_now;
        
        if (time_prev->tv_sec == 0 &&
            time_prev->tv_usec == 0) {
                return 0;
        }
        
        usecs_now = (time_now->tv_sec - time_prev->tv_sec) * 1000000;
        usecs_now += time_now->tv_usec - time_prev->tv_usec;;
        return usecs_now;
}

/*
 * Send a packet out the socket.  Tries to use scatter/gather
 * writes for efficiency, though this is only for efficiency
 * when using the program to measure overhead.  You could
 * just as easily use sendto() or other sending functions.
 *
 * Returns the number of bytes sent, or -1.
 * Note that this number INCLUDES the 8-byte UDP header.
 */

int send_pkt(int sock, struct sockaddr_in *dst,
	     char *extra_data, int extra_data_len,
	     int *seqno)
{
	struct timeval t;
	struct msghdr msg;
	struct iovec iov[2];
	
	struct example_pkt_hdr pkt_hdr;
	int rc;

	bzero(&msg, sizeof(msg));

	msg.msg_iov = iov;

	iov[0].iov_base = (char *)&pkt_hdr;
	iov[0].iov_len = sizeof(pkt_hdr);

	if (extra_data && extra_data_len) {
		iov[1].iov_base = extra_data;
		iov[1].iov_len = extra_data_len;
		msg.msg_iovlen = 2;
	} else {
		msg.msg_iovlen = 1;
	}

	msg.msg_name = dst;
	msg.msg_namelen = sizeof(*dst);
	
	gettimeofday(&t, NULL);
	pkt_hdr.seqno = htonl((*seqno)++);
	pkt_hdr.datalen = htonl(extra_data_len);
	pkt_hdr.send_tv_sec = htonl(t.tv_sec);
	pkt_hdr.send_tv_usec = htonl(t.tv_usec);

	rc = sendmsg(sock, &msg, 0);
#if 0
	if (rc == -1) {
		perror("sending");
	}
#endif
	if (rc > 0) { rc += 8; } /* Account for the UDP header */
	return rc;
}

/*
 * This is an INCOMPLETE function for receiving ACK packets from the
 * client.  It doesn't deal with returned data (it just ignores
 * it), and it treats all out-of-order packets as losses.
 * A real application would most likely want to re-implement this
 * function so it can react to what's happening with the client,
 * and be more fully featured.
 *
 * WARNING:  This function assumes that ALL packets sent to the server
 *           are the same length for purposes of computing the number of
 *           bytes lost in transmission.
 *
 * It attempts to use the SO_TIMESTAMP functionality, if present.
 * This is implemented in the Linux kernel distributed with the
 * CM, and in the latest 2.4-test6 betas and above.  SO_TIMESTAMP
 * may yield more accurate RTT values if your program spends a lot of
 * time processing its packets, or simply leaves them in the socket
 * buffer.
 */

int recv_pkt(int sock, int cmid, struct sockaddr_in *from,
	     char *extra_data, int *extra_data_len,
	     int *ack_seqno)
{
	unsigned int rtt;
	char pktbuf[2048];
	struct example_pkt_hdr *pkt_hdr;
	char *pkt_contents;
	int err, nlost, rc, nbytes, nrecd;
	struct timeval time_now, time_sent;

	struct msghdr msg;
	struct iovec iov;
	int gottime = 0;
	
#ifdef SO_TIMESTAMP
	char ctrl[CMSG_SPACE(sizeof(struct timeval))];
	struct cmsghdr *cmsg = (struct cmsghdr *)&ctrl;
	msg.msg_control = (caddr_t)ctrl;
	msg.msg_controllen = sizeof(ctrl);
#endif

	pkt_hdr = (struct example_pkt_hdr *)pktbuf;
	pkt_contents = pktbuf + sizeof(*pkt_hdr);

	msg.msg_name = &from;
	msg.msg_namelen = sizeof(from);
	msg.msg_iov = &iov;
	msg.msg_iovlen = 1;
	iov.iov_base = pktbuf;
	iov.iov_len = sizeof(pktbuf);

	err = recvmsg(sock, &msg, 0);

	if (err <= 0) {
		perror("recv");
		return -1;
	}
	/*
	 * Otherwise, assume that it's a valid ack
	 * from the server.
	 */

	if ((int)ntohl(pkt_hdr->seqno) < (int)*ack_seqno) {
		/*
		 * It was a re-ordered packet;  we already counted
		 * it as a loss, so discard it
		 */
		return 0;
	}

#ifdef SO_TIMESTAMP
	if (cmsg->cmsg_level == SOL_SOCKET &&
	    cmsg->cmsg_type == SCM_TIMESTAMP &&
	    cmsg->cmsg_len == CMSG_LEN(sizeof(time_now))) {
                                /* Copy to avoid alignment problems: */
	    memcpy(&time_now,CMSG_DATA(cmsg),sizeof(time_now));
	    gottime = 1;
	}
#endif
	if (!gottime) {
	    gettimeofday(&time_now, NULL);
	}
	time_sent.tv_sec = ntohl(pkt_hdr->send_tv_sec);
	time_sent.tv_usec = ntohl(pkt_hdr->send_tv_usec);

	rtt = timediff(&time_now, &time_sent);

	/* Any skipped sequence numbers are "lost" */
	nlost = ntohl(pkt_hdr->seqno) - *ack_seqno - ntohs(pkt_hdr->n_acks);
	*ack_seqno = ntohl(pkt_hdr->seqno);
	nbytes = sizeof(struct example_pkt_hdr) + ntohl(pkt_hdr->datalen);
	/*
	 * Please note what we're doing here.  The BUFUDP congestion
	 * manager implementation INCLUDES the 8-byte UDP header.
	 */
	nbytes+=8;
	nrecd = ntohs(pkt_hdr->n_acks) * nbytes;

	if (nlost == 0) {
		rc = cm_update(cmid, nrecd, 0, CM_NOLOSS, rtt);
	} else {
		rc = cm_update(cmid, nrecd, nlost*nbytes, CM_TRANSIENT, rtt);
	}
	if (rc) {
		perror("cm_update");
	}
	return 0;
}