/*
    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 client using select().
 *
 * Also tests asynchrononous connect (for proper timeout handling in case of
 * no connection).
 *
 * Doesn't work! WSAENONSOCK is returned when setting stdin to non-blocking.
 * So the BSD API is for porting pure socket code only, otherwise we should
 * not use this API!
 */

#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 */

#ifndef _WIN32
#include <unistd.h>			/* for close */
#include <fcntl.h>			/* for fcntl */
#include <sys/select.h>			/* for select */
#include <signal.h>			/* for signal */
#else
#include <io.h>
#include <Winsock2.h>
/* goes to port/unistd.h */
#define STDIN_FILENO	fileno(stdin)
#define STDOUT_FILENO	fileno(stdout)
#define STDERR_FILENO	fileno(stderr)

/* goes to a port/types.h */
typedef SSIZE_T ssize_t;
#endif

#include <errno.h>			/* for errno */
#include <assert.h>			/* for assertions */

#define DEBUG 0

#ifndef _WIN32
typedef int wolf_network_socket_handle_t;
#define wolf_network_closesocket( socket ) close( socket )

#define WOLF_NETWORK_ERRNO		errno

#define WOLF_NETWORK_EWOULDBLOCK	EWOULDBLOCK
#define WOLF_NETWORK_EINPROGRESS	EINPROGRESS
#define WOLF_NETWORK_EISCONN		EISCONN
#define WOLF_NETWORK_EINVAL		EINVAL

#else
typedef SOCKET wolf_network_socket_handle_t;
#define wolf_network_closesocket( socket ) closesocket( socket )

#define WOLF_NETWORK_ERRNO		WSAGetLastError( )

#define WOLF_NETWORK_EWOULDBLOCK	WSAEWOULDBLOCK
#define WOLF_NETWORK_EINPROGRESS	WSAEINPROGRESS
#define WOLF_NETWORK_EISCONN		WSAEISCONN
#define WOLF_NETWORK_EINVAL		WSAEINVAL

#define SHUT_RD 			SD_RECEIVE
#define SHUT_WR				SD_SEND
#endif

static bool wolf_network_sock_nonblocking( wolf_network_socket_handle_t fd ) {
#ifndef _WIN32
	int flags;

	flags = fcntl( fd, F_GETFL, 0 /* ignored */ );
	if( flags < 0 ) return false;
	flags |= O_NONBLOCK;
	flags = fcntl( fd, F_SETFL, flags );
	if( flags < 0 ) return false;
	return true;
#else
	int res;
	ULONG non_blocking;

	non_blocking = 1;
	res = ioctlsocket( fd, FIONBIO, &non_blocking );
	if( res == SOCKET_ERROR ) return false;
	return true;
#endif
}

#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 MAX_IDLE_TIMEOUT 10
#define WRITE_BUFFER_SIZE 4096
#define READ_BUFFER_SIZE 4096
#define CONNECT_TIMEOUT 2

#ifndef _WIN32
#define max(a,b) ((a) < (b) ? (b) : (a))
#endif

static volatile bool terminate = false;

static void term_handler( int sig ) {
	WOLF_UNUSED( sig );
	terminate = true;
}
 
int main( int argc, char *argv[] ) {
	char *host;
	char *service;
	struct addrinfo hints;
	struct addrinfo *result = NULL;
	int error;
	int fd = -1;
	int res;
	int idle_secs;
	char read_buffer[READ_BUFFER_SIZE];
	char write_buffer[WRITE_BUFFER_SIZE];
	size_t stdout_to_write;
	size_t stdout_written;
	size_t fd_to_write;
	size_t fd_written;
	size_t total_stdin_read;
	size_t total_fd_read;
	size_t total_fd_written;
	size_t total_stdout_written;
	int count;
	bool stdin_eof;
	bool fd_write_closed;
	bool fd_eof;
	int max_fd;
	fd_set read_set;
	fd_set write_set;
	struct timeval timeout;
	socklen_t peer_addr_len;
	wolf_network_sockaddr_union_t peer_addr;
	char peer_hostname[NI_MAXHOST] = "";
	char peer_service[NI_MAXSERV] = "";
#ifdef _WIN32
	WSADATA wsa_data;
#endif

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

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

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

	/* 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 client endpoint */
	fd = socket( result->ai_family, result->ai_socktype, result->ai_protocol );
	if( fd < 0 ) {
		fprintf( stderr, "socket failed: %s (%d)\n", strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
		goto FAIL;
	}

	/* set socket non-blocking for asynchronous connect */
	if( !wolf_network_sock_nonblocking( fd ) ) {
		fprintf( stderr, "set nonblocking failed for socket: %s (%d)\n", strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
		goto FAIL;
	}

	/* connect asynchronously, from now on no matter what happens (for instance an
	 * interrupt by a signal) the kernel tries to establish the connection! So we
	 * don't have to go back into a connect..
	 */
	res = connect( fd, result->ai_addr, result->ai_addrlen );
	if( res < 0 ) {
		if( WOLF_NETWORK_ERRNO == WOLF_NETWORK_EINPROGRESS ||
		    WOLF_NETWORK_ERRNO == WOLF_NETWORK_EWOULDBLOCK ) {
			/* this is ok, being connected, we loop later in a select */
			fprintf( stderr, "connecting in progress..\n" );
		} else if( WOLF_NETWORK_ERRNO == EINTR ) {
			/* interrupted, no problem, the connection goes on anyway */
		} else {
			fprintf( stderr, "connect failed: %s (%d)\n", strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
			goto FAIL;
		}
	} else {
		/* a fast connect */
		goto CONNECTED;
	}

	/* wait till we are connected */
CONNECT_SELECT_AGAIN:
	FD_ZERO( &read_set );
	FD_ZERO( &write_set );
	FD_SET( fd, &read_set );
	FD_SET( fd, &write_set );
	timeout.tv_sec = CONNECT_TIMEOUT;
	timeout.tv_usec = 0;
	res = select( fd + 1, &read_set, &write_set, NULL, &timeout );
	if( res < 0 ) {
		if( WOLF_NETWORK_ERRNO == EINTR || WOLF_NETWORK_ERRNO == EAGAIN ) {
			goto CONNECT_SELECT_AGAIN;
		} else {
			/* fatal errors */
			fprintf( stderr, "select failed: %s (%d)\n",
				strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
			goto FAIL;
		}
	} else if( res == 0 ) {
		/* connect timeout */
		fprintf( stderr, "Connect timeout..terminating\n" );
		goto FAIL;
	} else {
		if( FD_ISSET( fd, &read_set ) || FD_ISSET( fd, &write_set ) ) {
/* strategy second connect */
SECOND_CONNECT_AGAIN:
			res = connect( fd, result->ai_addr, result->ai_addrlen );
			if( res < 0 ) {
				if( WOLF_NETWORK_ERRNO == EINTR ) {
					/* interrupted, no problem, the connection goes on anyway */
					goto SECOND_CONNECT_AGAIN;
				} else if( WOLF_NETWORK_ERRNO == WOLF_NETWORK_EISCONN ||
					   WOLF_NETWORK_ERRNO == WOLF_NETWORK_EINVAL ) {
					/* already connected, first connect succeeded */
					goto CONNECTED;
				} else {
					fprintf( stderr, "second connect failed: %s (%d)\n", strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
					goto FAIL;
				}
			}

/* strategy setsockopt and SO_ERROR:

Not really realiable! Can result 0 though the server is not there!
			socklen_t sock_error_len;
			int sock_error;

			sock_error_len = sizeof( int );
			res = getsockopt( fd, SOL_SOCKET, SO_ERROR, &sock_error, &sock_error_len );
			if( res < 0 ) {
				fprintf( stderr, "getsockopt for connection check failed: %s (%d)\n",
					strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
				goto FAIL;
			} else {
				if( error == 0 ) {
					goto CONNECTED;
				} else {
					fprintf( stderr, "SO_ERROR is not ok: %s (%d)\n",
						strerror( error ), error );
					goto FAIL;
				}
			}
*/
		} else {
			fprintf( stderr, "Socket not ready after select\n" );
			goto FAIL;
		}
	}

CONNECTED:
	/* determine the address of the other end of the connection */
	peer_addr_len = sizeof( peer_addr );
	res = getsockname( fd, &peer_addr.addr, &peer_addr_len );
	if( res < 0 ) {
		fprintf( stderr, "getsockname failed: %s (%d)\n", strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
		goto FAIL;
	}
	res = getnameinfo( &peer_addr.addr, peer_addr_len,
		peer_hostname, NI_MAXHOST,
		peer_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, "Connected to %s, port %s (peer: %s, %s)\n", host, service,
		peer_hostname, peer_service );

	if( !wolf_network_sock_nonblocking( STDIN_FILENO ) ) {
		fprintf( stderr, "set nonblocking failed for stdin: %s (%d)\n", strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
		goto FAIL;
	}
	if( !wolf_network_sock_nonblocking( STDOUT_FILENO ) ) {
		fprintf( stderr, "set nonblocking failed for stdout: %s (%d)\n", strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
		goto FAIL;
	}

	idle_secs = 0;
	count = 0;
	stdout_to_write = 0;
	stdout_written = 0;
	fd_to_write = 0;
	fd_written = 0;
	total_stdin_read = 0;
	total_fd_read = 0;
	total_fd_written = 0;
	total_stdout_written = 0;
	stdin_eof = false;
	fd_write_closed = false;
	fd_eof = false;

#ifndef _WIN32
	signal( SIGTERM, term_handler );
	signal( SIGINT, term_handler );
	signal( SIGPIPE, SIG_IGN );
#else
	/* TODO: install Ctrl-C handler as in service */
#endif
	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;

		/* set up select mask and compute highest file descriptor number */
		max_fd = 0;
		max_fd = max( max_fd, STDIN_FILENO );
		max_fd = max( max_fd, STDOUT_FILENO );
		max_fd = max( max_fd, fd );

		if( !stdin_eof && fd_to_write == 0 ) {
			FD_SET( STDIN_FILENO, &read_set );
		}
		if( stdout_to_write > 0 ) {
			FD_SET( STDOUT_FILENO, &write_set );
		}
		if( !fd_eof && stdout_to_write == 0 ) {
			FD_SET( fd, &read_set );
		}
		if( fd_to_write > 0 ) {
			FD_SET( fd, &write_set );
		}

		/* wait for events or timeout */
#if DEBUG
fprintf( stderr, "select call read_fd: %s write_fd: %s wbuf: %d rbuf: %d stdin_eof: %d fdwcl: %d\n",
	fd_set_to_string( &read_set, 0, max_fd, buf, 10 ),
	fd_set_to_string( &write_set, 0, max_fd, buf2, 10 ),
	fd_to_write, stdout_to_write, stdin_eof, fd_write_closed );
#endif

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

#if DEBUG
fprintf( stderr, "select %04d read_fd: %s write_fd: %s wbuf: %d rbuf: %d stdin_eof: %d fdwcl: %d\n",
	res, fd_set_to_string( &read_set, 0, max_fd, buf, 10 ),
	fd_set_to_string( &write_set, 0, max_fd, buf2, 10 ),
	fd_to_write, stdout_to_write, stdin_eof, fd_write_closed );
#endif

		if( res < 0 ) {
			if( WOLF_NETWORK_ERRNO == EINTR || WOLF_NETWORK_ERRNO == EAGAIN ) {
				/* skip */
			} else {
				/* fatal errors */
				fprintf( stderr, "select failed: %s (%d)\n",
					strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_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;

			/* empty the socket read buffer to stdout */
			if( FD_ISSET( STDOUT_FILENO, &write_set ) ) {
				assert( stdout_to_write > 0 );
				rres = write( STDOUT_FILENO, read_buffer + stdout_written, stdout_to_write );
				if( rres < 0 ) {
					if( WOLF_NETWORK_ERRNO == EAGAIN || WOLF_NETWORK_ERRNO == WOLF_NETWORK_EWOULDBLOCK || WOLF_NETWORK_ERRNO == EINTR ) {
						/* skip */
					} else {
						fprintf( stderr, "write to stdout failed: %s (%d)\n",
							strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
						goto FAIL;
					}
				} else {
					assert( rres >= 0 );
					stdout_written += (size_t)rres;
					assert( stdout_to_write >= (size_t)rres );
					stdout_to_write -= (size_t)rres;
					total_stdout_written += (size_t)rres;
				}
			}

			/* empty the socket write buffer writing to fd */
			if( FD_ISSET( fd, &write_set ) ) {
				assert( fd_to_write > 0 );
				rres = write( fd, write_buffer + fd_written, fd_to_write );
				if( rres < 0 ) {
					if( WOLF_NETWORK_ERRNO == EAGAIN || WOLF_NETWORK_ERRNO == WOLF_NETWORK_EWOULDBLOCK || WOLF_NETWORK_ERRNO == EINTR ) {
						/* skip */
					} else {
						fprintf( stderr, "write to socket failed: %s (%d)\n",
							strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_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;
					total_fd_written += (size_t)rres;
				}
			}

			/* read from socket, fill socket read buffer */
			if( FD_ISSET( fd, &read_set ) ) {
				assert( stdout_to_write == 0 );
				rres = read( fd, read_buffer, READ_BUFFER_SIZE );
				if( rres < 0 ) {
					if( WOLF_NETWORK_ERRNO == EAGAIN || WOLF_NETWORK_ERRNO == WOLF_NETWORK_EWOULDBLOCK || WOLF_NETWORK_ERRNO == EINTR ) {
						/* skip */
					} else {
						fprintf( stderr, "read from socket failed: %s (%d)\n",
							strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
						goto FAIL;
					}
				} else if( rres == 0 ) {
					/* EOF on socket */
					shutdown( fd, SHUT_RD );
					fd_eof = true;
#if DEBUG
					fprintf( stderr, "EOF on socket\n" );
#endif
				} else {
					assert( rres > 0 );
					stdout_written = 0;
					stdout_to_write = (size_t)rres;
					total_fd_read += (size_t)rres;
				}
			}

			/* read from stdin, fill socket write buffer */
			if( FD_ISSET( STDIN_FILENO, &read_set ) ) {
				assert( fd_to_write == 0 );
				rres = read( STDIN_FILENO, write_buffer, WRITE_BUFFER_SIZE );
				if( rres < 0 ) {
					if( WOLF_NETWORK_ERRNO == EAGAIN || WOLF_NETWORK_ERRNO == WOLF_NETWORK_EWOULDBLOCK || WOLF_NETWORK_ERRNO == EINTR ) {
						/* skip */
					} else {
						fprintf( stderr, "read from stdin failed: %s (%d)\n",
							strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
						goto FAIL;
					}
				} else if( rres == 0 ) {
					/* EOF on STDIN */
					stdin_eof = true;

#if DEBUG
					fprintf( stderr, "EOF on stdin\n" );
#endif
				} else {
					assert( rres > 0 );
					fd_written = 0;
					fd_to_write = (size_t)rres;
					total_stdin_read += (size_t)rres;
				}
			}

			/* close writing side of the socket if we have emptied
			 * the socket write buffer and there is nothing more on
			 * the input (stdin)
			 */
			if( fd_to_write == 0 && stdin_eof && !fd_write_closed ) {
#if DEBUG
				fprintf( stderr, "socket SHUT_WR\n" );
#endif
				shutdown( fd, SHUT_WR );
				fd_write_closed = true;
			}
		}

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

END:
	fprintf( stderr, "Terminated stdin: %zd, stdout: %zd, fd-in: %zd, fd-out: %zd\n",
		total_stdin_read, total_stdout_written, total_fd_read, total_fd_written );

	res = wolf_network_closesocket( fd );
	if( res < 0 ) {
		fprintf( stderr, "close failed: %s (%d)\n", strerror( WOLF_NETWORK_ERRNO ), WOLF_NETWORK_ERRNO );
		goto FAIL;
	}
 
	goto OK;

FAIL:

	if( fd >= 0 ) wolf_network_closesocket( fd );
	if( result != NULL ) freeaddrinfo( result );


#ifdef _WIN32
	WSACleanup( );
#endif

	return EXIT_FAILURE;

OK:

	return EXIT_SUCCESS;
}