/*
    Copyright (C) 2008 Andreas Baumann <abaumann@yahoo.com>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

/* Unix traditional and most portable way of a TCP/IP server using select().
 *
 * Also tests asynchrononous accept (to avoid DoSAs on blocking accepts).
 */

#include "port/netdb.h"			/* for getaddrinfo, getnameinfo */
#include "network/network.h"		/* for networking stuff */
#include "port/stdlib.h"		/* for EXIT_XXX */
#include "port/stdio.h"			/* for fprintf */
#include "port/string.h"		/* for memset */
#include "port/stdbool.h"		/* for bool, true, false */
#include "port/unused.h"		/* for WOLF_UNUSED */

#include <unistd.h>			/* for close */
#include <errno.h>			/* for errno */
#include <sys/select.h>			/* for select */
#include <assert.h>			/* for assertions */
#include <signal.h>			/* for signal */

#define DEBUG 0

#define MAX_ACCEPT_IDLE_TIMEOUT 4
#define MAX_IDLE_TIMEOUT 10

#if DEBUG
static char *fd_set_to_string( fd_set *fd, int min_fd, int max_fd, char *buf, size_t buflen ) {
	int i;
	size_t j;

	assert( min_fd <= max_fd );

	for( i = min_fd, j = 0; i <= max_fd && j < buflen; i++, j++ ) {
		if( FD_ISSET( i, fd ) ) {
                                buf[j] = '1';
		} else {
			buf[j] = '0';
		}
	}

	buf[j] = '\0';

	return buf;
}
#endif

#define BUFFER_SIZE 4096

static volatile bool terminate = false;

static void term_handler( int sig ) {
	WOLF_UNUSED( sig );
	terminate = true;
}

int main( int argc, char* argv[] ) {
	int error;
	int serv_fd = -1;
	int res;
	char *host;
	char *service;
	struct addrinfo hints;
	struct addrinfo *result = NULL;
	wolf_network_sockaddr_union_t client_addr;
	socklen_t client_addr_len;
	fd_set read_set;
	fd_set write_set;
	struct timeval timeout;
	char client_hostname[NI_MAXHOST] = "";
	char client_service[NI_MAXSERV] = "";
	int on;
	int client_fd = -1;
	int idle_secs;
	char buffer[BUFFER_SIZE];
	size_t fd_to_write;
	size_t fd_written;
	size_t total_fd_read;
	size_t total_fd_written;
	int count;
	bool fd_eof;

	if( argc != 3 ) {
		fprintf( stderr, "usage: test2 <host> <port>\n" );
		goto FAIL;
	}

	host = argv[1];
	service = argv[2];

#ifdef _WIN32
	WSADATA wsa_data;
	WSAStartup( MAKEWORD( 2, 2 ), &wsa_data );
#endif

	signal( SIGTERM, term_handler );
	signal( SIGINT, term_handler );
	signal( SIGPIPE, SIG_IGN );

	/* tell getaddrinfo what we want */
	memset( &hints, 0, sizeof( struct addrinfo ) );
	hints.ai_flags = AI_PASSIVE;
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;
 
	/* resolve the domain name into a list of addresses */
	error = getaddrinfo( host, service, &hints, &result );
	if( error != 0 ) {
		fprintf( stderr, "getaddrinfo failed: %s (%d)\n",
			gai_strerror( error ), error );
		goto FAIL;
	}

	/* open the server endpoint */
	serv_fd = socket( result->ai_family, result->ai_socktype, result->ai_protocol );
	if( serv_fd < 0 ) {
		fprintf( stderr, "socket failed: %s (%d)\n", strerror( errno ), errno );
		goto FAIL;
	}

	/* setsockopt SO_REUSEADDR, this is handy while testing or deployment, in
	 * production in can lead to a certain type of rerouting packets to wrong
	 * servers
	 */
	on = 1;
	res = setsockopt( serv_fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof( on ) );
	if( res < 0 ) {
		(void)close( serv_fd );
		fprintf( stderr, "setsockopt(SO_REUSEADDR) failed: %s (%d)\n", strerror( errno ), errno );
		goto FAIL;
	}

	/* bind server socket */
	res = bind( serv_fd, result->ai_addr, result->ai_addrlen );
	if( res < 0 ) {
		fprintf( stderr, "bind failed: %s (%d)\n", strerror( errno ), errno );
		goto FAIL;
	}

	/* listen to the socket */
	res = listen( serv_fd, 10 );
	if( res < 0 ) {
		fprintf( stderr, "listen failed: %s (%d)\n", strerror( errno ), errno );
		goto FAIL;
	}

	/* set socket non-blocking for accepts (Stevens 15.6) */
	if( wolf_network_set_nonblocking( serv_fd ) != WOLF_OK ) {
		goto FAIL;
	}

	/* accept the connection in the kernel, select ready for ready tells us there
	 * is a connection ready. Then we have to call accept again and check the result
	 */
FIRST_ACCEPT_AGAIN:
	client_addr_len = sizeof( client_addr );
	client_fd = accept( serv_fd, &client_addr.addr, &client_addr_len );
	if( client_fd < 0 ) {
		if( errno == EINTR ) {
			if( terminate ) goto OK;
			goto FIRST_ACCEPT_AGAIN;
		} else if( errno  == EAGAIN || errno == EWOULDBLOCK ) {
			/* as expected, skip */
		} else {
			fprintf( stderr, "first non-blocking accept failed: %s (%d)\n", strerror( errno ), errno );
			goto FAIL;
		}
	}

	/* go into a select and wait till we are ready for read (this indicates a
	 * possibly new connection)
	 */
ACCEPT_SELECT_AGAIN:
	FD_ZERO( &read_set );
	FD_SET( serv_fd, &read_set );
	timeout.tv_sec = MAX_ACCEPT_IDLE_TIMEOUT;
	timeout.tv_usec = 0;

	res = select( serv_fd + 1, &read_set, NULL, NULL, &timeout );
	if( res < 0 ) {
		if( errno == EINTR ) {
			if( terminate ) goto OK;
			goto ACCEPT_SELECT_AGAIN;
		} else if( errno == EAGAIN ) {
			goto ACCEPT_SELECT_AGAIN;
		} else {
			/* fatal errors */
			fprintf( stderr, "accept/select failed: %s (%d)\n",
				strerror( errno ), errno );
			goto FAIL;
		}
	} else if( res == 0 ) {
		/* accept timeout, here we could to periodic stuff, also if we
		 * get blocked somewhere above we could recover..
		 */
		fprintf( stderr, "Idle. No new connections..\n" );
		goto ACCEPT_SELECT_AGAIN;
	} else {
		if( FD_ISSET( serv_fd, &read_set ) ) {
			/* this is a new connection */
		} else {
			/* can't happen! */
			assert( false );
		}
	}

	/* second accept, now we have to check for possible "short-and-run-away"
	 * connects.
	 */
ACCEPT_AGAIN:
	client_addr_len = sizeof( client_addr );
	client_fd = accept( serv_fd, &client_addr.addr, &client_addr_len );
	if( client_fd < 0 ) {
		if( errno == EINTR ) {
			if( terminate ) goto OK;
			/* interrupted, again */
			goto ACCEPT_AGAIN;
		} else if( errno == ECONNABORTED ) {
			/* connection run away */
			goto FIRST_ACCEPT_AGAIN;
		} else {
			fprintf( stderr, "second non-blocking accept failed: %s (%d)\n", strerror( errno ), errno );
			goto FAIL;
		}
	}

	/* determine where the request came from */
	res = getnameinfo( &client_addr.addr, client_addr_len,
		client_hostname, NI_MAXHOST,
		client_service, NI_MAXSERV,
		 NI_NUMERICSERV | NI_NUMERICHOST );
	if( res < 0 ) {
		fprintf( stderr, "getnameinfo failed: %s (%d)\n",
			gai_strerror( error ), error );
		goto FAIL;
	}

	fprintf( stderr, "New connection from %s, port %s\n", client_hostname, client_service );

	/* handle connection input/output as echo server */

	idle_secs = 0;
	count = 0;
	fd_to_write = 0;
	fd_written = 0;
	total_fd_read = 0;
	total_fd_written = 0;
	fd_eof = false;

	do {
		ssize_t rres;
#if DEBUG
		char buf[10];
		char buf2[10];
#endif
		FD_ZERO( &read_set );
		FD_ZERO( &write_set );
		timeout.tv_sec = 1;
		timeout.tv_usec = 0;
		if( fd_to_write == 0 ) {
			FD_SET( client_fd, &read_set );
		}
		if( fd_to_write > 0 ) {
			FD_SET( client_fd, &write_set );
		}

		/* wait for events or timeout */
#if DEBUG
fprintf( stderr, "select call read_fd: %s write_fd: %s buf: %d eof: %d\n",
	fd_set_to_string( &read_set, client_fd, client_fd, buf, 10 ),
	fd_set_to_string( &write_set, client_fd, client_fd, buf2, 10 ),
	fd_to_write, fd_eof );
#endif

		res = select( client_fd + 1, &read_set, &write_set, NULL, &timeout );

#if DEBUG
fprintf( stderr, "select %04d read_fd: %s write_fd: %s buf: %d eof: %d\n",
	res, fd_set_to_string( &read_set, client_fd, client_fd, buf, 10 ),
	fd_set_to_string( &write_set, client_fd, client_fd, buf2, 10 ),
	fd_to_write, fd_eof );
#endif

		if( res < 0 ) {
			if( errno == EINTR || errno == EAGAIN ) {
				/* skip */
			} else {
				/* fatal errors */
				fprintf( stderr, "select failed: %s (%d)\n",
					strerror( errno ), errno );
				goto FAIL;
			}
		} else if( res == 0 ) {
			/* timeout */
			idle_secs++;
			if( idle_secs > MAX_IDLE_TIMEOUT ) {
				fprintf( stderr, "Idle connection after %d seconds..terminating\n",
					idle_secs );
				goto END;
			}
		} else {
			/* something happened */
			idle_secs = 0;

			/* if we have something to write back to the client, do that */
			if( FD_ISSET( client_fd, &write_set ) ) {
				assert( fd_to_write > 0 );
				rres = write( client_fd, buffer + fd_written, fd_to_write );
				if( rres < 0 ) {
					if( errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR ) {
						/* skip */
					} else {
						fprintf( stderr, "write to stdout failed: %s (%d)\n",
							strerror( errno ), errno );
						goto FAIL;
					}
				} else {
					assert( rres >= 0 );
					fd_written += (size_t)rres;
					assert( fd_to_write >= (size_t)rres );
					fd_to_write -= (size_t)rres;
					fd_written += (size_t)rres;
					total_fd_written += (size_t)rres;
				}
			}

			/* if there is space in the read buffer and something to be read, do so */
			if( FD_ISSET( client_fd, &read_set ) ) {
				assert( fd_to_write == 0 );
				rres = read( client_fd, buffer, BUFFER_SIZE );
				if( rres < 0 ) {
					if( errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR ) {
						/* skip */
					} else {
						fprintf( stderr, "read from socket failed: %s (%d)\n",
							strerror( errno ), errno );
						goto FAIL;
					}
				} else if( rres == 0 ) {
					/* EOF on socket */
					shutdown( client_fd, SHUT_RD );
					fd_eof = true;
#if DEBUG
					fprintf( stderr, "EOF on socket\n" );
#endif
				} else {
					assert( rres > 0 );
					fd_written = 0;
					fd_to_write = (size_t)rres;
					total_fd_read += (size_t)rres;
				}
			}
		}

		count++;
		if( count % 10000 == 0 ) {
			fprintf( stderr, "Transfered nof_selects: %d, fd-in: %zd, fd-out: %zd\n",
				count, total_fd_read, total_fd_written );
		}
	} while( !terminate && !( fd_eof && fd_to_write == 0 ) );

END:

	fprintf( stderr, "Terminated nof_selects: %d, fd-in: %zd, fd-out: %zd\n",
		count, total_fd_read, total_fd_written );

	(void)close( client_fd );
	(void)close( serv_fd );
	freeaddrinfo( result );

	goto OK;

FAIL:

	if( client_fd >= 0 ) (void)close( client_fd );
	if( serv_fd >= 0 ) (void)close( serv_fd );
	if( result != NULL ) freeaddrinfo( result );

#ifdef _WIN32
	WSACleanup( );
#endif

	return EXIT_FAILURE;

OK:

#ifdef _WIN32
	WSACleanup( );
#endif

	return EXIT_SUCCESS;
}