/* Operating Systems Spring 2005
   Sample TCP socket code

   This program can run in one of two modes:
     1. as a server that creates a passive TCP socket
        and waits for connections. Once it gets a connection,
        the server simply reads from the socket and writes
	to stdout until an EOF is reached (indicating that the
	client has closed it's end of the connection).

		To run this program as a server:
		  > ./tcpplay -s portnum

     2. as a client. The client connects to the specified
        server (hostname:port entered on the command line).
	The client then reads a line at a time from stdin
	(using readline), and sends each line to the server.
	The client will quit if it sees EOF from stdin, or
        if it reads the string "quit".

		To run this program as a client:
		  > ./tcpplay -c server:port

     The server() and client() functions defined in the file common.c
     are called to do the actual work of transferring data between
     processes, the code here just sets up the socket descriptor before
     calling one of those functions.

*/

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h> /* socket specific definitions */
#include <netinet/in.h> /* INET constants and stuff */
#include <arpa/inet.h>  /* IP address conversion stuff */
#include <netdb.h>	/* gethostbyname */

#include "common.h"  /* client() and server() functions */


/* here is code that attempts to create a tcp socket
   and connect to the server specified.
   No data is read or sent, this just creates the 
   socket and connects to the server

   If everything works, this function returns a socket descriptor.
   Any errors are considered fatal (program exits).
*/

int create_tcp_connection(char *server_hostname,int port) {
  int sd;
  struct sockaddr_in skaddr;
  struct hostent *hp;           /* used for name lookup */

  /* create a socket
     IP protocol family (PF_INET)
     TCP protocol (SOCK_STREAM)
  */

  if ((sd = socket( PF_INET, SOCK_STREAM, 0 )) < 0) {
    perror("Problem creating socket");
    exit(1);
  }

  /* fill in an address structure that will be used to specify
     the address of the server we want to connect to
     address family is IP  (AF_INET)
     server IP address is found by calling gethostbyname with the
     name of the server 
  */
  skaddr.sin_family = AF_INET;

  /* convert hostname to a network byte order binary IP address */
  /* First try to convert using gethostbyname */
  if ((hp = gethostbyname(server_hostname))!=0) {

    /* Name lookup was successful - copy the IP address */
    memcpy( &skaddr.sin_addr.s_addr, hp->h_addr, hp->h_length);

  } else {
 
    /* Name lookup didn't work, try converting from dotted decimal */

    if (inet_aton(server_hostname,&skaddr.sin_addr)==0) {
      printf("Invalid IP address: %s\n",server_hostname);
      exit(1);
    }
  }
 
  /* port must be in network byte order! */
  skaddr.sin_port = htons(port);

  /* attempt to establish a connection with the server */
  if (connect(sd,(struct sockaddr *) &skaddr,sizeof(skaddr)) < 0 ) {
    perror("Problem connecting to server");
    exit(1);
  }
                                        
  return(sd);
}

/* set up a passive server socket on the specified port.
   This returns the socket descriptor (any errors are 
   considered fatal).
*/

int create_listening_socket(int port) {
  struct sockaddr_in skaddr;
  int ld;
  int len;
 
  /* create a socket
       IP protocol family (PF_INET)
       TCP protocol (SOCK_STREAM)
  */
   
  if ((ld = socket( PF_INET, SOCK_STREAM, 0 )) < 0) {
    perror("Problem creating socket\n");
    exit(1);
  }
   
  /* establish our address
   address family is AF_INET
   our IP address is INADDR_ANY (any of our IP addresses)
   the port must be in network byte order!
  */
 
  skaddr.sin_family = AF_INET;
  skaddr.sin_addr.s_addr = htonl(INADDR_ANY);
  skaddr.sin_port = htons(0);
 
  if (bind(ld, (struct sockaddr *) &skaddr, sizeof(skaddr))<0) {
    perror("Problem binding\n");
    exit(0);
  }
  len = sizeof(struct sockaddr);
  getsockname(ld,(struct sockaddr*)&skaddr,&len);
  printf("Port number assigned is %d\n",ntohs(skaddr.sin_port)); 

  /* put the socket into passive mode (waiting for connections) */
  if (listen(ld,5) < 0 ) {
    perror("Error calling listen\n");
    exit(1);
  }
                                                  
  return(ld);
}



/* when acting as the server, this function is called.
   tcp_server waits for a connection, then calls the server() function.
   
*/

void tcp_server(int port) {
  int ld,client;
  int addrlen;
  struct sockaddr_in from;

  /* create passive mode (listening) socket on the specified port
  */

  ld = create_listening_socket(port);

  /* do forever (handles multiple clients - iteratively) */
  while (1) {
    printf("Ready for a connection...\n");
    addrlen=sizeof(from);
    if ( (client = accept( ld, (struct sockaddr*) &from, &addrlen)) < 0) {
      perror("Problem with accept call");
      exit(1);
    }
    /* call the server function to handle the connection */
    server(client);
  }

}

/* ======================================================
   tcp_client is called if this is run as a client.
   The client connects to the server, then
   reads from stdin (using readline),
   and unless the line is "quit",
   sends it to the server. Also quits if stdin is closed...
*/

void tcp_client(char *server_hostname,int port) {
  int sd;

  /* create a connected tcp socket */
  sd = create_tcp_connection(server_hostname,port);

  /* call the client function to handle this connection */
  client(sd);
}


enum {NONE=0,CLIENT, SERVER};

int main(int argc, char **argv) {
  int mode = NONE;
  char *server=NULL;
  int i;
  int port=0;
  char c;

  while ((c=getopt(argc,argv,"c:s:"))>=0) {
    switch(c) {
    case 'c':
	  mode = CLIENT; 
	  /* parse the string to get the server name and port */
	  server = (char *) malloc(strlen(optarg));
	  i=0;
	  while (*optarg && *optarg!=':') {
	    server[i++]=*optarg;
	    optarg++;
	  }
	  server[i]=0;

	  if (*optarg != ':') {
	    printf("Invalid server address, should be server:port\n");
	    exit(1);
	  }
	  port = atoi(optarg+1);
	  break;
    case 's':
	  mode = SERVER; 
	  port = atoi(optarg);
	  break;
    case '?':  /* invalid option - getopt returns a '?' */
      exit(1);
    }
  }

  if ((mode==NONE) || (port<=0)) {
	printf("Invalid command line. Usage: %s (-s port | -c server:port)\n",argv[0]);
	exit(0);
  }

  if (mode==SERVER) {
    tcp_server(port);
  } else {
    tcp_client(server,port);
  }

  return(1);
}