gwj / socket · Files

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

#define DEFAULT_PORT 9734
#define MAX_CONN 16
#define MAX_EVENTS 32
#define BUF_SIZE 16
#define MAX_LINE 256

void server_run();
void client_run();

int main(int argc, char *argv[]) {
  int opt;
  char role = 's';
  while ((opt = getopt(argc, argv, "cs")) != -1) {
    switch (opt) {
    case 'c':
      role = 'c';
      break;
    case 's':
      break;
    default:
      printf("usage: %s [-cs]\n", argv[0]);
      exit(1);
    }
  }
  if (role == 's') {
    server_run();
  } else {
    client_run();
  }
  return 0;
}

/*
 * register events of fd to epfd
 */
static void epoll_ctl_add(int epfd, int fd, uint32_t events) {
  struct epoll_event ev;
  ev.events = events;
  ev.data.fd = fd;
  if (epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &ev) == -1) {
    perror("epoll_ctl()\n");
    exit(1);
  }
}

static void set_sockaddr(struct sockaddr_in *addr) {
  bzero((char *)addr, sizeof(struct sockaddr_in));
  addr->sin_family = AF_INET;
  addr->sin_addr.s_addr = INADDR_ANY;
  addr->sin_port = htons(DEFAULT_PORT);
}

/*
 * epoll echo server
 */
void server_run() {
  int i;
  int n;
  int epfd;
  int nfds;
  int listen_sock;
  int conn_sock;
  socklen_t socklen;
  char buf[BUF_SIZE];
  struct sockaddr_in srv_addr;
  struct sockaddr_in cli_addr;
  struct epoll_event events[MAX_EVENTS];

  listen_sock = socket(AF_INET, SOCK_STREAM, 0);

  if (listen_sock == -1) {
    perror("socket");
    exit(1);
  }

  set_sockaddr(&srv_addr);

  int reuseaddr = 1;
  if (setsockopt(listen_sock, SOL_SOCKET, SO_REUSEADDR, &reuseaddr,
                 sizeof(int)) < 0) {
    perror("setsockopt(SO_REUSEADDR)");
    exit(1);
  }

  if (bind(listen_sock, (struct sockaddr *)&srv_addr, sizeof(srv_addr)) < 0) {
    perror("bind");
    exit(1);
  }

  if (listen(listen_sock, MAX_CONN) < 0) {
    perror("listen");
    exit(1);
  }

  epfd = epoll_create(1);
  if (epfd == -1) {
    perror("epoll_create");
    exit(1);
  }

  epoll_ctl_add(epfd, listen_sock, EPOLLIN | EPOLLOUT);

  socklen = sizeof(cli_addr);
  for (;;) {
    do {
      nfds = epoll_wait(epfd, events, MAX_EVENTS, -1);
    } while (nfds == -1 && errno == EINTR);

    if (nfds == -1) {
      perror("epoll_wait");
      exit(1);
    }

    for (i = 0; i < nfds; i++) {
      if (events[i].data.fd == listen_sock) {
        /* handle new connection */

        do {
          conn_sock =
              accept(listen_sock, (struct sockaddr *)&cli_addr, &socklen);
        } while (conn_sock == -1 && errno == EINTR);

        if (conn_sock == -1) {
          perror("accept");
          exit(1);
        }

        inet_ntop(AF_INET, (char *)&(cli_addr.sin_addr), buf, sizeof(buf));
        printf("[+] connected with %s:%d\n", buf, ntohs(cli_addr.sin_port));

        epoll_ctl_add(epfd, conn_sock,
                      EPOLLIN | EPOLLET | EPOLLRDHUP | EPOLLHUP);
      } else if (events[i].events & EPOLLIN) {
        /* handle EPOLLIN event */
        for (;;) {
          n = recv(events[i].data.fd, buf, sizeof(buf) - 1, MSG_DONTWAIT);
          if (n == -1 && ((errno == EAGAIN) || (errno == EWOULDBLOCK))) {
            break;
          } else if (n == 0) {
            break;
          } else if (n == -1) {
            perror("recv");
            break;
          } else {
            buf[n] = '\0';
            printf("[+] data: %s\n", buf);
            n = 0;
            while (n != strlen(buf)) {
              int rv = send(events[i].data.fd, buf + n, strlen(buf) - n,
                            MSG_NOSIGNAL);
              if (rv == -1 && (errno == EINTR))
                continue; // FIXME: might lock up due to client not receiving -
                          // should wait using epoll
              if (rv == -1)
                break;
              n += rv;
            }
          }
        }
      } else {
        printf("[+] unexpected\n");
      }
      /* check if the connection is closing */
      if (events[i].events & (EPOLLRDHUP | EPOLLHUP)) {
        printf("[+] connection closed\n");

        if (epoll_ctl(epfd, EPOLL_CTL_DEL, events[i].data.fd, NULL) == -1) {
          perror("epoll_ctl");
          exit(1);
        }
        close(events[i].data.fd);
        continue;
      }
    }
  }
}

/*
 * test clinet
 */
void client_run() {
  int n;
  int c;
  int sockfd;
  char buf[MAX_LINE];
  struct sockaddr_in srv_addr;

  sockfd = socket(AF_INET, SOCK_STREAM, 0);
  if (sockfd == -1) {
    perror("socket");
    exit(1);
  }

  set_sockaddr(&srv_addr);

  if (connect(sockfd, (struct sockaddr *)&srv_addr, sizeof(srv_addr)) < 0) {
    perror("connect()");
    exit(1);
  }

  for (;;) {
    printf("input: ");
    fgets(buf, sizeof(buf), stdin);
    c = strlen(buf);

    n = 0;
    while (n != c) {
      int rv = send(sockfd, buf + n, c - n, MSG_NOSIGNAL);
      if (rv == -1 && (errno == EINTR))
        continue;
      if (rv == -1)
        exit(1);
      n += rv;
    }

    n = 0;
    while (n != c) {
      int rv = read(sockfd, buf + n, c - n);
      if (rv == -1 && (errno == EINTR))
        continue;
      if (rv == -1)
        exit(1);
      if (rv == 0)
        break;
      n += rv;
    }

    buf[c] = '\0';
    printf("echo: %s\n", buf);
  }
  close(sockfd);
}