server: Fix handling of multiple clients

Improve handling of multiple connected clients. This fixes issues related
to blocking of one client while another is connected as well as certain
race conditions.

Signed-off-by: Ian Moffett <ian@osmora.org>

1 files changed, 79 insertions, 8 deletions

diff --git a/lib/libostp/server.c b/lib/libostp/server.c
index aad91e6..3bede26 100644
--- a/lib/libostp/server.c
+++ b/lib/libostp/server.c
@@ -37,10 +37,19 @@
 #include <stdlib.h>
 #include <unistd.h>
 #include <stdio.h>
+#include <pthread.h>
 
 #define MAX_BACKLOG 4
 #define LISTEN_PORT 5352
 
+struct recv_args {
+    void *buf;
+    struct ostp_client *c;
+    struct ostp_listener *lp;
+    size_t n;
+    pthread_t td;
+};
+
 static int
 handle_client(struct sockaddr_in *caddr, struct ostp_client *c, struct ostp_listener *lp)
 {
@@ -77,6 +86,24 @@ handle_client(struct sockaddr_in *caddr, struct ostp_client *c, struct ostp_list
 }
 
 /*
+ * When we get data from a client, a thread is sent
+ * here to handle it.
+ */
+static void *
+recv_td(void *args)
+{
+    struct recv_args *ap = args;
+    struct ostp_listener *lp;
+
+    lp = ap->lp;
+    if (lp->on_recv != NULL) {
+        lp->on_recv(ap->c, ap->buf, ap->n);
+    }
+    free(args);
+    return NULL;
+}
+
+/*
  * Put a listener in a known state during
  * initialization.
  */
@@ -122,14 +149,17 @@ int
 listener_poll(struct ostp_listener *lp)
 {
     struct sockaddr_in caddr;
+    struct recv_args *recv_ap;
     struct ostp_client *c;
     socklen_t caddr_len;
     pthread_t client_td;
-    int client_sock, error = 0;
-    char *ip;
+    int client_sock, n, error = 0;
+    unsigned char *session_key;
+    char *ip, buf[4096];
 
     memset(lp->clients, -1, sizeof(lp->clients));
     lp->clients[0].sockfd = lp->serv_sockfd;
+    lp->client_count = 1;
 
     while (1) {
         FD_ZERO(&lp->client_fds);
@@ -155,23 +185,64 @@ listener_poll(struct ostp_listener *lp)
                 continue;
             }
 
+            if (lp->client_count >= MAX_CLIENTS) {
+                printf("New connection rejected, max clients reached\n");
+                close(client_sock);
+                continue;
+            }
+
             for (int i = 0; i < MAX_CLIENTS; ++i) {
                 c = &lp->clients[i];
-                if (lp->client_count >= MAX_CLIENTS) {
-                    printf("New connection rejected, max clients reached\n");
-                    continue;
-                }
                 if (c->sockfd < 0) {
                     c->sockfd = client_sock;
-                    ip = inet_ntoa(caddr.sin_addr);
+                    c->authed = 0;
+                    ++lp->client_count;
 
+                    ip = inet_ntoa(caddr.sin_addr);
                     printf("Incoming connection from %s\n", ip);
-                    ++lp->client_count;
                     handle_client(&caddr, c, lp);
                     break;
                 }
             }
         }
+
+        /* Handle data from clients */
+        for (int i = 1; i < MAX_CLIENTS; ++i) {
+            c = &lp->clients[i];
+            if (c->sockfd < 0 || !c->authed)
+                continue;
+            if (!FD_ISSET(c->sockfd, &lp->client_fds))
+                continue;
+
+            session_key = c->session.session_key;
+            n = recv_frame(c->sockfd, sizeof(buf) - 1, session_key, buf);
+            if (n < 0) {
+                printf("recv_frame() failure, packet lost\n");
+                continue;
+            }
+            if (n == 0) {
+                printf("Peer disconnected\n");
+                listener_close(lp, c);
+                continue;
+            }
+
+            if (lp->on_recv != NULL) {
+                recv_ap = malloc(sizeof(*recv_ap));
+                if (recv_ap == NULL)
+                    continue;
+
+                /* Prepare on_recv() args */
+                recv_ap->c = c;
+                recv_ap->buf = buf;
+                recv_ap->n = n;
+
+                error = pthread_create(&recv_ap->td, NULL, recv_td, recv_ap);
+                if (error != 0) {
+                    free(recv_ap);
+                    continue;
+                }
+            }
+        }
     }
 
     close(client_sock);