From bbb8b121bae59bef1c1aaf0bacdc37cc81ac2bed Mon Sep 17 00:00:00 2001
From: Virtually Nick <vnick@apache.org>
Date: Thu, 29 Aug 2024 16:46:06 -0400
Subject: [PATCH] GUACAMOLE-600: Correct handling of non-blocking socket for
 timeout.

---
 src/libguac/tcp.c | 75 +++++++++++++++++++++++++++++++++++------------
 1 file changed, 57 insertions(+), 18 deletions(-)

diff --git a/src/libguac/tcp.c b/src/libguac/tcp.c
index e4ae71f24..1351e7826 100644
--- a/src/libguac/tcp.c
+++ b/src/libguac/tcp.c
@@ -54,8 +54,7 @@ int guac_tcp_connect(const char* hostname, const char* port, const int timeout)
     }
 
     /* Attempt connection to each address until success */
-    current_address = addresses;
-    while (current_address != NULL) {
+    for (current_address = addresses; current_address != NULL; current_address = current_address->ai_next) {
 
         /* Resolve hostname */
         if ((retval = getnameinfo(current_address->ai_addr,
@@ -69,40 +68,80 @@ int guac_tcp_connect(const char* hostname, const char* port, const int timeout)
             continue;
         }
 
-        /* Get socket */
+        /* Get socket or return the error. */
         fd = socket(current_address->ai_family, SOCK_STREAM, 0);
         if (fd < 0) {
             freeaddrinfo(addresses);
             return fd;
         }
 
-        /* Set socket to non-blocking */
-        fcntl(fd, F_SETFL, O_NONBLOCK);
+        /* Variable to store current socket options. */
+        int opt;
+
+        /* Get current socket options */
+        if ((opt = fcntl(fd, F_GETFL, NULL)) < 0) {
+            guac_error = GUAC_STATUS_INVALID_ARGUMENT;
+            guac_error_message = "Failed to retrieve socket options.";
+            close(fd);
+            continue;
+        }
 
-        /* Set up timeout. */
-        fd_set fdset;
-        FD_ZERO(&fdset);
-        FD_SET(fd, &fdset);
+        /* Set socket to non-blocking */
+        if (fcntl(fd, F_SETFL, opt | O_NONBLOCK) < 0) {
+            guac_error = GUAC_STATUS_INVALID_ARGUMENT;
+            guac_error_message = "Failed to set non-blocking socket.";
+            close(fd);
+            continue;
+        }
 
+        /* Structure that stores our timeout setting. */
         struct timeval tv;
         tv.tv_sec = timeout;
         tv.tv_usec = 0;
 
         /* Connect and wait for timeout */
-        if (connect(fd, current_address->ai_addr, current_address->ai_addrlen) < 0) {
-            guac_error = GUAC_STATUS_REFUSED;
-            guac_error_message = "Unable to connect via socket.";
-            close(fd);
-            break;
+        if ((retval = connect(fd, current_address->ai_addr, current_address->ai_addrlen)) < 0) {
+            if (errno == EINPROGRESS) {
+                /* Set up timeout. */
+                fd_set fdset;
+                FD_ZERO(&fdset);
+                FD_SET(fd, &fdset);
+                
+                retval = select(fd + 1, NULL, &fdset, NULL, &tv);
+            }
+            
+            else {
+                guac_error = GUAC_STATUS_REFUSED;
+                guac_error_message = "Unable to connect via socket.";
+                close(fd);
+                continue;
+            }
         }
 
-        /* Check for the connection and break if successful */
-        if (select(fd + 1, NULL, &fdset, NULL, &tv) > 0)
+        /* Successful connection */
+        if (retval > 0) {
+            /* Restore previous socket options. */
+            if (fcntl(fd, F_SETFL, opt) < 0) {
+                guac_error = GUAC_STATUS_INVALID_ARGUMENT;
+                guac_error_message = "Failed to set non-blocking socket.";
+                close(fd);
+                continue;
+            }
+
             break;
+        }
+
+        if (retval == 0) {
+            guac_error = GUAC_STATUS_REFUSED;
+            guac_error_message = "Timeout connecting via socket.";
+        }
+        else {
+            guac_error = GUAC_STATUS_INVALID_ARGUMENT;
+            guac_error_message = "Error attempting to connect via socket.";
+        }
 
-        /* Connection not successful - free resources and go to the next address. */
+        /* Some error has occurred - free resources before next iteration. */
         close(fd);
-        current_address = current_address->ai_next;
 
     }