broadd fix udp connect + async connect_relay

1 files changed, 171 insertions, 73 deletions

diff --git a/broadd.c b/broadd.c
index c5a8dbf..cef3531 100644
--- a/broadd.c
+++ b/broadd.c
@@ -86,6 +86,7 @@
 #include <inttypes.h>
 #include <netdb.h>
 #include <poll.h>
+#include <pthread.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -121,6 +122,8 @@ int env_rate_limit_s = RATE_LIMIT_S;
 int env_wildcard     = 1;
 onlyudp(int env_client_timeout = CLIENT_TIMEOUT);
 
+pthread_mutex_t relay_mutex = PTHREAD_MUTEX_INITIALIZER;
+
 #define DEDUP_SIZE 65536
 #define DEDUP_MASK (DEDUP_SIZE - 1)
 
@@ -392,11 +395,21 @@ void broadcast_msg(uint64_t sender_id, char* payload, int msg_len) {
     if (sent > 0) log("sentto:%d\n", sent);
 }
 
-void connect_relay(const char* ip, const char* port) {
+void* connect_relay_thread(void* arg) {
+    char* ip_port = arg;
+
     if (relay_count >= MAX_RELAYS) {
         printf("relay(full)\n");
-        return;
+        free(ip_port);
+        return NULL;
+    }
+
+    char ip[ 64 ], port[ 16 ];
+    if (sscanf(ip_port, "%63s %15s", ip, port) != 2) {
+        free(ip_port);
+        return NULL;
     }
+    free(ip_port);
 
     // try tcp
     struct addrinfo hint = {}, *result;
@@ -404,47 +417,91 @@ void connect_relay(const char* ip, const char* port) {
     hint.ai_socktype     = SOCK_STREAM;
     if (getaddrinfo(ip, port, &hint, &result) != 0) {
         perror("relay: getaddrinfo failed\n");
-        return;
+        return NULL;
     }
 
     int  is_udp = 0;
     int  fd     = socket(result->ai_family, result->ai_socktype, result->ai_protocol);
     char gbuf[ 64 ];
 
-    if (fd >= 0 && connect(fd, result->ai_addr, result->ai_addrlen) == 0) {
-        // read greeting
-        int n = read(fd, gbuf, sizeof(gbuf) - 1);
-        if (n > 0 && gbuf[ 0 ] == 'I' && gbuf[ 1 ] == 'T') {
-            // confirmed tcp
-            char pbuf[ 9 ] = "P";
-            memcpy(pbuf + 1, &wildcard_pass, 8);
-            write(fd, pbuf, 9);
-            read(fd, gbuf, sizeof(gbuf)); // p_ok
-        } else {
-            // not a broadd tcp server
-            close(fd);
-            fd = -1;
-        }
+    if (fd < 0) {
+        perror("relay: socket");
+        freeaddrinfo(result);
+        return NULL;
+    }
+
+    // set nonblocking
+    int flags = fcntl(fd, F_GETFL);
+    if (flags == -1 || fcntl(fd, F_SETFL, flags | O_NONBLOCK) == -1) {
+        close(fd);
+        freeaddrinfo(result);
+        printf("relay: fcntl failed\n");
+        return NULL;
+    }
+
+    int conn_err = connect(fd, result->ai_addr, result->ai_addrlen);
+    if (conn_err == 0) goto tcp_connected;
+    if (errno != EINPROGRESS) {
+        close(fd);
+        freeaddrinfo(result);
+        printf("relay: connect fail\n");
+        return NULL;
+    }
+
+    struct pollfd pfd  = { .fd = fd, .events = POLLOUT };
+    int           pres = poll(&pfd, 1, 10000); // 10s
+
+    if (pres <= 0) {
+        // timeout or error
+        close(fd);
+        freeaddrinfo(result);
+        if (pres == 0) printf("relay: tcp timeout\n");
+        else { perror("relay: poll failed"); }
+        return NULL;
+    }
+
+    // check success
+    int       so_err = 0;
+    socklen_t so_len = sizeof(so_err);
+    if (getsockopt(fd, SOL_SOCKET, SO_ERROR, &so_err, &so_len) == -1 || so_err != 0) {
+        printf("relay: trying udp");
+        close(fd);
+        freeaddrinfo(result);
+        goto try_udp;
+    }
+
+tcp_connected:
+    fcntl(fd, F_SETFL, fcntl(fd, F_GETFL, 0) & ~O_NONBLOCK);
+    // read greeting
+    int n = read(fd, gbuf, sizeof(gbuf) - 1);
+    if (n > 0 && gbuf[ 0 ] == 'I' && gbuf[ 1 ] == 'T') {
+        // it's tcp
+        char pbuf[ 9 ] = "P";
+        memcpy(pbuf + 1, &wildcard_pass, 8);
+        write(fd, pbuf, 9);
+        read(fd, gbuf, sizeof(gbuf)); // p_ok
+        // set nonblock
+        fcntl(fd, F_SETFL, fcntl(fd, F_GETFL, 0) | O_NONBLOCK);
     } else {
-        if (fd >= 0) {
-            close(fd);
-            fd = -1;
-        }
+        close(fd);
+        fd = -1;
     }
+
     freeaddrinfo(result);
 
     // udp fallback
     if (fd < 0) {
+    try_udp:
         hint.ai_socktype = SOCK_DGRAM;
         if (getaddrinfo(ip, port, &hint, &result) != 0) {
             printf("relay: getaddrinfo failed\n");
-            return;
+            return NULL;
         }
         fd = socket(result->ai_family, result->ai_socktype, result->ai_protocol);
         if (fd < 0) {
             freeaddrinfo(result);
             printf("relay: socket failed\n");
-            return;
+            return NULL;
         }
 
         // prevent hanging
@@ -459,7 +516,7 @@ void connect_relay(const char* ip, const char* port) {
             close(fd);
             freeaddrinfo(result);
             printf("relay: udp connect failed\n");
-            return;
+            return NULL;
         }
 
         struct sockaddr_storage src;
@@ -470,7 +527,7 @@ void connect_relay(const char* ip, const char* port) {
             printf("relay: udp timeout\n");
             close(fd);
             freeaddrinfo(result);
-            return;
+            return NULL;
         }
         gbuf[ rn ]         = '\0';
         int remote_timeout = (gbuf[ 0 ] == 'T') ? atoi(gbuf + 1) : 30;
@@ -480,35 +537,73 @@ void connect_relay(const char* ip, const char* port) {
             close(fd);
             freeaddrinfo(result);
             printf("relay: invalid handshake\n");
-            return;
+            return NULL;
         }
 
+        fcntl(fd, F_SETFL, fcntl(fd, F_GETFL, 0) | O_NONBLOCK);
+
         char pbuf[ 9 ] = "P";
         memcpy(pbuf + 1, &wildcard_pass, 8);
         sendto(fd, pbuf, 9, 0, result->ai_addr, result->ai_addrlen);
         is_udp = 1;
 
-        struct relay* r   = &relays[ relay_count++ ];
-        r->fd             = fd;
-        r->is_udp         = 1;
-        r->last_hb        = time(NULL);
-        r->remote_timeout = remote_timeout;
-        memcpy(&r->addr, result->ai_addr, result->ai_addrlen);
-        r->addrlen = result->ai_addrlen;
+        pthread_mutex_lock(&relay_mutex);
+        if (relay_count >= MAX_RELAYS) {
+            printf("relay(full) %s:%s\n", ip, port);
+            close(fd);
+        } else {
+            struct relay* r   = &relays[ relay_count++ ];
+            r->fd             = fd;
+            r->is_udp         = 1;
+            r->last_hb        = time(NULL);
+            r->remote_timeout = remote_timeout;
+            memcpy(&r->addr, result->ai_addr, result->ai_addrlen);
+            r->addrlen = result->ai_addrlen;
+            printf("relay-connect %s:%s (%s)\n", ip, port, is_udp ? "udp" : "tcp");
+        }
+        pthread_mutex_unlock(&relay_mutex);
         freeaddrinfo(result);
 
         tv.tv_sec = 0;
         setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, (const char*) &tv, sizeof tv);
     } else {
-        struct relay* r   = &relays[ relay_count++ ];
-        r->fd             = fd;
-        r->is_udp         = 0;
-        r->last_hb        = time(NULL);
-        r->remote_timeout = 0;
+        pthread_mutex_lock(&relay_mutex);
+        if (relay_count >= MAX_RELAYS) {
+            printf("relay(full) %s:%s\n", ip, port);
+            close(fd);
+        } else {
+            struct relay* r   = &relays[ relay_count++ ];
+            r->fd             = fd;
+            r->is_udp         = 0;
+            r->last_hb        = time(NULL);
+            r->remote_timeout = 0;
+        }
+        pthread_mutex_unlock(&relay_mutex);
+    }
+
+    printf("successful connect\n");
+    return NULL;
+}
+
+void connect_relay(const char* ip, const char* port) {
+    if (relay_count >= MAX_RELAYS) {
+        printf("relay(full)\n");
+        return;
     }
 
-    fcntl(fd, F_SETFL, fcntl(fd, F_GETFL, 0) | O_NONBLOCK);
-    printf("relay-connect %s:%s (%s)\n", ip, port, is_udp ? "udp" : "tcp");
+    char* ip_port = malloc(128);
+    if (!ip_port) return;
+
+    snprintf(ip_port, 128, "%s %s", ip, port);
+
+    pthread_t th;
+    if (pthread_create(&th, NULL, connect_relay_thread, ip_port) != 0) {
+        perror("thread spawn failed\n");
+        free(ip_port);
+        return;
+    }
+
+    pthread_detach(th);
 }
 
 char* admin_msg;
@@ -523,44 +618,43 @@ int main(int argc, char** argv) {
     onlyudp(printf("hi from udp;" verb(" verbose mode") "\n"));
 
     // parse args
-    {
-        const char* e;
-        if ((e = getenv("MAX_CLIENTS"))) {
-            int v = atoi(e);
-            if (v > 0 && v <= 65536) env_max_clients = v;
-            else
-                printf("MAX_CLIENTS: ignored (must be 1...65536)\n");
-        }
-        if ((e = getenv("BUFS"))) {
-            int v = atoi(e);
-            if (v > 0 && v <= 65536) env_bufs = v;
-            else
-                printf("BUFS: ignored (must be 1...65536)\n");
-        }
-        if ((e = getenv("RATE_LIMIT_S"))) {
-            int v = atoi(e);
-            if (v > 0) env_rate_limit_s = v;
-            else
-                printf("RATE_LIMIT_S: ignored (must be >0)\n");
-        }
-        onlyudp(if ((e = getenv("CLIENT_TIMEOUT"))) {
-            int v = atoi(e);
-            if (v > 0) env_client_timeout = v;
-            else
-                printf("CLIENT_TIMEOUT: ignored (must be >0)\n");
-        });
-        if ((e = getenv("WILDCARD"))) {
-            int v        = atoi(e);
-            env_wildcard = (v != 0);
-        }
+    const char* e;
+    if ((e = getenv("MAX_CLIENTS"))) {
+        int v = atoi(e);
+        if (v > 0 && v <= 65536) env_max_clients = v;
+        else
+            printf("MAX_CLIENTS: ignored (must be 1...65536)\n");
+    }
+    if ((e = getenv("BUFS"))) {
+        int v = atoi(e);
+        if (v > 0 && v <= 65536) env_bufs = v;
+        else
+            printf("BUFS: ignored (must be 1...65536)\n");
+    }
+    if ((e = getenv("RATE_LIMIT_S"))) {
+        int v = atoi(e);
+        if (v > 0) env_rate_limit_s = v;
+        else
+            printf("RATE_LIMIT_S: ignored (must be >0)\n");
+    }
+    onlyudp(if ((e = getenv("CLIENT_TIMEOUT"))) {
+        int v = atoi(e);
+        if (v > 0) env_client_timeout = v;
+        else
+            printf("CLIENT_TIMEOUT: ignored (must be >0)\n");
+    });
+    if ((e = getenv("WILDCARD"))) {
+        int v        = atoi(e);
+        env_wildcard = (v != 0);
     }
 
-    admin_msg = malloc(env_bufs);
+    // alloc admin msg
+    admin_msg      = malloc(env_bufs);
+    admin_msg[ 0 ] = 'A';
 
     set_stdin_nonblocking();
 
-    admin_msg[ 0 ] = 'A';
-
+    // fill passwords
     memcpy(&default_pass, "00000000", 8);
     memcpy(&wildcard_pass, "00000001", 8);
 
@@ -623,6 +717,7 @@ int main(int argc, char** argv) {
             n++;
         });
 
+        pthread_mutex_lock(&relay_mutex);
         for (int i = 0; i < relay_count; i++) {
             fds[ n ].fd     = relays[ i ].fd;
             fds[ n ].events = POLLIN;
@@ -639,6 +734,7 @@ int main(int argc, char** argv) {
                 if (poll_timeout < 0 || t < poll_timeout) poll_timeout = t;
             }
         }
+        pthread_mutex_unlock(&relay_mutex);
         poll(fds, n, poll_timeout);
 
         time_t now = time(NULL);
@@ -799,6 +895,7 @@ int main(int argc, char** argv) {
         }
 
         // process relay fds
+        pthread_mutex_lock(&relay_mutex);
         for (int i = relay_count - 1; i >= 0; i--) {
             struct relay*  r  = &relays[ i ];
             struct pollfd* pf = &fds[ relay_base + i ];
@@ -827,6 +924,7 @@ int main(int argc, char** argv) {
             // broadcast relay msg with id1
             if (rbuf[ 0 ] == 'M' && rrd >= 9) { broadcast_msg(1, rbuf + 1, rrd - 1); }
         }
+        pthread_mutex_unlock(&relay_mutex);
 
         onlyudp(if (unlikely(got_msg)) {
             for (int i = 0; i < active; i++)