1 files changed, 57 insertions, 19 deletions

diff --git a/src/tcp_server.c b/src/tcp_server.c
index 2e65bc5..a6263f7 100644
--- a/src/tcp_server.c
+++ b/src/tcp_server.c
@@ -6,6 +6,8 @@
  */
 
 #include "tcp_server.h"
+#include "compiler.h"
+#include "event_loop.h"
 #include "log.h"
 #include "net.h"
 #include "signal.h"
@@ -16,18 +18,24 @@
 
 struct tcp_server {
 	int fd;
+	tcp_server_conn_handler conn_handler;
+	void *arg;
 };
 
-int tcp_server_create(struct tcp_server **_server, const char *port)
+int tcp_server_create(struct tcp_server **_server, const char *port,
+                      tcp_server_conn_handler conn_handler, void *arg)
 {
 	int ret = 0;
 
-	struct tcp_server *server = malloc(sizeof(struct tcp_server));
+	struct tcp_server *server = calloc(1, sizeof(struct tcp_server));
 	if (!server) {
 		log_errno("malloc");
 		return -1;
 	}
 
+	server->conn_handler = conn_handler;
+	server->arg = arg;
+
 	ret = net_bind(port);
 	if (ret < 0)
 		goto free;
@@ -50,7 +58,7 @@ void tcp_server_destroy(struct tcp_server *server)
 
 struct child_context {
 	int fd;
-	tcp_server_conn_handler handler;
+	tcp_server_conn_handler conn_handler;
 	void *arg;
 };
 
@@ -65,7 +73,7 @@ static void *connection_thread(void *_ctx)
 	if (ret < 0)
 		goto free_ctx;
 
-	ctx->handler(ctx->fd, ctx->arg);
+	ctx->conn_handler(ctx->fd, ctx->arg);
 
 free_ctx:
 	free(ctx);
@@ -73,7 +81,7 @@ free_ctx:
 	return NULL;
 }
 
-int tcp_server_accept(const struct tcp_server *server, tcp_server_conn_handler handler, void *arg)
+static int create_connection_thread(int fd, tcp_server_conn_handler conn_handler, void *arg)
 {
 	sigset_t old_mask;
 	pthread_t child;
@@ -81,24 +89,20 @@ int tcp_server_accept(const struct tcp_server *server, tcp_server_conn_handler h
 
 	struct child_context *ctx = calloc(1, sizeof(*ctx));
 	if (!ctx) {
-		log_errno("malloc");
+		log_errno("calloc");
 		return -1;
 	}
 
-	ctx->handler = handler;
+	ctx->fd = fd;
+	ctx->conn_handler = conn_handler;
 	ctx->arg = arg;
 
-	ret = net_accept(server->fd);
-	if (ret < 0)
-		goto free_ctx;
-	ctx->fd = ret;
-
 	/* Block all signals (we'll unblock them later); the child thread will
 	 * have all signals blocked initially. This allows the main thread to
 	 * handle SIGINT/SIGTERM/etc. */
 	ret = signal_block_all(&old_mask);
 	if (ret < 0)
-		goto close_conn;
+		goto free_ctx;
 
 	ret = pthread_create(&child, NULL, connection_thread, ctx);
 	if (ret) {
@@ -106,19 +110,53 @@ int tcp_server_accept(const struct tcp_server *server, tcp_server_conn_handler h
 		goto restore_mask;
 	}
 
+restore_mask:
 	/* Restore the previously-enabled signals for handling in the main thread. */
 	signal_restore(&old_mask);
 
 	return ret;
 
-restore_mask:
-	signal_restore(&old_mask);
-
-close_conn:
-	net_close(ctx->fd);
-
 free_ctx:
 	free(ctx);
 
 	return ret;
 }
+
+int tcp_server_accept(const struct tcp_server *server)
+{
+	int fd = -1, ret = 0;
+
+	ret = net_accept(server->fd);
+	if (ret < 0)
+		return ret;
+	fd = ret;
+
+	ret = create_connection_thread(fd, server->conn_handler, server->arg);
+	if (ret < 0)
+		goto close_conn;
+
+	return ret;
+
+close_conn:
+	net_close(fd);
+
+	return ret;
+}
+
+static int tcp_server_event_loop_handler(UNUSED struct event_loop *loop, UNUSED int fd,
+                                         UNUSED short revents, void *_server)
+{
+	struct tcp_server *server = (struct tcp_server *)_server;
+	return tcp_server_accept(server);
+}
+
+int tcp_server_add_to_event_loop(struct tcp_server *server, struct event_loop *loop)
+{
+	struct event_fd entry = {
+	    .fd = server->fd,
+	    .events = POLLIN,
+	    .handler = tcp_server_event_loop_handler,
+	    .arg = server,
+	};
+	return event_loop_add(loop, &entry);
+}