chat_tcp_server.c

linux 2.6内核下编写的一个聊天程序

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <errno.h>
#include <arpa/inet.h>

#include "buffer.h"

#define LISTENQ 5
#define BUFFER_SIZE 1024
#define BUFFER_MAXSIZE 4096

////////////////////////////////////////////////////////////////////////////////
//
////////////////////////////////////////////////////////////////////////////////

typedef struct connection
{
  int fd;

  buffer_t *rbuf;
  buffer_t *wbuf;

} connection_t;

////////////////////////////////////////////////////////////////////////////////
// the global defination
////////////////////////////////////////////////////////////////////////////////

fd_set global_read_set;
fd_set global_write_set;
fd_set global_except_set;

int global_max_fd = 0;

connection_t **global_connection_list = NULL;
int global_connection_count = 0;
int global_support_connection_num;

int network_init (int n);
connection_t * connection_create (int fd);
void connection_destroy (connection_t * c);
void network_finial (void);
void network_register_read (int fd);
void network_unregister_read (int fd);
void network_register_write (int fd);
void network_unregister_write (int fd);
void network_register_except (int fd);
void network_unregister_except (int fd);
void data_exchange(int src);

////////////////////////////////////////////////////////////////////////////////
// the main function
////////////////////////////////////////////////////////////////////////////////
int
main (int argc, char **argv)
{
  if (argc < 3)
    {
      fprintf (stderr, "Usage:%s <IP><PORT>\n", strerror (errno));
      exit (1);
    }
  //XXX: step 1 create a socket
  int server_listen_socket;

  if ((server_listen_socket = socket (PF_INET, SOCK_STREAM, 0)) < 0)
    {
      fprintf (stderr, "socket() failed:%s\n", strerror (errno));
      exit (1);
    }
  //XXX: step 2 create a bind
  int on;
  on = 1;

  if (setsockopt
      (server_listen_socket, SOL_SOCKET, SO_REUSEADDR, &on, sizeof (on)) < 0)
    {
      fprintf (stderr, "setsockopt failed:%s\n", strerror (errno));
      exit (1);
    }
#ifdef _SO_REUSEPORT_
  on = 1;
  if (setsockopt
      (server_listen_socket, SOL_SOCKET, SO_REUSEPORT, &on, sizeof (on)) < 0)
    {
      fprintf (stderr, "setsockopt failed:%s\n", strerror (errno));
      exit (1);
    }
#endif
  struct sockaddr_in server_address;
  memset (&server_address, 0, sizeof (server_address));

  server_address.sin_family = PF_INET;
  server_address.sin_port = htons (atoi (argv[2]));
  server_address.sin_addr.s_addr = inet_addr (argv[1]);

  if (bind
      (server_listen_socket, (struct sockaddr *) &server_address,
       sizeof (server_address)) < 0)
    {
      fprintf (stderr, "bind() failed:%s\n", strerror (errno));
      exit (1);
    }
  //XXX: step 3: create a listen
  if (listen (server_listen_socket, LISTENQ) < 0)
    {
      fprintf (stderr, "listen() failed:%s\n", strerror (errno));
      exit (1);
    }

  //XXX: below is the function of select
  network_init (1024);

  network_register_read (server_listen_socket);

  struct timeval to;

  int n;

  fd_set rset;
  fd_set wset;
  fd_set eset;

  //XXX: the main loop

  for (;;)
    {
      rset = global_read_set;
      wset = global_write_set;
      eset = global_except_set;

      to.tv_sec = 1;
      to.tv_usec = 0;

      if ((n = select (global_max_fd + 1, &rset, &wset, &eset, &to)) < 0)
	{
	  if (errno == EINTR)
	    {
	      fprintf (stdout, "closed by signal,continue.....\n");
	      continue;
	    }
	  else if (errno == EAGAIN)
	  {
		  continue;
	  }
	  else
	    {
	      fprintf (stderr, "select() failed:%s\n", strerror (errno));
	      exit (1);
	    }
	}
      else if (n == 0)
	{
	  fprintf (stdout, "timeout......\n");
	}
      else
	{
	  fprintf (stdout,
		   "%d fd is waiting for reading or writing or excepting\n",
		   n);

	  int i;

	  for (i = 0; i <= global_max_fd; i++)
	    {
	      if (FD_ISSET (i, &rset))
		{
		  if (i == server_listen_socket)
		    {
		      fprintf (stdout, "fd %d is ready for accepting\n", i);
		      //XXX: below is the accept

		      int server_accept_socket;

		      struct sockaddr_in peer_address;
		      socklen_t peer_address_length;

		      peer_address_length = sizeof (peer_address);

		      if ((server_accept_socket =
			   accept (server_listen_socket,
				   (struct sockaddr *) &peer_address,
				   &peer_address_length)) < 0)
			{
			  fprintf (stderr, "accept() failed:%s\n",
				   strerror (errno));
			  exit (1);
			}
		      else
			{
			  fprintf (stdout,
				   "New tcp connection %d from (%s:%d)\n",
				   server_accept_socket,
				   inet_ntoa (peer_address.sin_addr),
				   ntohs (peer_address.sin_port));

			  network_register_read (server_accept_socket);

			  connection_t *conn;

			  //TODO: the connect_create
			  conn = connection_create(server_accept_socket);
			  global_connection_list[server_accept_socket] = conn;

			}

		    }
		  else
		    {
		      //can do read
		      ssize_t readn;
		      char buffer[BUFFER_SIZE];
		      fprintf(stdout, "fd %d is ready for reading\n", i);
		    again:
		      if ((readn = read (i, buffer, BUFFER_SIZE)) < 0)
			{
			  if (errno == EINTR)
			    {
			      goto again;
			    }
			  else
			    {
			      fprintf (stderr, "read error on socket %d\n",
				       i);
			    }
			}
		      else if (readn == 0)
			{
			  fprintf (stdout, "connection was closed by peer\n");
			  close (i);
			  network_unregister_read (i);
			  network_unregister_write (i);
			  network_unregister_except (i);

			  connection_destroy(global_connection_list[i]);
			  global_connection_list[i] = NULL;
			}
		      else
			{
			  buffer[readn] = '\0';
			  fprintf (stdout, "read %d bytes on socket %d,%s\n",
				   readn, i, buffer);
			  buffer_append_data(global_connection_list[i]->rbuf,
					      buffer, readn);

			  data_exchange(i);
			}
		    }
		}
	      //TODO: check the writable
	      if (FD_ISSET (i, &wset))
		{
		  fprintf (stdout, "fd %d is ready for writing\n", i);

		  ssize_t written;
		retry:
		  if ((written =
		       write (i, global_connection_list[i]->wbuf->p,
			      global_connection_list[i]->wbuf->length)) < 0)
		    {
		      if (errno == EINTR)
			{
			  goto retry;
			}
		      else
			{
			  fprintf (stderr, "write failed on socket:%d:%s\n",
				   i, strerror (errno));
			}
		    }
		  else
		    {
		      connection_t *current;
		      current = global_connection_list[i];
		      buffer_remove_data (current->wbuf, written);
		      if (current->wbuf->length == 0)
			{
			  network_unregister_write(i);
			}
		    }
		}
	    }
	}
    }

  close (server_listen_socket);
  network_finial();

  return 0;
}


int
network_init (int n)
{
  global_support_connection_num = n;

  if ((global_connection_list =
       malloc (global_support_connection_num * sizeof (connection_t *))) ==
      NULL)
    {
      fprintf (stderr, "malloc() failed:%s\n", strerror (errno));
      exit (1);
    }

  int i;

  for (i = 0; i < global_support_connection_num; i++)
    {
      global_connection_list[i] = NULL;
    }
  global_connection_count = 0;

  FD_ZERO (&global_read_set);
  FD_ZERO (&global_write_set);
  FD_ZERO (&global_except_set);

  return 0;

}

connection_t * connection_create (int fd)
{
  connection_t *conn;

  if ((conn = (connection_t *) malloc (sizeof (connection_t))) == NULL)
    {
      fprintf (stderr, "malloc failed:%s\n", strerror (errno));
      return NULL;
    }

  conn->fd = fd;

  conn->rbuf = buffer_create (BUFFER_SIZE, BUFFER_MAXSIZE);
  conn->wbuf = buffer_create (BUFFER_SIZE, BUFFER_MAXSIZE);

  return conn;
}

void connection_destroy (connection_t * c)
{
  if (c)
    {
      if (c->rbuf)
	{
	  buffer_destory (c->rbuf);
	  c->rbuf = NULL;
	}
      if (c->wbuf)
	{
	  buffer_destory (c->wbuf);
	  c->wbuf = NULL;
	}

      free (c);
      c = NULL;

    }
}

void network_finial (void)
{
  int i;
  for (i = 0; i < global_support_connection_num; i++)
    {
      if (global_connection_list[i] != NULL)
	{
	  connection_destroy (global_connection_list[i]);
	}
    }
}

void network_register_read (int fd)
{
  FD_SET (fd, &global_read_set);

  if (fd > global_max_fd)
    {
      global_max_fd = fd;
    }
}
void network_unregister_read (int fd)
{
  FD_CLR (fd, &global_read_set);
}

void network_register_write (int fd)
{
  FD_SET (fd, &global_write_set);

  if (fd > global_max_fd)
    {
      global_max_fd = fd;
    }
}
void network_unregister_write (int fd)
{
  FD_CLR (fd, &global_write_set);
}

void network_register_except (int fd)
{
  FD_SET (fd, &global_except_set);

  if (fd > global_max_fd)
    {
      global_max_fd = fd;
    }
}
void network_unregister_except (int fd)
{
  FD_CLR (fd, &global_except_set);
}

void data_exchange(int src)
{
  int i;
  char *data;
  unsigned long n;

  connection_t *target;
  connection_t *source;

  source = global_connection_list[src];

  // FIXME: check return value
  //char *buffer_fetch_data(buffer_t * buf, unsigned long length, unsigned long *rlength)
  data = buffer_fetch_data(source->rbuf, source->rbuf->length, &n);

  if (data)
  {
    for (i = 0; i < global_support_connection_num; i++)
    {
      if (i != src && global_connection_list[i] != NULL)
      {
	target = global_connection_list[i];

	// FIXME: check return value
	//int buffer_append_data(buffer_t * buf, char *data, unsigned long length)
	buffer_append_data(target->wbuf, data, n);
	network_register_write(i);
      }
    }

    free(data);
  }
}