koekeishiya#1107 fix rare occurence where an incoming socket connecti…

…on just drops

CHANGELOG.md

@@ -9,6 +9,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 - New config *window_origin_display* to specify which display a window should become managed at upon creation [#951](https://github.com/koekeishiya/yabai/issues/951)
 
 ### Changed
+- Fixed an issue with the way unix sockets were handled that would cause an incoming connection to drop in rare occasions [#1107](https://github.com/koekeishiya/yabai/issues/1107)
 - Update scripting addition to support macOS 12.1.0 [#1054](https://github.com/koekeishiya/yabai/issues/1054)
 - Fixed an issue that would cause the target window to snap back to its previous position when moved between displays using the cursor inside mission-control [#820](https://github.com/koekeishiya/yabai/issues/820)
 - Properly drain autoreleased objects (from Apple frameworks) [#751](https://github.com/koekeishiya/yabai/issues/751)

src/event.c

@@ -1187,9 +1187,11 @@ static EVENT_CALLBACK(EVENT_HANDLER_DAEMON_MESSAGE)
 
             fflush(rsp);
             fclose(rsp);
+
+            return EVENT_SUCCESS;
         }
     }
 
     socket_close(param1);
-    return EVENT_SUCCESS;
+    return EVENT_FAILURE;
 }

src/manifest.m

@@ -39,7 +39,6 @@
 #define HASHTABLE_IMPLEMENTATION
 #include "misc/hashtable.h"
 #undef HASHTABLE_IMPLEMENTATION
-#include "misc/socket.h"
 
 #include "osax/common.h"
 

src/message.c

@@ -2495,6 +2495,8 @@ bool message_loop_begin(char *socket_path)
         return false;
     }
 
+    fcntl(g_message_loop.sockfd, F_SETFD, FD_CLOEXEC | fcntl(g_message_loop.sockfd, F_GETFD));
+
     g_message_loop.is_running = true;
     pthread_create(&g_message_loop.thread, NULL, &message_loop_run, NULL);
 

src/misc/helpers.h

@@ -3,6 +3,24 @@
 
 const CFStringRef kAXEnhancedUserInterface = CFSTR("AXEnhancedUserInterface");
 
+static inline bool socket_connect(int *sockfd, char *socket_path)
+{
+    struct sockaddr_un socket_address;
+    socket_address.sun_family = AF_UNIX;
+
+    *sockfd = socket(AF_UNIX, SOCK_STREAM, 0);
+    if (*sockfd == -1) return false;
+
+    snprintf(socket_address.sun_path, sizeof(socket_address.sun_path), "%s", socket_path);
+    return connect(*sockfd, (struct sockaddr *) &socket_address, sizeof(socket_address)) != -1;
+}
+
+static inline void socket_close(int sockfd)
+{
+    shutdown(sockfd, SHUT_RDWR);
+    close(sockfd);
+}
+
 static inline char *json_optional_bool(int value)
 {
     if (value == 0) return "null";

src/misc/macros.h

@@ -14,6 +14,8 @@
 #define in_range_ei(a, b, c) (((a) >  (b)) && ((a) <= (c)))
 #define in_range_ee(a, b, c) (((a) >  (b)) && ((a) <  (c)))
 
+#define FAILURE_MESSAGE "\x07"
+
 #define MAXLEN 512
 
 #define REGEX_MATCH_UD  0

src/sa.m

@@ -194,8 +194,8 @@ static bool scripting_addition_request_handshake(char *version, uint32_t *attrib
     bool result = false;
     char rsp[BUFSIZ] = {};
 
-    if (socket_connect_un(&sockfd, g_sa_socket_file)) {
-        if (socket_write(sockfd, "handshake")) {
+    if (socket_connect(&sockfd, g_sa_socket_file)) {
+        if (send(sockfd, "handshake", 9, 0) != -1) {
             int length = recv(sockfd, rsp, sizeof(rsp)-1, 0);
             if (length <= 0) goto out;
 
@@ -540,11 +540,12 @@ int scripting_addition_load(void)
 static bool scripting_addition_run_command(char *message)
 {
     int sockfd;
+    char dummy;
     bool result = false;
 
-    if (socket_connect_un(&sockfd, g_sa_socket_file)) {
-        if (socket_write(sockfd, message)) {
-            socket_wait(sockfd);
+    if (socket_connect(&sockfd, g_sa_socket_file)) {
+        if (send(sockfd, message, strlen(message), 0) != -1) {
+            recv(sockfd, &dummy, 1, 0);
             result = true;
         }
     }

src/yabai.c

@@ -58,14 +58,6 @@ static int client_send_message(int argc, char **argv)
         error("yabai-msg: 'env USER' not set! abort..\n");
     }
 
-    int sockfd;
-    char socket_file[MAXLEN];
-    snprintf(socket_file, sizeof(socket_file), SOCKET_PATH_FMT, user);
-
-    if (!socket_connect_un(&sockfd, socket_file)) {
-        error("yabai-msg: failed to connect to socket..\n");
-    }
-
     int message_length = argc;
     int argl[argc];
 
@@ -74,7 +66,7 @@ static int client_send_message(int argc, char **argv)
         message_length += argl[i];
     }
 
-    char message[message_length];
+    char *message = malloc(message_length);
     char *temp = message;
 
     for (int i = 1; i < argc; ++i) {
@@ -84,36 +76,37 @@ static int client_send_message(int argc, char **argv)
     }
     *temp++ = '\0';
 
-    if (!socket_write_bytes(sockfd, message, message_length)) {
+    int sockfd;
+    char socket_file[MAXLEN];
+    snprintf(socket_file, sizeof(socket_file), SOCKET_PATH_FMT, user);
+
+    if (!socket_connect(&sockfd, socket_file)) {
+        error("yabai-msg: failed to connect to socket..\n");
+    }
+
+    if (send(sockfd, message, message_length, 0) == -1) {
         error("yabai-msg: failed to send data..\n");
     }
 
     shutdown(sockfd, SHUT_WR);
+    free(message);
 
     int result = EXIT_SUCCESS;
-    int byte_count = 0;
+    FILE *output = stdout;
+    int bytes_read = 0;
     char rsp[BUFSIZ];
 
-    struct pollfd fds[] = {
-        { sockfd, POLLIN, 0 }
-    };
+    while ((bytes_read = read(sockfd, rsp, sizeof(rsp)-1)) > 0) {
+        rsp[bytes_read] = '\0';
 
-    while (poll(fds, 1, -1) > 0) {
-        if (fds[0].revents & POLLIN) {
-            if ((byte_count = recv(sockfd, rsp, sizeof(rsp)-1, 0)) <= 0) {
-                break;
-            }
-
-            rsp[byte_count] = '\0';
-
-            if (rsp[0] == FAILURE_MESSAGE[0]) {
-                result = EXIT_FAILURE;
-                fprintf(stderr, "%s", rsp + 1);
-                fflush(stderr);
-            } else {
-                fprintf(stdout, "%s", rsp);
-                fflush(stdout);
-            }
+        if (rsp[0] == FAILURE_MESSAGE[0]) {
+            result = EXIT_FAILURE;
+            output = stderr;
+            fprintf(output, "%s", rsp + 1);
+            fflush(output);
+        } else {
+            fprintf(output, "%s", rsp);
+            fflush(output);
         }
     }