Chat server using the poll() API

server.c

#include <arpa/inet.h>
#include <poll.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

// This is the static configuration. Since this is just a small demo program,
// there is no dynamic config loading. If you want to change any settings, just
// recompile the program.

#define PORT 1234
#define MSG_LIMIT 256
#define MAX_CONNECTIONS 4096

#define VERIFY(x)                                                              \
  do {                                                                         \
    if (!(x)) {                                                                \
      fprintf(stderr, "%s:%d Assertion failed on function '%s'\n", __FILE__,   \
              __LINE__, __func__);                                             \
      exit(-1);                                                                \
    }                                                                          \
  } while (0)

static int create_listening_socket(uint16_t port, int backlog) {
  int sock = socket(AF_INET, SOCK_STREAM, 0);
  VERIFY(sock > 0);

  int enable = 1;
  VERIFY(!setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(enable)));

  struct sockaddr_in addr;
  addr.sin_family = AF_INET;
  addr.sin_addr.s_addr = INADDR_ANY;
  addr.sin_port = htons(port);

  VERIFY(bind(sock, (const struct sockaddr *)&addr, sizeof(addr)) == 0);

  listen(sock, backlog);

  return sock;
}

typedef struct {
  int fd;
  char readBuf[MSG_LIMIT];
  char nick[MSG_LIMIT];
  size_t readHead;
} UserConnection;

// Send text to every client that is connected and has entered a nick
static void emit_to_all(UserConnection *connections, const char *msg, size_t len) {
  for (size_t i = 0; i < MAX_CONNECTIONS; i++)
    if (connections[i].fd > 0 && connections[i].nick[0] != '\0')
      write(connections[i].fd, msg, len);
}

int main() {
  int listener = create_listening_socket(PORT, 5);

  UserConnection *connections = calloc(MAX_CONNECTIONS, sizeof(UserConnection));
  VERIFY(connections != NULL);

  for (size_t i = 0; i < MAX_CONNECTIONS; i++)
    connections[i].fd = -1;

  printf("Listening socket FD = %d\n", listener);

  while (true) {
    struct pollfd *polls = calloc(MAX_CONNECTIONS + 1, sizeof(struct pollfd));
    VERIFY(polls != NULL);

    polls[0].fd = listener;
    polls[0].events = POLLIN;

    for (size_t i = 0; i < MAX_CONNECTIONS; i++) {
        polls[i + 1].fd = connections[i].fd;
        polls[i + 1].events = POLLIN;
    }

    VERIFY(poll(polls, MAX_CONNECTIONS + 1, -1) > 0);

    for (size_t pollId = 1; pollId < MAX_CONNECTIONS + 1; pollId++) {
      struct pollfd p = polls[pollId];
      UserConnection *uc = &connections[pollId - 1];

      if (p.revents & POLLHUP || p.revents & POLLERR) {
        connections[pollId - 1].fd = 0;
        close(p.fd);
        uc->fd = 0;
        printf("%d had an oopsie\n", p.fd);
      }

      if (p.revents & POLLIN) {
        char buf[1];
        if (read(p.fd, &buf, 1) < 1) {
          close(p.fd);
          uc->fd = 0;
          if (uc->nick[0] != '\0') {
            emit_to_all(connections, "* ", 2);
            emit_to_all(connections, uc->nick, strlen(uc->nick));
            emit_to_all(connections, " has disconnected\n", strlen(" has disconnected\n"));
          }
          uc->nick[0] = '\0';
          continue;
        }

        if (buf[0] == '\n') {
          uc->readBuf[uc->readHead] = '\0';
          if (uc->nick[0] == '\0') {
            strcpy(uc->nick, uc->readBuf);
            uc->readHead = 0;
            emit_to_all(connections, "* ", 2);
            emit_to_all(connections, uc->nick, strlen(uc->nick));
            emit_to_all(connections, " has joined\n", strlen(" has joined\n"));
            continue;
          }

          emit_to_all(connections, "<", 1);
          emit_to_all(connections, uc->nick, strlen(uc->nick));
          emit_to_all(connections, "> ", 2);
          emit_to_all(connections, uc->readBuf, uc->readHead);
          emit_to_all(connections, "\n", 1);
          uc->readHead = 0;
          continue;
        }
        uc->readBuf[uc->readHead++] = buf[0];

        if (uc->readHead >= MSG_LIMIT) {
          write(p.fd, "* Message too long, disconnecting...\n", strlen("* Message too long, disconnecting...\n"));
          close(p.fd);
          uc->fd = 0;
          uc->nick[0] = '\0';
        }
      }
    }

    if (polls[0].revents & POLLIN) {
      printf("Accepted a connection\n");
      int fd = accept(listener, NULL, NULL);
      VERIFY(fd != -1);

      bool found = false;
      for (size_t i = 0; i < MAX_CONNECTIONS; i++) {
        UserConnection *conn = &connections[i];
        if (conn->fd == -1) {
          conn->fd = fd;
          conn->readHead = 0;
          found = true;
          
          const char *c = "Please enter your nick: ";
          write(fd, c, strlen(c));
          break;
        }
      }

      // If we do not have any spare connection slots, send an apology message
      // explaining the situation.
      
      if (!found) {
        const char *str = "We are unable to accept your connection as we have "
                          "too many users right now. Please try again later.\n";
        write(fd, str, strlen(str));
        close(fd);
      }

      printf("Accepted connection with FD = %d\n", fd);
    }

    free(polls);
  }

  return 0;
}