/* 
 * echoserv.c - a simple single-threaded TCP server test
 * by //YorHel
 *
 * Copyright 2006 Y. Heling,
 * February 2006
 * License: MIT
 *
 * Use it at your own risk
 */

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

#define LISTEN_PORT 1337

#define MAX(x,y) ((x) > (y) ? (x) : (y))
;
#define MAX_CONNECTIONS 10
#define READ_BUFFER_SIZE 512
                                     
#define CST_FREE 0
#define CST_READ 1
#define CST_WRITE 2

typedef struct {
  int fd;
  char state;
  char *buf;
} connections;
static connections *conns;

void error(const char *);
void close_and_free(int);

int main() {
  printf("Simple TCP server by //YorHel\n\n");

  printf("* Creating socket\n");
  int l;
  if((l = socket(AF_INET, SOCK_STREAM, 0)) < 0)
    error("Can't create socket");

  printf("* Setting SO_REUSEADDR on main socket\n");
  int set = 1;
  if(setsockopt(l, SOL_SOCKET, SO_REUSEADDR, (void *) &set, sizeof(set)) < 0)
    error("Can't set SO_REUSEADDR on main socket"); 

  printf("* Binding socket\n");
  struct sockaddr_in serv_addr;
  memset(&serv_addr, 0, sizeof(struct sockaddr_in));
  serv_addr.sin_family = AF_INET;
  serv_addr.sin_port = htons(LISTEN_PORT);
  serv_addr.sin_addr.s_addr = INADDR_ANY;
  if(bind(l, (struct sockaddr *) &serv_addr, sizeof(serv_addr)) < 0)
    error("Can't bind socket");
  
  printf("* Listening to socket\n");
  listen(l, 5);

  printf("* Creating connections table\n");
  conns = (connections *) malloc( sizeof(connections) * MAX_CONNECTIONS );
  if(conns == (connections *) 0)
    error("Can't allocate memory for connections table");
  int cnum;
  for(cnum=MAX_CONNECTIONS; cnum--;)
    conns[cnum].state = CST_FREE;
  
  printf("* Accepting connections\n");
 /* re-use the same vars over-and-over again
  * is a little faster than defining them each time */
  int n, cn, newid, clifd;
  char stop = 0;
  fd_set rd, wr;
  struct sockaddr_in cliaddr;
  unsigned int clilen;
 /* main loop */
  while(!stop) {
   /* set FDs for select() */
    FD_ZERO(&rd); FD_ZERO(&wr);
    n = 0;
    FD_SET(l, &rd);
    n = MAX(l, n);
   
   /* set FDs for the connections */
    for(cn=MAX_CONNECTIONS;cn--;) {
      if(conns[cn].state == CST_READ) {
        FD_SET(conns[cn].fd, &rd);
        n = MAX(n, conns[cn].fd);
      } else if(conns[cn].state == CST_WRITE) {
        FD_SET(conns[cn].fd, &wr);
        n = MAX(n, conns[cn].fd);
      }
    }

    n = select(n + 1, &rd, &wr, (fd_set *) NULL, (struct timeval *) 0);
    if(n < 0)
      error("select() failed...");
    if(n == 0)
      continue;
    
   /* Something happend, handle it :) */
    if(FD_ISSET(l, &rd)) { /* new connection */
     /* get free slot, if one */
      newid = -1;
      for(cn=MAX_CONNECTIONS;cn--;)
        if(conns[cn].state == CST_FREE) {
          newid = cn; break;
        }
      if(newid < 0)
        printf(" *WARNING: Too many connections\n");
      else {
       /* accept the connection */
        clilen = sizeof(cliaddr);
        memset(&cliaddr, 0, clilen);
        clifd = accept(l, (struct sockaddr *) &cliaddr, &clilen);
        if(clifd < 0)
          error("Can't accept connection");
        printf(" [%d] We have a connection!!\n", newid);
        conns[newid].state = CST_WRITE;
        if((conns[newid].buf = malloc(READ_BUFFER_SIZE)) < 0)
          error("Can't allocate memory");
        sprintf(conns[newid].buf, "Hello world!\n");
        conns[newid].fd = clifd;
      }
    }

   /* checking active sockets */
    for(cn=MAX_CONNECTIONS;cn--;) {
     /* we can write */
      if(conns[cn].state == CST_WRITE && FD_ISSET(conns[cn].fd, &wr)) {
        if(write(conns[cn].fd, conns[cn].buf, strlen(conns[cn].buf)) <= 0)
          close_and_free(cn);
        printf(" [%d] Sent: %s", cn, conns[cn].buf);
        conns[cn].state = CST_READ;
      }
     /* we can read */
      if(conns[cn].state == CST_READ && FD_ISSET(conns[cn].fd, &rd)) {
        memset((void *) conns[cn].buf, 0, READ_BUFFER_SIZE);
        if(read(conns[cn].fd, conns[cn].buf, READ_BUFFER_SIZE) <= 0)
          close_and_free(cn);
        printf(" [%d] G0t: %s", cn, conns[cn].buf);
        conns[cn].state = CST_WRITE;
        if(strstr(conns[cn].buf, "exit") == conns[cn].buf) {
          printf(" [%d] Closing connection\n", cn);
          close_and_free(cn);
        } else if(strstr(conns[cn].buf, "die()") == conns[cn].buf) {
          printf("Got die() from connection #%d, dying!\n", cn);
          stop = 1;
        }
      }
    }  
  }
 /* close all connections before dying */
  for(cnum=MAX_CONNECTIONS;cnum--;)
    if(conns[cnum].state != CST_FREE)
      close_and_free(cnum);
  close(l);


  return 0;
}
void close_and_free(int cn) {
  close(conns[cn].fd);
  conns[cn].state = CST_FREE;
  free(conns[cn].buf); 
}


void error(const char *msg) {
  perror(msg);
  exit(1);
}