/*
 * Copyright 2010, Atis Elsts, the.kfx@gmail.com
 * Distributed under the terms of the MIT license.
 */


#include <unistd.h>
#include <memory.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>


const unsigned short TEST_PORT = 40000;


void
usage()
{
	printf("server [tcp|udp] [4|6] [local-address]\n");
	exit(1);
}


void
recvLoop(int fd)
{
	for (;;)	{
		char buffer[1000];
		int status = recv(fd, buffer, sizeof(buffer) - 1, 0);
		if (status < 0) {
			perror("recv");
			exit(-1);
		}
		if (status == 0) {
			printf("received EOF!\n");
			break;
		} else {
			buffer[status] = 0;
			printf("received %d bytes: \"%s\"\n", status, buffer);
		}
	}
}


int
main(int argc, char *argv[])
{
	int socketType = SOCK_DGRAM;
	int socketFamily = AF_INET;
	if (argc > 1) {
		if (!strcmp(argv[1], "tcp"))
			socketType = SOCK_STREAM;
		else if (!strcmp(argv[1], "udp"))
			socketType = SOCK_DGRAM;
		else
			usage();
	}
	if (argc > 2) {
		switch (atoi(argv[2])) {
		case 4:
			socketFamily = AF_INET;
			break;
		case 6:
			socketFamily = AF_INET6;
			break;
		default:
			usage();
		}
	}

	sockaddr_storage localAddress;
	memset(&localAddress, 0, sizeof(localAddress));
	localAddress.ss_family = socketFamily;
	((sockaddr_in *) &localAddress)->sin_port = htons(TEST_PORT);

	if (argc > 3) {
		do {
			void *dstBuffer = &((sockaddr_in *) &localAddress)->sin_addr;
			if (inet_pton(AF_INET, argv[3], dstBuffer) == 1) {
				printf("using IPv4 local address\n");
				localAddress.ss_family = AF_INET;
				break;
			}

			dstBuffer = &((sockaddr_in6 *) &localAddress)->sin6_addr;
			if (inet_pton(AF_INET6, argv[3], dstBuffer) == 1) {
				printf("using IPv6 local address\n");
				localAddress.ss_family = AF_INET6;
				break;
			}

			usage();
		} while (false);
	}

	int fd = socket(socketFamily, socketType, 0);
	if (fd < 0) {
		perror("socket");
		return -1;
	}

	if (bind(fd, (sockaddr *)&localAddress, localAddress.ss_family == AF_INET ?
			sizeof(sockaddr_in) : sizeof(sockaddr_in6)) < 0) {
		perror("bind");
		return -1;
	}

	switch (socketType) {
	case SOCK_DGRAM:
		for (;;)
			recvLoop(fd);
		break;
	case SOCK_STREAM:
		if (listen(fd, 5) < 0) {
			perror("listen");
			return 1;
		}
		for (;;) {
			int clientfd = accept(fd, NULL, 0);
			if (clientfd < 0) {
				perror("accept");
				return 1;
			}
			printf("TCP server: got some client!\n");
			if (fork() != 0) {
				// parent code
				close(clientfd);
				continue;
			}
			// child code
			close(fd);
			recvLoop(clientfd);
			exit(0);
		}
		break;
	}
}