Skip to content

Instantly share code, notes, and snippets.

@gkbrk
Created June 25, 2021 22:31
Show Gist options
  • Save gkbrk/c15672a002167eddd1eb2a038c53f57e to your computer and use it in GitHub Desktop.
Save gkbrk/c15672a002167eddd1eb2a038c53f57e to your computer and use it in GitHub Desktop.
Chat server using the poll() API
#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;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment