redesign, broadcast server crash fix

source/server.cpp

@@ -26,20 +26,36 @@ static TransferState    g_server_state;
 static std::atomic<int> g_server_client_sock{-1};
 static std::atomic<int> g_server_listen_sock{-1};
 static std::atomic<int> g_broadcast_sock{-1};
+static std::atomic<bool> g_accept_thread_active{false};
+static std::atomic<bool> g_broadcast_thread_active{false};
+static pthread_t         g_broadcast_thread{};
 
 bool isServerTransferDone()       { return g_server_state.done.load(); }
 bool isServerTransferCancelled()  { return g_server_state.cancelled.load(); }
+bool isServerWorkersIdle()        { return !g_accept_thread_active.load() && !g_broadcast_thread_active.load(); }
 double getServerProgress()        { return g_server_state.progress(); }
 std::string getServerStatusText() { return g_server_state.getStatus(); }
 
 void cancelServerTransfer() {
     g_server_state.cancelled.store(true);
-    int sock = g_server_client_sock.load();
-    if (sock >= 0) shutdown(sock, SHUT_RDWR);
-    int lsock = g_server_listen_sock.load();
-    if (lsock >= 0) shutdown(lsock, SHUT_RDWR);
-    int bsock = g_broadcast_sock.load();
-    if (bsock >= 0) shutdown(bsock, SHUT_RDWR);
+    int sock = g_server_client_sock.exchange(-1);
+    if (sock >= 0) {
+        shutdown(sock, SHUT_RDWR);
+        close(sock);
+    }
+    int lsock = g_server_listen_sock.exchange(-1);
+    if (lsock >= 0) {
+        shutdown(lsock, SHUT_RDWR);
+        close(lsock);
+    }
+    int bsock = g_broadcast_sock.exchange(-1);
+    if (bsock >= 0) {
+        shutdown(bsock, SHUT_RDWR);
+        close(bsock);
+    }
+    if (g_broadcast_thread_active.load()) {
+        pthread_cancel(g_broadcast_thread);
+    }
 }
 
 #ifdef __SWITCH__
@@ -166,13 +182,17 @@ static void* handle_client(void* socket_desc) {
         receive_file(client_socket, filename_str, file_size);
     }
 
-    close(client_socket);
+    int owned_client = g_server_client_sock.exchange(-1);
+    if (owned_client == client_socket) {
+        close(client_socket);
+    }
     return nullptr;
 }
 
 struct AcceptArgs { int server_fd; };
 
 static void* accept_and_handle(void* arg) {
+    g_accept_thread_active.store(true);
     int server_fd = static_cast<AcceptArgs*>(arg)->server_fd;
     delete static_cast<AcceptArgs*>(arg);
 
@@ -180,8 +200,10 @@ static void* accept_and_handle(void* arg) {
     sockaddr_in client_addr{};
     socklen_t   client_len = sizeof(client_addr);
     int client_socket = accept(server_fd, (sockaddr*)&client_addr, &client_len);
-    g_server_listen_sock.store(-1);
-    close(server_fd);
+    int owned_listen = g_server_listen_sock.exchange(-1);
+    if (owned_listen == server_fd) {
+        close(server_fd);
+    }
 
     if (client_socket >= 0) {
         g_server_client_sock.store(client_socket);
@@ -191,23 +213,27 @@ static void* accept_and_handle(void* arg) {
         } else {
             close(client_socket);
         }
-        g_server_client_sock.store(-1);
     }
 
     g_server_state.done.store(true);
+    g_accept_thread_active.store(false);
     return nullptr;
 }
 
 static void* broadcast_listener(void* arg) {
-    Socket udp(socket(AF_INET, SOCK_DGRAM, 0));
-    if (!udp.valid()) {
+    g_broadcast_thread_active.store(true);
+    pthread_setcancelstate(PTHREAD_CANCEL_ENABLE, nullptr);
+    pthread_setcanceltype(PTHREAD_CANCEL_DEFERRED, nullptr);
+    int udp = socket(AF_INET, SOCK_DGRAM, 0);
+    if (udp < 0) {
         perror("broadcast_listener: socket");
+        g_broadcast_thread_active.store(false);
         return nullptr;
     }
 
-    g_broadcast_sock.store(udp.fd);
+    g_broadcast_sock.store(udp);
 
-    struct timeval tv{0, 100000};  // 100ms poll so cancel is detected quickly
+    struct timeval tv{0, 20000};  // 20ms poll so cancel/exit wins race with socketExit
     setsockopt(udp, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
 
     sockaddr_in addr{};
@@ -217,7 +243,9 @@ static void* broadcast_listener(void* arg) {
 
     if (bind(udp, (sockaddr*)&addr, sizeof(addr)) < 0) {
         perror("broadcast_listener: bind");
-        g_broadcast_sock.store(-1);
+        int owned = g_broadcast_sock.exchange(-1);
+        if (owned == udp) close(udp);
+        g_broadcast_thread_active.store(false);
         return nullptr;
     }
 
@@ -226,7 +254,9 @@ static void* broadcast_listener(void* arg) {
     group.imr_interface.s_addr = htonl(INADDR_ANY);
     if (setsockopt(udp, IPPROTO_IP, IP_ADD_MEMBERSHIP, &group, sizeof(group)) < 0) {
         perror("broadcast_listener: setsockopt");
-        g_broadcast_sock.store(-1);
+        int owned = g_broadcast_sock.exchange(-1);
+        if (owned == udp) close(udp);
+        g_broadcast_thread_active.store(false);
         return nullptr;
     }
 
@@ -250,21 +280,28 @@ static void* broadcast_listener(void* arg) {
         }
     }
 
-    g_broadcast_sock.store(-1);
+    int owned = g_broadcast_sock.exchange(-1);
+    if (owned == udp) close(udp);
+    g_broadcast_thread_active.store(false);
     return nullptr;
 }
 
 int startSendingThread() {
+    g_server_state.reset();
+    g_server_state.setStatus("Waiting for connection...");
+
     pthread_t broadcast_thread;
     if (pthread_create(&broadcast_thread, nullptr, broadcast_listener, nullptr) != 0) {
         perror("startSendingThread: broadcast thread");
         return 1;
     }
+    g_broadcast_thread = broadcast_thread;
     pthread_detach(broadcast_thread);
 
     Socket server(socket(AF_INET, SOCK_STREAM, 0));
     if (!server.valid()) {
         perror("startSendingThread: socket");
+        cancelServerTransfer();
         return 1;
     }
 
@@ -278,20 +315,20 @@ int startSendingThread() {
 
     if (bind(server, (sockaddr*)&addr, sizeof(addr)) < 0) {
         perror("startSendingThread: bind");
+        cancelServerTransfer();
         return 1;
     }
     if (listen(server, 3) < 0) {
         perror("startSendingThread: listen");
+        cancelServerTransfer();
         return 1;
     }
 
-    g_server_state.reset();
-    g_server_state.setStatus("Waiting for connection...");
-
     AcceptArgs* acc_args = new AcceptArgs{server.fd};
     pthread_t   accept_thread;
     if (pthread_create(&accept_thread, nullptr, accept_and_handle, acc_args) != 0) {
         delete acc_args;
+        cancelServerTransfer();
         return 1;
     }
     pthread_detach(accept_thread);