initial commit

1 files changed, 532 insertions, 0 deletions

diff --git a/broadd.c b/broadd.c
new file mode 100644
index 0000000..b72c362
--- /dev/null
+++ b/broadd.c
@@ -0,0 +1,532 @@
+/*
+-- broadd - tcp/udp broadcast server
+------ env vars:
+---------- MAX_CLIENTS    -> maximum number of hypothetical connections
+---------- BUFS           -> buffer size, values >MTU are not recommended on udp
+---------- RATE_LIMIT_S   -> rate limit expiry window
+---------- CLIENT_TIMEOUT -> how long to wait before evicting a udp client
+------ difference:
+---------- udp mode does not send m_ok or p_ok
+---------- send `N\n` to initiate connection in udp mode
+---------- you get kicked after TIMEOUT seconds with no packets in udp mode
+---------- send an `N\n` packet before TIMEOUT seconds are elapsed, to prevent getting your id reaped
+------ protocol:
+---------- all server replies except M for incoming messages are suffixed with "\n"
+----------
+---------- connect
+-------------- if in udp mode, server sends `T<timeout_sec>\n` to tell you when timeout will occur.
+-------------- server sends `I<id>`
+-------------- this is your id
+-------------- if server is full, it instead sends "full"
+---------- send
+-------------- `P<8-byte password>`
+------------------ server will only broadcast to you messages starting with that password
+------------------ the password will be stripped before it is sent to you
+-------------- `M<BUFS_max-byte buffer>`
+------------------ broadcast message to all clients
+------------------ only clients with `password==first_8_bytes` receive your msg
+------------------ message shorter than 8 returns "Err"
+------------------ default password is "00000000" ASCII
+-------------- `S<BUFS_max-byte buffer>`
+------------------ broadcast using default password "00000000"
+-------------- `N<whatever>`
+------------------ does nothing. use this for initiating your UDP connection.
+---------- receive
+---------------- when you send P for password
+-------------------- server sends back "p_ok"
+---------------- when you send M or S for message
+-------------------- server sends back "m_ok"
+---------------- when server is shutting down
+-------------------- server sends "bye"
+-------------------- close the socket when you get this msg
+---------------- when admin sends you msg
+-------------------- server sends `A<BUFS-1_max-byte buffer>`
+-------------------- this message is sent to all, regardless of password
+*/
+
+// ======== change these ========
+#if 1
+    // max buffer size
+    #define UDP_BUFS 1024
+    #define TCP_BUFS 8192
+
+    // rate limit: BUFS bytes / RLM seconds
+    #define RATE_LIMIT_S 1
+
+    // max clients count
+    #define MAX_CLIENTS 8192
+
+    // client timeout for udp mode in seconds
+    #define CLIENT_TIMEOUT 30
+#endif
+// ==============================
+
+#if UDP == 1
+    #define BUFS UDP_BUFS
+#else
+    #define BUFS TCP_BUFS
+#endif
+
+#include <errno.h>
+#include <fcntl.h>
+#include <inttypes.h>
+#include <netdb.h>
+#include <poll.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/socket.h>
+#include <sys/uio.h>
+#include <time.h>
+#include <unistd.h>
+
+#define p_stringify(...) #__VA_ARGS__
+#define stringify(...)   p_stringify(__VA_ARGS__)
+
+#if DEBUG == 1
+    #define log(...)  printf(__VA_ARGS__)
+    #define verb(...) __VA_ARGS__
+#else
+    #define log(...)
+    #define verb(...)
+#endif
+
+#if UDP == 1
+    #define onlyudp(...) __VA_ARGS__
+    #define onlytcp(...)
+#else
+    #define onlyudp(...)
+    #define onlytcp(...) __VA_ARGS__
+#endif
+
+int env_max_clients  = MAX_CLIENTS;
+int env_bufs         = BUFS;
+int env_rate_limit_s = RATE_LIMIT_S;
+onlyudp(int env_client_timeout = CLIENT_TIMEOUT);
+
+// stdin nonblocking
+#define NB_EOF   -1
+#define NB_AGAIN -2
+
+static char   nb_buffer[ BUFSIZ ];
+static size_t nb_buf_pos  = 0;
+static size_t nb_buf_fill = 0;
+
+#define fold inline __attribute__((always_inline)) static
+
+fold void set_stdin_nonblocking(void) {
+    static int done = 0;
+    if (!done) {
+        int flags = fcntl(STDIN_FILENO, F_GETFL);
+        if (flags != -1) { fcntl(STDIN_FILENO, F_SETFL, flags | O_NONBLOCK); }
+        done = 1;
+    }
+}
+
+fold int nb_getstring(char str[], size_t max_size) {
+    if (max_size <= 1) {
+        if (max_size == 1) str[ 0 ] = '\0';
+        return 0;
+    }
+
+    size_t pos = 0;
+
+    // drain buffer
+    size_t remaining = max_size - 1;
+    size_t buf_avail = nb_buf_fill - nb_buf_pos;
+    size_t copy_len  = buf_avail < remaining ? buf_avail : remaining;
+    if (copy_len > 0) {
+        memcpy(str + pos, nb_buffer + nb_buf_pos, copy_len);
+        nb_buf_pos += copy_len;
+        pos += copy_len;
+    }
+
+    if (pos < max_size - 1 && nb_buf_pos >= nb_buf_fill) {
+        ssize_t n;
+        do { n = read(STDIN_FILENO, nb_buffer, sizeof(nb_buffer)); } while (n == -1 && errno == EINTR);
+
+        if (n <= 0) {
+            if (n == 0) {
+                if (pos == 0) return NB_EOF;
+            } else { // n == -1
+                if (errno == EAGAIN || errno == EWOULDBLOCK) {
+                    if (pos == 0) return NB_AGAIN;
+                } else {
+                    if (pos == 0) return NB_EOF;
+                }
+            }
+        } else {
+            nb_buf_pos  = 0;
+            nb_buf_fill = (size_t) n;
+
+            remaining = max_size - 1 - pos;
+            buf_avail = nb_buf_fill - nb_buf_pos;
+            copy_len  = buf_avail < remaining ? buf_avail : remaining;
+            if (copy_len > 0) {
+                memcpy(str + pos, nb_buffer + nb_buf_pos, copy_len);
+                nb_buf_pos += copy_len;
+                pos += copy_len;
+            }
+        }
+    }
+
+    str[ pos ] = '\0';
+    return (int) pos;
+}
+
+#define likely(x)   __builtin_expect(!!(x), 1)
+#define unlikely(x) __builtin_expect(!!(x), 0)
+
+uint64_t default_pass;
+
+onlytcp(struct client {
+    int      fd;       // descriptor
+    uint64_t id;       // identifier
+    uint64_t pass;     // password
+    time_t   last_rl;  // ratelimit timestamp
+    int      rl_bytes; // bytes used
+});
+onlyudp(struct client {
+    struct sockaddr_storage addr;      // client address
+    socklen_t               addrlen;   // address len
+    uint64_t                id;        // identifier
+    uint64_t                pass;      // password
+    time_t                  last_rl;   // ratelimit timestamp
+    int                     rl_bytes;  // bytes used
+    time_t                  last_seen; // for kicking
+});
+struct client* clients = NULL;
+int            active  = 0;
+int            cap     = 0;
+uint64_t       nid     = 1; // next id to assign
+
+// accept new client
+onlytcp(void add_client(int fd) {
+    if (active >= env_max_clients) {
+        // refuse: too many clients
+        write(fd, "full\n", 5);
+        close(fd);
+        return;
+    }
+    if (active >= cap) {
+        cap = cap ? cap * 2 : 64;
+        if (cap > env_max_clients) cap = env_max_clients;
+        clients = realloc(clients, cap * sizeof(struct client));
+    }
+    onlytcp(clients[ active ].fd = fd);
+    clients[ active ].id       = nid++;
+    clients[ active ].pass     = default_pass;
+    clients[ active ].last_rl  = time(NULL);
+    clients[ active ].rl_bytes = 0;
+    char buf[ 64 ];
+    int  len = sprintf(buf, "I%" PRIu64 "\n", clients[ active ].id);
+    write(fd, buf, len);
+    active++;
+    log("clients:%d\n", active);
+});
+
+// disconnect client at index i
+void del_client(int i) {
+    onlytcp(close(clients[ i ].fd));
+    if (i < active - 1) clients[ i ] = clients[ active - 1 ];
+    active--;
+    log("clients:%d\n", active);
+}
+
+int sock;
+
+// broadcast to all clients whose password matches
+void broadcast_msg(int si, char* payload, int msg_len, int use_default) {
+    char pay_buf[ 64 ];
+    int  pay_len = sprintf(pay_buf, "M%" PRIu64 ":", clients[ si ].id);
+    int  sent    = 0;
+
+    uint64_t p;
+    char*    body;
+    int      body_len;
+
+    if (use_default) {
+        p        = default_pass;
+        body     = payload;
+        body_len = msg_len;
+    } else {
+        if (msg_len < 8) return;
+        memcpy(&p, payload, 8);
+        body     = payload + 8;
+        body_len = msg_len - 8;
+    }
+
+    struct iovec iov[ 2 ] = { { pay_buf, (size_t) pay_len }, { body, (size_t) body_len } };
+
+    for (int j = 0; j < active; j++) {
+        if (j != si && p == clients[ j ].pass) {
+            onlytcp(writev(clients[ j ].fd, iov, 2));
+            onlyudp({
+                char out[ 64 + env_bufs ];
+                memcpy(out, pay_buf, pay_len);
+                memcpy(out + pay_len, body, body_len);
+                sendto(sock, out, pay_len + body_len, 0, (struct sockaddr*) &clients[ j ].addr, clients[ j ].addrlen);
+            });
+            sent++;
+        }
+    }
+    if (sent > 0) log("sentto:%d\n", sent);
+}
+
+char *admin_msg;
+
+int main(int argc, char** argv) {
+    if (argc < 3) {
+        printf("more args.\n%s <ip> <port>\n", argv[ 0 ]);
+        exit(1);
+    }
+
+    admin_msg = malloc(env_bufs);
+
+    onlytcp(printf("hi from tcp;" verb(" verbose mode") "\n"));
+    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");
+        });
+    }
+
+    set_stdin_nonblocking();
+
+    admin_msg[ 0 ] = 'A';
+
+    memcpy(&default_pass, "00000000", 8);
+
+    struct addrinfo hint = {}, *result;
+    hint.ai_family       = AF_UNSPEC;
+    onlytcp(hint.ai_socktype = SOCK_STREAM);
+    onlyudp(hint.ai_socktype = SOCK_DGRAM);
+    hint.ai_flags = AI_PASSIVE;
+    if (getaddrinfo(argv[ 1 ], argv[ 2 ], &hint, &result) != 0) {
+        perror("getaddrinfo failed");
+        exit(1);
+    }
+
+    sock = socket(result->ai_family, result->ai_socktype, result->ai_protocol);
+
+    int option = 1;
+    setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &option, sizeof(option));
+    if (bind(sock, result->ai_addr, result->ai_addrlen) < 0) {
+        perror("bind failed");
+        exit(1);
+    }
+
+    freeaddrinfo(result);
+    onlytcp(listen(sock, 32));
+    fcntl(sock, F_SETFL, fcntl(sock, F_GETFL, 0) | O_NONBLOCK);
+
+    // fds[0] = server socket, fds[1] = stdin, fds[2..] = clients
+    struct pollfd fds[ env_max_clients + 2 ];
+    fds[ 1 ].fd     = STDIN_FILENO;
+    fds[ 1 ].events = POLLIN;
+
+    while (1) {
+        int run_shutdown = 0;
+        int got_msg      = 0;
+        int resval       = 0;
+
+        fds[ 0 ].fd     = sock;
+        fds[ 0 ].events = POLLIN;
+        int n           = 2;
+        onlytcp(for (int i = 0; i < active; i++) {
+            fds[ n ].fd     = clients[ i ].fd;
+            fds[ n ].events = POLLIN;
+            n++;
+        });
+
+        poll(fds, n, onlyudp(env_client_timeout * 1000) onlytcp(-1));
+
+        time_t now = time(NULL);
+
+        // sweep udp connections
+        onlyudp(for (int i = active - 1; i >= 0; i--) {
+            if (now - clients[ i ].last_seen > env_client_timeout) {
+                log("evict:%" PRIu64 "\n", clients[ i ].id);
+                del_client(i);
+            }
+        });
+
+        // check admin stdin
+        if (fds[ 1 ].revents & POLLIN) {
+            resval = nb_getstring(admin_msg + 1, env_bufs - 1);
+            if (resval == NB_EOF) {
+                printf("dying\n");
+                run_shutdown = 1;
+            } else if (resval > 0) {
+                got_msg = 1;
+            }
+        }
+
+        // shutdown: notify all clients and exit
+        if (unlikely(run_shutdown)) {
+            for (int i = 0; i < active; i++) {
+                onlytcp(write(clients[ i ].fd, "bye\n", 4));
+                onlytcp(close(clients[ i ].fd));
+                onlyudp(sendto(sock, "bye\n", 4, 0, (struct sockaddr*) &clients[ i ].addr, clients[ i ].addrlen));
+            }
+            active = 0;
+            break;
+        }
+
+        // accept new connections
+        onlytcp(if (fds[ 0 ].revents & POLLIN) {
+            int c;
+            while ((c = accept(sock, 0, 0)) >= 0) add_client(c);
+        });
+
+        onlyudp(
+            struct sockaddr_storage udp_src; //
+            socklen_t               udp_srclen = sizeof(udp_src);
+            char                    buf[ env_bufs ];
+            int                     rd;
+            int                     j = -1);
+
+        // process existing clients
+        onlytcp(for (int i = n - 1; i >= 2; i--)) //
+            onlyudp(
+                while (fds[ 0 ].revents & POLLIN && //
+                       (udp_srclen = sizeof(udp_src))
+                       && (rd = recvfrom(sock, buf, sizeof(buf), 0, (struct sockaddr*) &udp_src, &udp_srclen)) > 0)) //
+        {
+            onlytcp(int j = i - 2);
+
+            onlyudp(
+                // find existing client
+                j = -1; for (int k = 0; k < active; k++) {
+                    if (clients[ k ].addrlen == udp_srclen && memcmp(&clients[ k ].addr, &udp_src, udp_srclen) == 0) {
+                        j = k;
+                        break;
+                    }
+                }
+                // register new client
+                if (j < 0) {
+                    if (active >= env_max_clients) {
+                        sendto(sock, "full\n", 5, 0, (struct sockaddr*) &udp_src, udp_srclen);
+                        continue;
+                    }
+                    if (active >= cap) {
+                        cap = cap ? cap * 2 : 64;
+                        if (cap > env_max_clients) cap = env_max_clients;
+                        clients = realloc(clients, cap * sizeof(struct client));
+                    }
+                    j                     = active;
+                    clients[ j ].addr     = udp_src;
+                    clients[ j ].addrlen  = udp_srclen;
+                    clients[ j ].id       = nid++;
+                    clients[ j ].pass     = default_pass;
+                    clients[ j ].last_rl  = now;
+                    clients[ j ].rl_bytes = 0;
+                    active++;
+                    // send timeout amount
+                    sendto(
+                        sock,
+                        "T" stringify(env_client_timeout) "\n",
+                        strlen(stringify(env_client_timeout)) + 2,
+                        0,
+                        (struct sockaddr*) &udp_src,
+                        udp_srclen);
+                    char ibuf[ 32 ];
+                    int  ilen = sprintf(ibuf, "I%" PRIu64 "\n", clients[ j ].id);
+                    sendto(sock, ibuf, ilen, 0, (struct sockaddr*) &udp_src, udp_srclen);
+
+                    log("clients:%d\n", active);
+                }
+
+                onlyudp(clients[ j ].last_seen = now); //
+            );
+
+            onlytcp(if (likely(fds[ i ].revents & POLLIN))) {
+                onlytcp(char buf[ env_bufs ]; int rd = read(fds[ i ].fd, buf, sizeof(buf)); if (rd <= 0) {
+                    del_client(j);
+                    if (got_msg && j < active) write(clients[ j ].fd, admin_msg, resval + 1);
+                    continue;
+                });
+
+                if (buf[ 0 ] == 'P' && rd >= 9) {
+                    // set password
+                    memcpy(&clients[ j ].pass, buf + 1, 8);
+                    onlytcp(write(clients[ j ].fd, "p_ok\n", 5));
+
+                } else if (buf[ 0 ] == 'M' && rd >= 9) {
+                    // M: broadcast with password
+                    if (likely(now - clients[ j ].last_rl >= env_rate_limit_s)) {
+                        clients[ j ].last_rl  = now;
+                        clients[ j ].rl_bytes = 0;
+                    }
+                    int plen = rd - 1;
+                    if (unlikely(clients[ j ].rl_bytes + plen > env_bufs)) {
+                        onlytcp(write(clients[ j ].fd, "Rate\n", 5));
+                        onlyudp(sendto(sock, "Rate\n", 5, 0, (struct sockaddr*) &udp_src, udp_srclen));
+                    } else {
+                        clients[ j ].rl_bytes += plen;
+                        broadcast_msg(j, buf + 1, plen, 0);
+                        onlytcp(write(clients[ j ].fd, "m_ok\n", 5));
+                    }
+
+                } else if (buf[ 0 ] == 'S' && rd >= 2) {
+                    // S: broadcast with default password
+                    if (likely(now - clients[ j ].last_rl >= env_rate_limit_s)) {
+                        clients[ j ].last_rl  = now;
+                        clients[ j ].rl_bytes = 0;
+                    }
+                    int plen = rd - 1;
+                    if (unlikely(clients[ j ].rl_bytes + plen > env_bufs)) {
+                        onlytcp(write(clients[ j ].fd, "Rate\n", 5));
+                        onlyudp(sendto(sock, "Rate\n", 5, 0, (struct sockaddr*) &udp_src, udp_srclen));
+                    } else {
+                        clients[ j ].rl_bytes += plen;
+                        broadcast_msg(j, buf + 1, plen, 1);
+                        onlytcp(write(clients[ j ].fd, "m_ok\n", 6));
+                    }
+
+                } else if (buf[ 0 ] == 'N') {
+                    // no-op
+                } else {
+                    onlytcp(write(clients[ j ].fd, "Err\n", 4));
+                    onlyudp(sendto(sock, "Err\n", 4, 0, (struct sockaddr*) &udp_src, udp_srclen));
+                }
+            }
+
+            // send admin message
+            onlytcp(if (unlikely(got_msg)) write(clients[ j ].fd, admin_msg, resval + 1));
+        }
+
+        onlyudp(if (unlikely(got_msg)) {
+            for (int i = 0; i < active; i++)
+                sendto(sock, admin_msg, resval + 1, 0, (struct sockaddr*) &clients[ i ].addr, clients[ i ].addrlen);
+        });
+    }
+
+    close(sock);
+    free(clients);
+    return 0;
+}
\ No newline at end of file