From 420b989adb46b60d7cf38f686b78138a63ecd8b8 Mon Sep 17 00:00:00 2001
From: Valtteri Koskivuori <vkoskiv@gmail.com>
Date: Wed, 25 May 2022 00:43:22 +0300
Subject: [PATCH] Add optional poll() implementation under MG_ENABLE_POLL

This adds the ability to substitute select() with poll() on supported
unix environments when MG_ENABLE_POLL is defined. A make flag called
USE_POLL is provided as well.
Using poll() removes the limitation of FD_SETSIZE concurrent sockets,
generally 1024 on Linux environments.
---
 mongoose.c      | 35 ++++++++++++++++++++++++++++++++++-
 mongoose.h      |  8 ++++++++
 src/arch_unix.h |  4 ++++
 src/config.h    |  4 ++++
 src/sock.c      | 35 ++++++++++++++++++++++++++++++++++-
 5 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/mongoose.c b/mongoose.c
index ff996417..2e69cc1e 100644
--- a/mongoose.c
+++ b/mongoose.c
@@ -4301,7 +4301,7 @@ static void accept_conn(struct mg_mgr *mgr, struct mg_connection *lsn) {
 #endif
       MG_ERROR(("%lu accept failed, errno %d", lsn->id, MG_SOCK_ERRNO));
 #if (MG_ARCH != MG_ARCH_WIN32) && (MG_ARCH != MG_ARCH_FREERTOS_TCP) && \
-    (MG_ARCH != MG_ARCH_TIRTOS)
+    (MG_ARCH != MG_ARCH_TIRTOS) && !(MG_ENABLE_POLL)
   } else if ((long) fd >= FD_SETSIZE) {
     MG_ERROR(("%ld > %ld", (long) fd, (long) FD_SETSIZE));
     closesocket(fd);
@@ -4393,6 +4393,39 @@ static void mg_iotest(struct mg_mgr *mgr, int ms) {
     FreeRTOS_FD_CLR(c->fd, mgr->ss,
                     eSELECT_READ | eSELECT_EXCEPT | eSELECT_WRITE);
   }
+#elif MG_ENABLE_POLL
+  size_t i = 0, n = 0;
+  for (struct mg_connection *c = mgr->conns; c != NULL; c = c->next) n++;
+  struct pollfd fds[n == 0 ? 1 : n]; // Avoid zero-length VLA
+
+  memset(fds, 0, sizeof(fds));
+  for (struct mg_connection *c = mgr->conns; c != NULL; c = c->next, i++) {
+    if (c->is_closing || c->is_resolving || FD(c) == INVALID_SOCKET) {
+      // Socket not valid, ignore
+    } else {
+      fds[i].fd = FD(c);
+      fds[i].events |= POLLIN;
+      if (c->is_connecting || (c->send.len > 0 && c->is_tls_hs == 0)) {
+          fds[i].events |= POLLOUT;
+      }
+    }
+  }
+
+  if (poll(fds, n, ms) < 0) {
+    MG_ERROR(("poll failed, errno: %d", MG_SOCK_ERRNO));
+    return;
+  }
+
+  i = 0;
+  for (struct mg_connection *c = mgr->conns; c != NULL; c = c->next, i++) {
+    if (c->is_closing || c->is_resolving || FD(c) == INVALID_SOCKET) {
+      // Socket not valid, ignore
+    } else {
+      c->is_readable = (unsigned)(fds[i].revents & POLLIN ? true : false);
+      c->is_writable = (unsigned)(fds[i].revents & POLLOUT ? true : false);
+      fds[i].revents = 0;
+    }
+  }
 #else
   struct timeval tv = {ms / 1000, (ms % 1000) * 1000}, tv_zero = {0, 0};
   struct mg_connection *c;
diff --git a/mongoose.h b/mongoose.h
index dbb992e5..b260d721 100644
--- a/mongoose.h
+++ b/mongoose.h
@@ -436,7 +436,11 @@ extern int SockSet(SOCKET hSock, int Type, int Prop, void *pbuf, int size);
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#if defined(MG_ENABLE_POLL) && MG_ENABLE_POLL
+#include <poll.h>
+#else
 #include <sys/select.h>
+#endif
 #include <sys/socket.h>
 #include <sys/stat.h>
 #include <sys/time.h>
@@ -578,6 +582,10 @@ int sscanf(const char *, const char *, ...);
 #define MG_ENABLE_MIP 0
 #endif
 
+#ifndef MG_ENABLE_POLL
+#define MG_ENABLE_POLL 0
+#endif
+
 #ifndef MG_ENABLE_FATFS
 #define MG_ENABLE_FATFS 0
 #endif
diff --git a/src/arch_unix.h b/src/arch_unix.h
index dd696e22..424257dc 100644
--- a/src/arch_unix.h
+++ b/src/arch_unix.h
@@ -22,7 +22,11 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#if defined(MG_ENABLE_POLL) && MG_ENABLE_POLL
+#include <poll.h>
+#else
 #include <sys/select.h>
+#endif
 #include <sys/socket.h>
 #include <sys/stat.h>
 #include <sys/time.h>
diff --git a/src/config.h b/src/config.h
index 0d68c468..6dda7931 100644
--- a/src/config.h
+++ b/src/config.h
@@ -4,6 +4,10 @@
 #define MG_ENABLE_MIP 0
 #endif
 
+#ifndef MG_ENABLE_POLL
+#define MG_ENABLE_POLL 0
+#endif
+
 #ifndef MG_ENABLE_FATFS
 #define MG_ENABLE_FATFS 0
 #endif
diff --git a/src/sock.c b/src/sock.c
index 700067be..7c8e6d10 100644
--- a/src/sock.c
+++ b/src/sock.c
@@ -383,7 +383,7 @@ static void accept_conn(struct mg_mgr *mgr, struct mg_connection *lsn) {
 #endif
       MG_ERROR(("%lu accept failed, errno %d", lsn->id, MG_SOCK_ERRNO));
 #if (MG_ARCH != MG_ARCH_WIN32) && (MG_ARCH != MG_ARCH_FREERTOS_TCP) && \
-    (MG_ARCH != MG_ARCH_TIRTOS)
+    (MG_ARCH != MG_ARCH_TIRTOS) && !(MG_ENABLE_POLL)
   } else if ((long) fd >= FD_SETSIZE) {
     MG_ERROR(("%ld > %ld", (long) fd, (long) FD_SETSIZE));
     closesocket(fd);
@@ -475,6 +475,39 @@ static void mg_iotest(struct mg_mgr *mgr, int ms) {
     FreeRTOS_FD_CLR(c->fd, mgr->ss,
                     eSELECT_READ | eSELECT_EXCEPT | eSELECT_WRITE);
   }
+#elif MG_ENABLE_POLL
+  size_t i = 0, n = 0;
+  for (struct mg_connection *c = mgr->conns; c != NULL; c = c->next) n++;
+  struct pollfd fds[n == 0 ? 1 : n]; // Avoid zero-length VLA
+
+  memset(fds, 0, sizeof(fds));
+  for (struct mg_connection *c = mgr->conns; c != NULL; c = c->next, i++) {
+    if (c->is_closing || c->is_resolving || FD(c) == INVALID_SOCKET) {
+      // Socket not valid, ignore
+    } else {
+      fds[i].fd = FD(c);
+      fds[i].events |= POLLIN;
+      if (c->is_connecting || (c->send.len > 0 && c->is_tls_hs == 0)) {
+          fds[i].events |= POLLOUT;
+      }
+    }
+  }
+
+  if (poll(fds, n, ms) < 0) {
+    MG_ERROR(("poll failed, errno: %d", MG_SOCK_ERRNO));
+    return;
+  }
+
+  i = 0;
+  for (struct mg_connection *c = mgr->conns; c != NULL; c = c->next, i++) {
+    if (c->is_closing || c->is_resolving || FD(c) == INVALID_SOCKET) {
+      // Socket not valid, ignore
+    } else {
+      c->is_readable = (unsigned)(fds[i].revents & POLLIN ? true : false);
+      c->is_writable = (unsigned)(fds[i].revents & POLLOUT ? true : false);
+      fds[i].revents = 0;
+    }
+  }
 #else
   struct timeval tv = {ms / 1000, (ms % 1000) * 1000}, tv_zero = {0, 0};
   struct mg_connection *c;