From fca39ffb4ac6a72a83833688a66adc6c2001f3a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=B8ren=20Holm?= <sgh@sgh.dk>
Date: Fri, 29 Apr 2022 07:48:27 +0200
Subject: [PATCH] Added handshakeType to allow increase wait for ws/wss
 handshake.

For long latency clients then may not be able to deliver the websocket
handshake before we fall back to regular the RFB protocol.

So to not introduce a multisecond delay for all conenction a handshakeType
setting is introduced to tell which type of handshake we expect.

The problem is that a RFB client will be quiet "forever" until the server
sends it's signature. On the other hand a websocket client will send it's
HTTP websocket handshake immediately.

Before the server for sure know the the type of client the connect procedure
cannot continue. There is no other option than to wait for some amount of time
when the client type is auto detected.
---
 libvncserver/websockets.c | 21 ++++++++++++++-------
 rfb/rfb.h                 | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/libvncserver/websockets.c b/libvncserver/websockets.c
index 77791e5ff..2d657375f 100644
--- a/libvncserver/websockets.c
+++ b/libvncserver/websockets.c
@@ -74,8 +74,7 @@ Connection: Upgrade\r\n\
 Sec-WebSocket-Accept: %s\r\n\
 \r\n"
 
-#define WEBSOCKETS_CLIENT_CONNECT_WAIT_MS 100
-#define WEBSOCKETS_CLIENT_SEND_WAIT_MS 100
+#define WEBSOCKETS_CLIENT_CONNECT_WAIT_MS 500
 #define WEBSOCKETS_MAX_HANDSHAKE_LEN 4096
 
 #if defined(__linux__) && defined(NEED_TIMEVAL)
@@ -119,8 +118,17 @@ webSocketsCheck (rfbClientPtr cl)
     char bbuf[4], *scheme;
     int ret;
 
-    ret = rfbPeekExactTimeout(cl, bbuf, 4,
-                                   WEBSOCKETS_CLIENT_CONNECT_WAIT_MS);
+    int timeout = WEBSOCKETS_CLIENT_CONNECT_WAIT_MS;
+    switch (cl->screen->handshake_type) {
+        case RFB_HANDSHAKE_AUTO:
+            timeout = cl->screen->maxClientWait ? cl->screen->maxClientWait : rfbMaxClientWait;
+            break;
+        case RFB_HANDSHAKE_WEBSOCKET:
+            timeout = WEBSOCKETS_CLIENT_CONNECT_WAIT_MS;
+            break;
+    }
+
+    ret = rfbPeekExactTimeout(cl, bbuf, 4, timeout);
     if ((ret < 0) && (errno == ETIMEDOUT)) {
       rfbLog("Normal socket connection\n");
       return TRUE;
@@ -138,7 +146,7 @@ webSocketsCheck (rfbClientPtr cl)
 	  rfbErr("webSocketsHandshake: rfbssl_init failed\n");
 	  return FALSE;
 	}
-	ret = rfbPeekExactTimeout(cl, bbuf, 4, WEBSOCKETS_CLIENT_CONNECT_WAIT_MS);
+	ret = rfbPeekExactTimeout(cl, bbuf, 4, timeout);
         scheme = "wss";
     } else {
         scheme = "ws";
@@ -183,8 +191,7 @@ webSocketsHandshake(rfbClientPtr cl, char *scheme)
     }
 
     while (len < WEBSOCKETS_MAX_HANDSHAKE_LEN-1) {
-        if ((n = rfbReadExactTimeout(cl, buf+len, 1,
-                                     WEBSOCKETS_CLIENT_SEND_WAIT_MS)) <= 0) {
+        if ((n = rfbReadExact(cl, buf+len, 1)) <= 0) {
             if ((n < 0) && (errno == ETIMEDOUT)) {
                 break;
             }
diff --git a/rfb/rfb.h b/rfb/rfb.h
index b5f37e075..4e0b11285 100644
--- a/rfb/rfb.h
+++ b/rfb/rfb.h
@@ -100,6 +100,11 @@ enum rfbSocketState {
 	RFB_SOCKET_SHUTDOWN
 };
 
+enum rfbHandshakeType {
+	RFB_HANDSHAKE_AUTO,
+	RFB_HANDSHAKE_WEBSOCKET,
+};
+
 typedef void (*rfbKbdAddEventProcPtr) (rfbBool down, rfbKeySym keySym, struct _rfbClientRec* cl);
 typedef void (*rfbKbdReleaseAllKeysProcPtr) (struct _rfbClientRec* cl);
 typedef void (*rfbPtrAddEventProcPtr) (int buttonMask, int x, int y, struct _rfbClientRec* cl);
@@ -372,6 +377,16 @@ typedef struct _rfbScreenInfo
 #ifdef LIBVNCSERVER_HAVE_LIBZ
     rfbSetXCutTextUTF8ProcPtr setXCutTextUTF8;
 #endif
+    /** This setting specifies the handshake type to expect.
+     * WebSocket clients are detected by the presence of a special handshake.
+     * The handshake can only be detected by waiting some time for it. Normal
+     * RFB clients do not transmit an initial handshake.
+     *
+     * RFB_HANDSHAKE_AUTO - wait 100ms for WS handshake. Then fall back to normal RFB.
+     * RFB_HANDSHAKE_WEBSOCKET - wait 500ms for WS handshake.
+     *
+     */
+    enum rfbHandshakeType handshake_type;
 } rfbScreenInfo, *rfbScreenInfoPtr;