From 2c1c6b50b72912c3b1981d570fad6e3d4fc4f2c6 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Mon, 22 Jan 2024 08:03:03 +0100
Subject: [PATCH] Add pre_open event

---
 src/context.c           |  9 +++++++++
 src/internal/internal.h |  1 +
 src/libusockets.h       |  2 ++
 src/loop.c              | 18 ++++++++++++------
 4 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/context.c b/src/context.c
index 7bf9396..4b57d89 100644
--- a/src/context.c
+++ b/src/context.c
@@ -225,6 +225,9 @@ struct us_socket_context_t *us_create_socket_context(int ssl, struct us_loop_t *
     context->long_timestamp = 0;
     context->global_tick = 0;
 
+    /* Some new events must be set to null for backwards compatibility */
+    context->on_pre_open = 0;
+
     us_internal_loop_link(loop, context);
 
     /* If we are called from within SSL code, SSL code will make further changes to us */
@@ -414,6 +417,12 @@ struct us_socket_t *us_socket_context_adopt_socket(int ssl, struct us_socket_con
     return new_s;
 }
 
+/* For backwards compatibility, this function will be set to nullptr by default. */
+void us_socket_context_on_pre_open(int ssl, struct us_socket_context_t *context, LIBUS_SOCKET_DESCRIPTOR (*on_pre_open)(LIBUS_SOCKET_DESCRIPTOR fd)) {
+    /* For this event, there is no difference between SSL and non-SSL */
+    context->on_pre_open = on_pre_open;
+}
+
 void us_socket_context_on_open(int ssl, struct us_socket_context_t *context, struct us_socket_t *(*on_open)(struct us_socket_t *s, int is_client, char *ip, int ip_length)) {
 #ifndef LIBUS_NO_SSL
     if (ssl) {
diff --git a/src/internal/internal.h b/src/internal/internal.h
index be002d2..2bb4213 100644
--- a/src/internal/internal.h
+++ b/src/internal/internal.h
@@ -130,6 +130,7 @@ struct us_socket_context_t {
     struct us_socket_t *iterator;
     struct us_socket_context_t *prev, *next;
 
+    LIBUS_SOCKET_DESCRIPTOR (*on_pre_open)(LIBUS_SOCKET_DESCRIPTOR fd);
     struct us_socket_t *(*on_open)(struct us_socket_t *, int is_client, char *ip, int ip_length);
     struct us_socket_t *(*on_data)(struct us_socket_t *, char *data, int length);
     struct us_socket_t *(*on_writable)(struct us_socket_t *);
diff --git a/src/libusockets.h b/src/libusockets.h
index d16624a..344c68c 100644
--- a/src/libusockets.h
+++ b/src/libusockets.h
@@ -155,6 +155,8 @@ struct us_socket_context_t *us_create_socket_context(int ssl, struct us_loop_t *
 void us_socket_context_free(int ssl, struct us_socket_context_t *context);
 
 /* Setters of various async callbacks */
+void us_socket_context_on_pre_open(int ssl, struct us_socket_context_t *context,
+    LIBUS_SOCKET_DESCRIPTOR (*on_pre_open)(LIBUS_SOCKET_DESCRIPTOR fd));
 void us_socket_context_on_open(int ssl, struct us_socket_context_t *context,
     struct us_socket_t *(*on_open)(struct us_socket_t *s, int is_client, char *ip, int ip_length));
 void us_socket_context_on_close(int ssl, struct us_socket_context_t *context,
diff --git a/src/loop.c b/src/loop.c
index ca9d2d4..ef10e22 100644
--- a/src/loop.c
+++ b/src/loop.c
@@ -275,13 +275,19 @@ void us_internal_dispatch_ready_poll(struct us_poll_t *p, int error, int events)
                     /* Todo: stop timer if any */
 
                     do {
-                        /* adopt the newly accepted socket */
-                        us_adopt_accepted_socket(0, us_socket_context(0, &listen_socket->s),
-                            client_fd, listen_socket->socket_ext_size, bsd_addr_get_ip(&addr), bsd_addr_get_ip_length(&addr));
+                        struct us_socket_context_t *context = us_socket_context(0, &listen_socket->s);
+                        /* See if we want to export the FD or keep it here (this event can be unset) */
+                        if (context->on_pre_open == 0 || context->on_pre_open(client_fd) == client_fd) {
+
+                            /* Adopt the newly accepted socket */
+                            us_adopt_accepted_socket(0, context,
+                                client_fd, listen_socket->socket_ext_size, bsd_addr_get_ip(&addr), bsd_addr_get_ip_length(&addr));
+
+                            /* Exit accept loop if listen socket was closed in on_open handler */
+                            if (us_socket_is_closed(0, &listen_socket->s)) {
+                                break;
+                            }
 
-                        /* Exit accept loop if listen socket was closed in on_open handler */
-                        if (us_socket_is_closed(0, &listen_socket->s)) {
-                            break;
                         }
 
                     } while ((client_fd = bsd_accept_socket(us_poll_fd(p), &addr)) != LIBUS_SOCKET_ERROR);