#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <netdb.h> 


#define DEFAULT_SEND_SIZE 256
#define MIN_BUFFER_SIZE (2*1024)


/*
 * Supplied string can contain up to 32 hex chars.  If it contains less, the
 * rest are assumed to be zero.
 */
static int
str_to_key(uint8_t *key, const char *str)
{
	unsigned int len, i;
	unsigned int byte_num;
	unsigned int nibble_num;
	char c;
	uint8_t val;
	
	len = strlen(str);
	if (len > 32)
		return (0);


	memset(key, 0, 16);
	
	byte_num = 0;
	nibble_num = 0;
	for (i = 0; i < len; i++) {
		c = tolower(str[i]);
		if (c >= '0' && c <= '9')
			val = c - '0';
		else if (c >= 'a' && c <= 'f')
			val = c - 'a' + 10;
		else
			return (0);

		key[byte_num] |= val << ((1 - nibble_num) * 4);
		nibble_num = (nibble_num + 1) % 2;
		if (nibble_num == 0)
			byte_num++;
	}

	return (1);
}

void
errno_fatal(const char *msg)
{
	perror(msg);
	exit(1);
}


void
usage(const char *progname)
{
	printf("usage: %s [options] hostname port\n", progname);
	printf("\n");
	printf("Options:\n");
	printf("  -k key       Use given pre-shared key (up to 32 hex digits) for TFO connections\n");
	printf("  -n           Use standard (non-TFO) connect (default is TFO)\n");
	printf("  -r size      Specify amount of data to wait for (default %u)\n", DEFAULT_SEND_SIZE);
	printf("  -s size      Specify amount of data to send (default %u)\n", DEFAULT_SEND_SIZE);
	printf("\n");
}


int
main(int argc, char *argv[])
{
	int i, sockfd;
	unsigned int portno;
	unsigned int send_size, receive_size;
	struct sockaddr_in serv_addr;
	struct hostent *server;
	char *buffer;
	unsigned int buffer_size;
	int rflag;
	int c;
	int use_tfo = 1;
	int rv;
	int sendto_flags;
	int val;
	int use_psk = 0;
	uint8_t psk[16];
	
	send_size = DEFAULT_SEND_SIZE;
	rflag = 0;
	while ((c = getopt (argc, argv, "k:nr:s:")) != -1) {
		switch (c)
			{
			case 'n':
				use_tfo = 0;
				break;
			case 'k':
#ifdef __FreeBSD__
				if (!str_to_key(psk, optarg)) {
					printf("Bad key format\n");
					return (1);
				}
				use_psk = 1;
#else
				printf("Pre-shared key mode is only available on FreeBSD.\n");
				return (1);
#endif
				break;
			case 'r':
				receive_size = strtoul(optarg, NULL, 10);
				rflag = 1;
				break;
			case 's':
				send_size = strtoul(optarg, NULL, 10);
				break;
			default:
			case '?':
				usage(argv[0]);
				return (1);
			}
	}
	if (!rflag)
		receive_size = send_size;

	if (argc - optind < 2) {
		usage(argv[0]);
		return (1);
	}
	portno = strtoul(argv[optind + 1], NULL, 10);
	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (sockfd < 0) 
		errno_fatal("Failed to open socket");
	server = gethostbyname(argv[optind]);
	if (server == NULL) {
		printf("No such host\n");
		return (1);
	}
	bzero((char *) &serv_addr, sizeof(serv_addr));
	serv_addr.sin_family = AF_INET;
	bcopy((char *)server->h_addr, 
	      (char *)&serv_addr.sin_addr.s_addr,
	      server->h_length);
	serv_addr.sin_port = htons(portno);

	buffer_size = (send_size > MIN_BUFFER_SIZE) ? send_size : MIN_BUFFER_SIZE;
	buffer = malloc(buffer_size);
	if (buffer == NULL) {
		printf("Failed to allocate buffer (size=%u)\n", buffer_size);
		return (1);
	}

	for (i = 0; i < send_size; i++) 
		buffer[i] = 'a' + (i % 26);

	printf("Sending %u bytes using %sTFO connection\n", send_size, use_tfo ? "" : "non-");
	
	if (!use_tfo) {
		if (connect(sockfd,(struct sockaddr *) &serv_addr,sizeof(serv_addr)) < 0) 
			errno_fatal("Failed to connect");
		sendto_flags = 0;
	} else {
#ifdef __FreeBSD__
		sendto_flags = 0;
		val = 1;
		if (use_psk) {
			struct tcp_fastopen tfo_opt;

			tfo_opt.enable = 1;
			memcpy(tfo_opt.psk, psk, sizeof(tfo_opt.psk));
			rv = setsockopt(sockfd, IPPROTO_TCP, TCP_FASTOPEN, &tfo_opt, sizeof(tfo_opt));
		} else
			rv = setsockopt(sockfd, IPPROTO_TCP, TCP_FASTOPEN, &val, sizeof(val));
		if (rv < 0)
			errno_fatal("Failed to set TCP_FASTOPEN sockopt");
#endif
#ifdef __linux__
		sendto_flags = MSG_FASTOPEN;
#endif
#ifdef __APPLE__
		serv_addr.sin_len = sizeof(serv_addr);

		sa_endpoints_t endpoints;
		endpoints.sae_srcif = 0;
		endpoints.sae_srcaddr = NULL;
		endpoints.sae_srcaddrlen = 0;
		endpoints.sae_dstaddr = (struct sockaddr *)&serv_addr;
		endpoints.sae_dstaddrlen = sizeof(serv_addr);
		rv = connectx(sockfd, &endpoints, SAE_ASSOCID_ANY, CONNECT_RESUME_ON_READ_WRITE | CONNECT_DATA_IDEMPOTENT, NULL, 0, NULL, NULL);
		if (rv < 0)
			errno_fatal("Failed to connect");
		sendto_flags = 0;
#endif
	}
	
	if (sendto(sockfd, buffer, send_size, sendto_flags, (struct sockaddr *)&serv_addr, sizeof(struct sockaddr)) < 0)
		errno_fatal("sendto failed");

	if (receive_size) {
		int bytes_remaining, total_received;
		
		printf("Waiting for %u bytes\n", receive_size);
		bytes_remaining = receive_size;
		total_received = 0;
		while ((bytes_remaining > 0) &&
		       ((rv = recv(sockfd, buffer, receive_size, 0)) > 0)) {
			bytes_remaining -= rv;
			total_received += rv;
			printf("%d bytes received\n", rv);
		}
		if (rv < 0)
			errno_fatal("recv failed\n");
		printf("%d total bytes received\n", total_received);
	}

	close(sockfd);
	return 0;
}