tcp_chat_server.c

使用select实现的UNIX/Linux环境下的TCP聊天程序

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

#include "buffer.h"

//////////////////////////////////////////////////////////////////////
// Definitions
//////////////////////////////////////////////////////////////////////

#define LISTENQ		10
#define BUFFER_SIZE	1024
#define BUFFER_MAXSIZE	8192

//////////////////////////////////////////////////////////////////////
// Type Definitions
//////////////////////////////////////////////////////////////////////
typedef struct connection
{
  int fd;

  buffer_t *rbuf;
  buffer_t *wbuf;

  int write_failed_count;
  int read_failed_count;

  // state
} connection_t;

//////////////////////////////////////////////////////////////////////
// Global variables
//////////////////////////////////////////////////////////////////////
fd_set global_read_set;
fd_set global_write_set;
fd_set global_exception_set;

int global_max_fd = 0;

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

//////////////////////////////////////////////////////////////////////
// Prototypes
//////////////////////////////////////////////////////////////////////
int network_init(int n);
void network_finalize(void);

void network_register_read(int fd);
void network_register_write(int fd);
void network_register_exception(int fd);

void network_unregister_read(int fd);
void network_unregister_write(int fd);
void network_unregister_exception(int fd);

connection_t *connection_create(int fd);
void connection_destroy(connection_t * c);
void dump_connection_list(void);

void data_exchange(int src);

//////////////////////////////////////////////////////////////////////
// Functions
//////////////////////////////////////////////////////////////////////
int main(int argc, char **argv)
{
  if (argc < 3)
  {
    fprintf(stdout, "Usage: %s <ip> <port>\n", argv[0]);
    exit(1);
  }

  int server_listening_socket;

  // XXX: step 1, socket();
  //int socket(int domain, int type, int protocol);
  if ((server_listening_socket = socket(PF_INET, SOCK_STREAM, 0)) < 0)
  {
    fprintf(stderr, "socket() failed: %s\n", strerror(errno));
    exit(1);
  }

  // XXX: step 1.5, reuse address
  int on;

  on = 1;
  // FIXME: check return value of setsockoipt
  // int setsockopt(int s, int level, int optname, const void *optval, socklen_t optlen);
  setsockopt(server_listening_socket, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));

#ifdef SO_REUSEPORT
  on = 1;
  // FIXME: check return value of setsockoipt
  // int setsockopt(int s, int level, int optname, const void *optval, socklen_t optlen);
  setsockopt(server_listening_socket, SOL_SOCKET, SO_REUSEPORT, &on, sizeof(on));
#endif

  // XXX: step 2, bind();
  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]);

  // int bind(int sockfd, struct sockaddr *my_addr, socklen_t addrlen);
  if (bind(server_listening_socket, (struct sockaddr *) &server_address, sizeof(server_address)) < 0)
  {
    fprintf(stderr, "bind() failed: %s\n", strerror(errno));
    // FIXME:
    exit(1);
  }
  // XXX: step 3, listen();
  // int listen(int s, int backlog);
  if (listen(server_listening_socket, LISTENQ) < 0)
  {
    fprintf(stderr, "listen() failed: %s\n", strerror(errno));
    // FIXME:
    exit(1);
  }

#if 0
  FD_CLR(int fd, fd_set * set);
  FD_ISSET(int fd, fd_set * set);
  FD_SET(int fd, fd_set * set);

  FD_ZERO(fd_set * set);
#endif

  network_init(1024);

  network_register_read(server_listening_socket);

  struct timeval to;
  int n;

  fd_set rset;
  fd_set wset;
  fd_set eset;

  // XXX: main loop
  for (;;)
  {
    rset = global_read_set;
    wset = global_write_set;
    eset = global_exception_set;

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

    //int select(int n, fd_set * readfds, fd_set * writefds, fd_set * exceptfds, struct timeval *timeout);
    if ((n = select(global_max_fd + 1, &rset, &wset, &eset, &to)) < 0)
    {
      // XXX: error
      if (errno == EINTR)
      {
	fprintf(stdout, "select() broken by signal, continue ...\n");
	continue;
      }
      else if (errno == EAGAIN)
      {
	continue;
      }
      else
      {
	fprintf(stderr, "select() failed: %s\n", strerror(errno));
	// FIXME: 
	exit(1);
      }
    }
    else if (n == 0)
    {
      // XXX: timeout
      // TODO: do something
      fprintf(stdout, "timeout ...\n");
    }
    else
    {
      fprintf(stdout, "%d fd is ready for reading or writting or exception.\n", n);

      int i;

      for (i = 0; i <= global_max_fd; i++)
      {
	// XXX: check readable
	if (FD_ISSET(i, &rset))
	{
	  //fprintf(stdout, "fd %d is ready for reading.\n", i);
	  if (i == server_listening_socket)
	  {
	    fprintf(stdout, "fd %d is ready for accepting.\n", i);
	    // accept()
	    // XXX: step 4, accept();
	    int new_accepted_socket;
	    struct sockaddr_in peer_address;
	    socklen_t peer_address_length;

	    peer_address_length = sizeof(peer_address);

	    // int accept(int s, struct sockaddr *addr, socklen_t *addrlen);
	    if ((new_accepted_socket = accept(server_listening_socket, (struct sockaddr *) &peer_address, &peer_address_length)) < 0)
	    {
	      fprintf(stderr, "accept() failed: %s\n", strerror(errno));
	      // FIXME:
	      exit(1);
	    }
	    else
	    {
	      fprintf(stdout, "New connection %d from %s:%d\n", new_accepted_socket, inet_ntoa(peer_address.sin_addr), ntohs(peer_address.sin_port));

	      network_register_read(new_accepted_socket);
	      //network_register_write(new_accepted_socket);
	      //network_register_exception(new_accepted_socket);

	      connection_t *conn;

	      conn = connection_create(new_accepted_socket);
	      global_connection_list[new_accepted_socket] = conn;

#ifdef _DEBUG_
	      dump_connection_list();
#endif
	    }
	  }
	  else
	  {
	    // read()
	    ssize_t read_bytes;
	    char buffer[BUFFER_SIZE];

	    fprintf(stdout, "fd %d is ready for reading.\n", i);
	  again:
	    //ssize_t read(int fd, void *buf, size_t count);
	    if ((read_bytes = read(i, buffer, BUFFER_SIZE)) < 0)
	    {
	      if (errno == EINTR)
		goto again;
	      else
	      {
		fprintf(stderr, "read() error on socket %d\n", i);
		// FIXME:
	      }
	    }
	    else if (read_bytes == 0)
	    {
	      fprintf(stdout, "Connection closed by peer.\n");

	      close(i);

	      network_unregister_read(i);
	      network_unregister_write(i);
	      network_unregister_exception(i);

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

	      //int buffer_append_data(buffer_t * buf, char *data, unsigned long length)
	      buffer_append_data(global_connection_list[i]->rbuf, buffer, read_bytes);

	      data_exchange(i);
	    }
	  }
	}

	// XXX: check writable
	if (FD_ISSET(i, &wset))
	{
	  ssize_t written;

	  // write()
	  fprintf(stdout, "fd %d is ready for writting.\n", i);

	re_writting:
	  //ssize_t write(int fd, const void *buf, size_t count);
	  if ((written = write(i, global_connection_list[i]->wbuf->p, global_connection_list[i]->wbuf->length)) < 0)
	  {
	    if (errno == EINTR)
	    {
	      goto re_writting;
	    }
	    else
	    {
	      fprintf(stderr, "write failed on socket %d: %s\n", i, strerror(errno));
	      // FIXME: 
	    }
	  }
	  else
	  {
	    connection_t *current;

	    current = global_connection_list[i];
#if 0
	    // void *memmove(void *dest, const void *src, size_t n);
	    memmove(current->wbuf->p, current->wbuf->p + written, current->wbuf->length - written);
	    current->wbuf->length -= written;
#endif
	    buffer_remove_data(current->wbuf, written);

	    if (current->wbuf->length == 0)
	    {
	      network_unregister_write(i);
	    }
	  }
	}

	// XXX: check exception
	if (FD_ISSET(i, &eset))
	{
	  // TODO: exception handler
	}
      }
    }
  }

  // XXX: step 5, r/w on new accepted socket
  // XXX: step 6, close new accepted socket
  // XXX: step 7, close listening socket
  close(server_listening_socket);

  network_finalize();

  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, "allocate memory failed: %s\n", strerror(errno));
    return -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_exception_set);

  return 0;
}

void network_finalize(void)
{
  // TODO: realease all items in global_connection_list
  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_exception(int fd)
{
  FD_SET(fd, &global_exception_set);

  if (fd > global_max_fd)
    global_max_fd = fd;
}

void network_unregister_exception(int fd)
{
  FD_CLR(fd, &global_exception_set);
}

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

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

  conn->fd = fd;

  //buffer_t *buffer_create(unsigned long size, unsigned long maxsize)
  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_destroy(c->rbuf);
      c->rbuf = NULL;
    }

    if (c->wbuf)
    {
      buffer_destroy(c->wbuf);
      c->wbuf = NULL;
    }

    free(c);
    c = NULL;
  }
}

void dump_connection_list(void)
{
  int i;
  connection_t *c;

  for (i = 0; i < global_support_connection_num; i++)
  {
    c = global_connection_list[i];

    if (c != NULL)
    {
      fprintf(stdout, "global_connection_list[%d]: object = %p, fd = %d\n", i, c, c->fd);
    }
  }
}

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);
  }
}