/**************************************************************
 * connection-oriented server programming 
 * 
 * Function: iterative echo server using select 
 * 
 * Reference: Richard Stevens, Unix Network Programming, 1998
 *
 * Created by: Hairong Qi
 **************************************************************/
#include <stdio.h>        /* for standard I/O */
#include <unistd.h>       /* for read() and write() */
#include <strings.h>      /* for bzero() */
#include <stdlib.h>       /* for atoi() */
#include <sys/types.h>    /* for accept() and connect() */
#include <sys/socket.h>
#include <sys/errno.h>    /* for perror() */
#include <sys/select.h>   /* for select() */
#include <netinet/in.h>   /* for htonl() */

#define LISTENQ 2
#define BSIZE 256
#define USAGE "./iservertcp_select portno\n"

/**
 ** output the error message and exit
 **/
void error(char *msg)
{
  perror(msg);
  exit(1);
}


/**
 ** the main function 
 **/
int main(int argc, char *argv[])
{
  int listenfd, connfd;              /* socket file descriptor */
  int portno;
  socklen_t clilen;
  char buffer[BSIZE];
  struct sockaddr_in serv_addr, cli_addr;
  int n;

  /* variable declarations for select() */
  int i, maxi, maxfd, sockfd;
  int nready, client[FD_SETSIZE];
  fd_set rset, allset;               /* the bitmap */
 
  /* check the correctness of command-line inputs */
  if (argc < 2) 
    error(USAGE);

  /* create a listening socket */
  listenfd = socket(PF_INET, SOCK_STREAM, 0);
  if (listenfd < 0) 
    error("ERROR opening socket");

  /* specify local address */
  bzero((char *) &serv_addr, sizeof(serv_addr));
  serv_addr.sin_family = AF_INET;
  serv_addr.sin_addr.s_addr = htonl(INADDR_ANY); 
  portno = atoi(argv[1]);
  serv_addr.sin_port = htons(portno);

  /* bind the address to the listening socket */
  if (bind(listenfd, (struct sockaddr *) &serv_addr, sizeof(serv_addr)) < 0) 
    error("ERROR on binding");

  /* specify a queue length for the server */
  listen(listenfd, LISTENQ);

  /* start accepting new connections or checking existing connections */
  maxfd = listenfd;                  /* the maximum fd so far */
  maxi = -1;                         /* index to the client[] array */
  for (i=0; i<FD_SETSIZE; i++)       /* initialize the client array */
    client[i] = -1;
  FD_ZERO(&allset);                  /* initialize the descriptor set */
  FD_SET(listenfd, &allset);         /* set the corresponding bit to on */

  while (1) {
    /*
    ** select waits for something to happen:
    ** either the establishment of a new client connection 
    ** or the arrival of data, a FIN, or an RST on an existing connection
    */
    rset = allset;                   /* reset rset (a dynamic bitmap) */
    nready = select(maxfd+1, &rset, NULL, NULL, NULL);

    /*
    ** check if the listening socket has activity, if it is, then
    ** a new connection has been established, and accept needs to be called
    */
    if (FD_ISSET(listenfd, &rset)) {
      /* accept connections */
      clilen = sizeof(cli_addr);
      connfd = accept(listenfd, (struct sockaddr *) &cli_addr, &clilen);
      if (connfd < 0) 
	error("ERROR on accept");
      
      /* 
      ** update the client array
      ** look for the first unused entry (-1) and record the connected socket
      */
      for (i=0; i<FD_SETSIZE; i++)
	if (client[i] < 0) {
	  client[i] = connfd;
	  break;
	}
      if (i == FD_SETSIZE)
	error("ERROR too many clients");

      FD_SET(connfd, &allset);        /* update the descriptor set (a static bitmap) */
      if (connfd > maxfd)             /* maxfd always records the largest fd */
	maxfd = connfd;
      if (i > maxi)                   /* update the index to the client array */
	maxi = i;

      /*
      ** the number ready descriptors is decremented; and if it is 0,
      ** we can avoid the next for loop below.
      ** this lets us use the return value from select to avoid checking
      ** descriptors that are not ready
      */
      if (--nready <= 0) 
	continue;
    }
    /*
    ** check for existing connections
    ** if there is data readable on an existing socket, echo it back
    */
    for (i=0; i<=maxi; i++) {
      if ((sockfd = client[i]) < 0)
	continue;

      if (FD_ISSET(sockfd, &rset)) {
	/* start reading in data and echo back */
	bzero(buffer, BSIZE);
	n = read(sockfd, buffer, BSIZE-1);
	if (n < 0) 
	  error("ERROR reading from socket");
	if (n == 0) {                   /* receives a FIN package, close connection */
	  close(sockfd);                /* close the socket */
	  FD_CLR(sockfd, &allset);      /* clear off the descriptor set */
	  client[i] = -1;               /* reset that entry in client array */
	} else {                        /* if n>0, echo back */
	  printf("Here is the message: %s\n", buffer);
	  fflush(stdout);
	  n = write(sockfd, "I got your message", 18);
	  if (n < 0) 
	    error("ERROR writing to socket");
	  /*
	  ** if there's no more readable descriptors, then go back to select
	  */
	  if (--nready <= 0) 
	    break;
	}
      }
    }
  }

  return 0; 
}