//-------------------------------------------------------------------
//	tweaktos.cpp
//
//	This program utilizes an Internet Protocol socket-option to
//	adjust the Type-of-Service field in an outgoing UDP packet; 
//	it uses the 'sendmsg()' function to receive ancillary data.
//
//	    compile using:  $ g++ tweaktos.cpp -o tweaktos
//	    execute using:  $ ./tweaktos <hostname> <port> <tos>
//
//	programmer: ALLAN CRUSE
//	written on: 15 FEB 2009
//-------------------------------------------------------------------

#include <netdb.h>	// for gethostbyname()
#include <stdio.h>	// for printf(), perror() 
#include <stdlib.h>	// for exit() 
#include <string.h>	// for strncpy()
#include <arpa/inet.h>	// for inet_ntoa()

int main( int argc, char **argv )
{
	//-----------------------------------------------
	// check for the required command-line arguments
	//----------------------------------------------- 
	if ( argc < 4 ) 
		{
		printf( "usage: %s <hostname> <port> <tos> \n", argv[0] );
		exit(1);
		}

	//----------------------------------------------
	// get the hostname for the destination station
	//---------------------------------------------- 
	char	peername[ 64 ] = {0}; 
	strncpy( peername, argv[ 1 ], 63 );

	//-------------------------------------------
	// get the port-number for the socket to use 
	//-------------------------------------------
	unsigned short	port = atoi( argv[ 2 ] );

	//---------------------------------------------
	// get the value for the Type-of-Service field
	//--------------------------------------------- 
	long	tos = strtol( argv[ 3 ], NULL, 0 )&0xFF;
	int	len = sizeof( tos );

	//-------------------------------------
	// get the peer's IPv4 network address
	//------------------------------------- 
	char	peeraddr[ 16 ] = {0};
	struct hostent	*pp = gethostbyname( peername );
	if ( !pp ) { herror( "gethostbyname" ); exit(1); }
	strcpy( peeraddr, inet_ntoa( *(in_addr*)pp->h_addr ) );
	
	//----------------------------------------------------
	// initialize an internet socket-address for the peer 
	//---------------------------------------------------- 
	struct sockaddr_in	paddr = { 0 };
	socklen_t		palen = sizeof( paddr );
	paddr.sin_family 	= AF_INET;
	paddr.sin_port 		= htons( port );
	paddr.sin_addr.s_addr 	= *(uint32_t*)pp->h_addr;

	//------------------------------------------------------
	// ask the kernel to create an internet datagram socket
	//------------------------------------------------------ 
	int	sock = socket( AF_INET, SOCK_DGRAM, IPPROTO_IP );
	if ( sock < 0 ) { perror( "socket" ); exit(1); }	

	//--------------------------------------------------------
	// set the value of this socket's 'Type-of-Service' field 
	//--------------------------------------------------------
	if ( setsockopt( sock, SOL_IP, IP_TOS, &tos, len ) < 0 )
		{ perror( "setsockopt TOS" ); exit(1); }

	//---------------------------------------------
	// create the message-string to be transmitted
	//--------------------------------------------- 
	char	msg[] = " ABCDEFGHIJKLMNOPQRSTUVWXYZ 1234567890 *";
	int	mlen = strlen( msg );

	//--------------------------------------------------
	// try to send this message to the destination host 
	//-------------------------------------------------- 
	int	tx = sendto( sock, msg, mlen, 0, (sockaddr*)&paddr, palen ); 
	if ( tx < 0 ) { perror( "sendto" ); exit(1); }

	//----------------------------------------------------------------
	// show confirmation of the message's size, port, target, and TOS
	//---------------------------------------------------------------- 
	printf( " sent %d bytes toward port %d", tx, port );
	printf( " on \'%s\' (%s)", peername, peeraddr );
	printf( " with TOS=0x%02X \n", tos );	

	//------------------------------------------------------------
	// now erase the outgoing message so we can reuse its storage
	//------------------------------------------------------------ 
	bzero( msg, mlen );

	//-----------------------------------------------------------
	// use the 'recvmsg()' function to get the server's response
	// together with ancillary data which includes the TOS value
	//----------------------------------------------------------- 
	int	oval = 1;
	int	olen = sizeof( oval );
	if ( setsockopt( sock, SOL_IP, IP_RECVTOS, &oval, olen ) < 0 )
		{ perror( "setsockopt RECVTOS" ); exit(1); }

	char	cbuf[ 80 ] = {0};
	struct iovec	myiov[1] = { { msg, mlen } };
	
	struct msghdr	mymsghdr;
	mymsghdr.msg_name 	= &paddr;
	mymsghdr.msg_namelen 	= sizeof( palen );
	mymsghdr.msg_iov	= myiov;
	mymsghdr.msg_iovlen	= 1;
	mymsghdr.msg_control	= &cbuf;
	mymsghdr.msg_controllen	= sizeof( cbuf );
	mymsghdr.msg_flags	= 0;

	int	rx = recvmsg( sock, &mymsghdr, 0 );
	if ( rx < 0 ) { perror( "recvmsg" ); exit(1); }

	//-------------------------------------------------------------
	// extract the 'Type-of-Service' value from the ancillary data
	//------------------------------------------------------------- 
	int	rxtos = 0;
	struct cmsghdr	*cmsg;
	for ( cmsg = CMSG_FIRSTHDR( &mymsghdr ); cmsg != NULL;
		cmsg = CMSG_NXTHDR( &mymsghdr, cmsg ) )
		{
		if (( cmsg->cmsg_level == SOL_IP )
			&&( cmsg->cmsg_type == IP_TOS ))
			memcpy( &rxtos, CMSG_DATA( cmsg ), 1 );
		}

	//-----------------------------------------------------------
	// report a message's arrival and show its 'Type-of-Service'
	//-----------------------------------------------------------
	strcpy( peeraddr, inet_ntoa( *(in_addr*)&paddr.sin_addr ) );
	printf( "\n received %d bytes from %s", rx, peeraddr );
	printf( " with TOS=0x%02X \n", rxtos );
}