/*
 * LINUX version.
 * Compile with:
 *	g++ -Wall -Werror -o socket socket.cpp
 *
 * --------------------------------------------------------------------
 *
 * Socket test code by Oliver "kfs1" Smith <oliver@kfs.org>
 *
 * See http://kfsone.wordpress.com/2011/09/06/sockets-data-black-holes/
 *
 * In short:
 *
 *  Client fills a TCP send buffer of X bytes with data to a
 *  server with a receive buffer of Y bytes (where Y < half X).
 *  The write is doing synchronously (blocking).
 *  When write returns, data has been copied into the send buffer
 *  but not all of it has been transmitted yet due to the size
 *  of the server's much smaller receive buffer.
 *
 *  What happens, then, if the client app now calls setsockopt
 *  to reduce it's send buffer?
 *
 *  - Generate an error to indicate that the resulting buffer
 *   would not be large enough to contain the existing queue?
 *  - Apply the new size only to additional writes?
 *  - Resize the current buffer and lose some of the data?
 *
 */

// ---------- Includes.

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

// ---------- Configuration.

// Port we will be using.
static const int portno = 8899 ;

// A message sent to help with synchronization.
static const size_t helloMessageSize = 64 ;

// How large should the server's receive buffer be (bytes).
static const size_t serverReceiveBufferSize = 2048 ;

// How large should the send buffer be?
static const size_t clientSendBufferSize = serverReceiveBufferSize * 8 ;

// How large should the big message be?
// (Make it smaller than the send buffer to avoid write errors)
static const size_t bigMessageSize = clientSendBufferSize - 256 ;

// Message we send at the end.
static const char finMessage[] = "[fin]" ;
// How large is that message?
static const size_t finMessageSize = sizeof(finMessage) + 1 ;

// ---------- Cheap & cheerful global buffer.

// A buffer for the data.
static char bytes[bigMessageSize] = "" ;

//////////////////////////////////////////////////////////////////////
// Helper: Display usage.

static void
_usageAndExit()
{
	printf("Usage: socket client|server\n") ;
	exit(-1) ;
}

//////////////////////////////////////////////////////////////////////
// Helper: Report an error and terminate.

static void
_die(const char* const reason)
{
	printf("ERROR: %s\n", reason) ;
	exit(-1) ;
}

//////////////////////////////////////////////////////////////////////
// Validate a read/write operation result.

static void
_validateOperation(const char* const operation, const char* const purpose, const int retval, const unsigned int size)
{
	if ( retval < 0 )
	{
		printf("ERROR: %s %s failed: %d, errno: %d\n", purpose, operation, retval, errno) ;
		exit(-1) ;
	}
	if ( (unsigned int)retval < size )
	{
		printf("ERROR: %s %s short: %d of %d bytes\n", purpose, operation, retval, size) ;
		exit(-1) ;
	}
	printf("- %s %s: %d bytes\n", purpose, operation, retval) ;
}

//////////////////////////////////////////////////////////////////////
// Send "size" bytes through the given socket.

static void
_write(const int socket, const unsigned int size, const char* const purpose)
{
	const int retval = write(socket, bytes, size) ;
	_validateOperation("write", purpose, retval, size) ;
}

//////////////////////////////////////////////////////////////////////
// Read upto "size" bytes through the given socket.

static void
_read(const int socket, const unsigned int size, const char* const purpose)
{
	const int retval = read(socket, bytes, size) ;
	_validateOperation("read", purpose, retval, size) ;
}

//////////////////////////////////////////////////////////////////////
// Query the operating system's current value for the send buffer.

static void
_printSndBufSize(const int socket)
{
	int bufSize ;
	socklen_t optlen = sizeof(bufSize) ;
	getsockopt(socket, SOL_SOCKET, SO_SNDBUF, &bufSize, &optlen) ;
	printf("- getsockopt says bufsize = %u (BSD sockets double the requested buf size)\n", bufSize) ;
}

//////////////////////////////////////////////////////////////////////
// Test code.

int
main(const int argc, const char* const argv[])
{
	bool isServer = false ;

	if ( argc != 2 )
		_usageAndExit() ;

	if ( strcmp(argv[1], "server") == 0 )
	{
		// Act as server.
		isServer = true ;
	}
	else if ( strcmp(argv[1], "client") == 0 )
	{
		// Act as client.
		isServer = false ;
	}
	else
	{
		// Didn't recognize arguments.
		_usageAndExit() ;
	}

	struct hostent* const server = gethostbyname("localhost") ;
	if ( server == NULL )
		_die("unable to resolve 'localhost'") ;

	// Construct the sockaddr_in with localhost and our desired port number.
	struct sockaddr_in serv_addr ;
	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) ;

	// Create a TCP socket.
	const int s = socket(AF_INET, SOCK_STREAM, 0) ;
	if ( s < 0 )
		_die("socket failed") ;

	if ( isServer == true )
	{
		/* Server code */

		// Allow the listen address to be reused when we exit.
		int reuse = 1 ;
		setsockopt(s, SOL_SOCKET, SO_REUSEADDR, (char *)&reuse, sizeof(reuse));

		// Give ourselves a smallish receive buffer.
		int bufSize = serverReceiveBufferSize ;
		int ssorv = setsockopt(s, SOL_SOCKET, SO_RCVBUF, &bufSize, sizeof(bufSize)) ;
		if ( ssorv < 0 )
			_die("setsockopt failed" ) ;

		// Set the socket up for listening.
		if ( bind(s, (struct sockaddr*) &serv_addr, sizeof(serv_addr)) < 0 )
			_die("bind failed") ;
		listen(s, 2) ;

		sockaddr_in cli_addr ;
		socklen_t val = sizeof(serv_addr) ;

		printf("- Listening\n") ;

		// Get the incoming connection.
		const int c = accept(s, (struct sockaddr*) &cli_addr, &val) ;
		if ( c < 0 )
			_die("accept failed") ;

		printf("- Connected\n") ;

		// Receive the client's handshake.
		_read(c, helloMessageSize, "handshake") ;

		if ( bytes[0] != 123 || bytes[helloMessageSize-1] != 123 )
			_die("handshake didn't have the right data") ;

		_write(c, helloMessageSize, "handshake response") ;

		// Give the client chance to receive our response and
		// start sending the big message; we used a read to
		// block until data starts to arrive, then we sleep to
		// force our very small receive buffer to fill.
		static const size_t slurpSize = 100 ;
		_read(c, slurpSize, "first 1k of big message") ;

		static const unsigned int delay_seconds = 2 ;
		printf("~~ sleeping for %u seconds ~~\n", delay_seconds) ;
		sleep(delay_seconds) ;

		static const size_t remainderSize = bigMessageSize - slurpSize ;
		_read(c, remainderSize, "remainder of big message") ;

		// Send the fin message.
		strcpy(bytes, finMessage) ;
		_write(c, 8, "FIN") ;

		shutdown(c, SHUT_RDWR) ;
		close(c) ;
	}
	else
	{
		/* Client code */
		printf("- Connecting\n") ;
		if ( connect(s, (struct sockaddr*) &serv_addr, sizeof(serv_addr)) < 0 )
			_die("connect failed") ;

		// Give ourselves the big send buffer.
		int bufSize = clientSendBufferSize ;
		int ssorv = setsockopt(s, SOL_SOCKET, SO_SNDBUF, &bufSize, sizeof(bufSize)) ;
		if ( ssorv < 0 )
			_die("setsockopt failed") ;
		_printSndBufSize(s) ;

		printf("- Connected\n") ;

		// Handshake.
		bytes[0] = bytes[helloMessageSize-1] = 123 ;
		_write(s, helloMessageSize, "handshake") ;
		_read(s, helloMessageSize, "handshake reply") ;

		bytes[0] = bytes[helloMessageSize-1] = 0 ;

		// Now write the big buffer. This puts the data into
		// the socket's send buffer, but the server doesn't
		// read it all (to simulate network transit time etc)
		// so the bulk of it is still sitting in the send buffer.
		_write(s, bigMessageSize, "big write") ;

		// Remember: the bulk of the data is still in the
		// send buffer. So what happens if we shorten the send buffer?
		bufSize = serverReceiveBufferSize ;
		ssorv = setsockopt(s, SOL_SOCKET, SO_SNDBUF, &bufSize, sizeof(bufSize)) ;
		// Did it produce an error? Did it set errno to something?
		printf("reducing buf size returned %d errno = %d\n", ssorv, errno) ;

		// Did it correctly change the send buffer size?
		_printSndBufSize(s) ;

		// See if we can get those last 8 bytes.
		_read(s, finMessageSize, "FIN") ;
		if ( strcmp(bytes, finMessage) != 0 )
			_die("'fin' message was bad.\n") ;

		printf("Done.\n") ;
	}

	shutdown(s, SHUT_RDWR) ;
	close(s) ;

	return 0 ;
}