libvncserver: add WebSockets handshake mode

CMakeLists.txt

@@ -754,6 +754,13 @@ if(LIBVNCSERVER_WITH_WEBSOCKETS AND WITH_LIBVNCSERVER)
   set_target_properties(test_wstest PROPERTIES OUTPUT_NAME wstest)
   set_target_properties(test_wstest PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/test)
   target_link_libraries(test_wstest vncserver ${ADDITIONAL_TEST_LIBS})
+
+  add_executable(test_wscheckmode_test
+    ${TESTS_DIR}/wscheckmode_test.c
+    )
+  set_target_properties(test_wscheckmode_test PROPERTIES OUTPUT_NAME wscheckmode_test)
+  set_target_properties(test_wscheckmode_test PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/test)
+  target_link_libraries(test_wscheckmode_test vncserver ${ADDITIONAL_TEST_LIBS})
 endif(LIBVNCSERVER_WITH_WEBSOCKETS AND WITH_LIBVNCSERVER)
 
 if(WITH_LIBVNCSERVER)
@@ -772,6 +779,7 @@ if(WITH_JPEG AND FOUND_LIBJPEG_TURBO)
 endif(WITH_JPEG AND FOUND_LIBJPEG_TURBO)
 if(LIBVNCSERVER_WITH_WEBSOCKETS AND WITH_LIBVNCSERVER)
     add_test(NAME wstest COMMAND test_wstest)
+    add_test(NAME wscheckmode COMMAND test_wscheckmode_test)
 endif(LIBVNCSERVER_WITH_WEBSOCKETS AND WITH_LIBVNCSERVER)
 
 endif(WITH_TESTS)

README.md

@@ -191,6 +191,18 @@ a URL to point your web browser to. There, you can click on the noVNC-Button to
 connect using the noVNC viewer git submodule (installable via
 `git submodule update --init`).
 
+By default, LibVNCServer auto-detects whether an incoming connection starts
+with a WebSocket handshake or with the regular RFB handshake. If the expected
+transport is known, the detection mode can be selected explicitly:
+
+    ../examples/example -websocketmode auto
+    ../examples/example -websocketmode rfb
+    ../examples/example -websocketmode ws
+
+`auto` keeps the default behavior. `rfb` skips WebSocket probing and treats the
+connection as regular RFB immediately. `ws` expects a WebSocket handshake before
+starting the RFB session.
+
 ### Using Secure Websockets
 
 If you don't already have an SSL cert that's trusted by your browser, the most

include/rfb/rfb.h

@@ -204,6 +204,12 @@ typedef struct _rfbExtensionData {
  * rfbProcessEvents for each of these.
  */
 
+typedef enum {
+    rfbWebSocketsHandshakeAuto = 0,
+    rfbWebSocketsHandshakeRfb,
+    rfbWebSocketsHandshakeWebSockets
+} rfbWebSocketsHandshakeMode;
+
 typedef struct _rfbScreenInfo
 {
     /** this structure has children that are scaled versions of this screen */
@@ -351,6 +357,8 @@ typedef struct _rfbScreenInfo
     rfbXvpHookPtr xvpHook;
     char *sslkeyfile;
     char *sslcertfile;
+    /** How WebSockets connections are detected on the RFB socket. */
+    rfbWebSocketsHandshakeMode webSocketsHandshakeMode;
     int ipv6port; /**< The port to listen on when using IPv6.  */
     char* listen6Interface;
     /* We have an additional IPv6 listen socket since there are systems that

src/libvncserver/cargs.c

@@ -46,6 +46,7 @@ rfbUsage(void)
 #ifdef LIBVNCSERVER_WITH_WEBSOCKETS
     fprintf(stderr, "-sslkeyfile path       set path to private key file for encrypted WebSockets connections\n");
     fprintf(stderr, "-sslcertfile path      set path to certificate file for encrypted WebSockets connections\n");
+    fprintf(stderr, "-websocketmode mode    WebSockets detection mode: auto, rfb, ws\n");
 #endif
     fprintf(stderr, "-httpdir dir-path      enable http server using dir-path home\n");
     fprintf(stderr, "-httpport portnum      use portnum for http connection\n");
@@ -215,6 +216,23 @@ rfbProcessArguments(rfbScreenInfoPtr rfbScreen,int* argc, char *argv[])
 		return FALSE;
 	    }
             rfbScreen->sslcertfile = argv[++i];
+        } else if (strcmp(argv[i], "-websocketmode") == 0) {  /* -websocketmode auto|rfb|ws */
+            if (i + 1 >= *argc) {
+		rfbUsage();
+		return FALSE;
+	    }
+            i++;
+            if (strcmp(argv[i], "auto") == 0) {
+                rfbScreen->webSocketsHandshakeMode = rfbWebSocketsHandshakeAuto;
+            } else if (strcmp(argv[i], "rfb") == 0) {
+                rfbScreen->webSocketsHandshakeMode = rfbWebSocketsHandshakeRfb;
+            } else if (strcmp(argv[i], "ws") == 0 || strcmp(argv[i], "websocket") == 0) {
+                rfbScreen->webSocketsHandshakeMode = rfbWebSocketsHandshakeWebSockets;
+            } else {
+                rfbErr("invalid -websocketmode value: %s\n", argv[i]);
+                rfbUsage();
+                return FALSE;
+            }
 #endif
         } else {
 	    rfbProtocolExtension* extension;

src/libvncserver/main.c

@@ -971,6 +971,7 @@ rfbScreenInfoPtr rfbGetScreen(int* argc,char** argv,
    screen->httpListenSock=RFB_INVALID_SOCKET;
    screen->httpListen6Sock=RFB_INVALID_SOCKET;
    screen->httpSock=RFB_INVALID_SOCKET;
+   screen->webSocketsHandshakeMode = rfbWebSocketsHandshakeAuto;
 
    screen->desktopName = "LibVNCServer";
    screen->alwaysShared = FALSE;

src/libvncserver/websockets.c

@@ -118,12 +118,49 @@ webSocketsCheck (rfbClientPtr cl)
 {
     char bbuf[4], *scheme;
     int ret;
+    int timeout;
+    rfbWebSocketsHandshakeMode handshakeMode;
 
-    ret = rfbPeekExactTimeout(cl, bbuf, 4,
-                                   WEBSOCKETS_CLIENT_CONNECT_WAIT_MS);
+    handshakeMode = rfbWebSocketsHandshakeAuto;
+    if (cl->screen)
+        handshakeMode = cl->screen->webSocketsHandshakeMode;
+
+    if (handshakeMode == rfbWebSocketsHandshakeRfb) {
+        rfbLog("Normal socket connection\n");
+        return TRUE;
+    }
+
+    timeout = WEBSOCKETS_CLIENT_CONNECT_WAIT_MS;
+    if (handshakeMode == rfbWebSocketsHandshakeWebSockets) {
+        timeout = rfbMaxClientWait;
+        if (cl->screen && cl->screen->maxClientWait)
+            timeout = cl->screen->maxClientWait;
+    }
+
+    if (handshakeMode == rfbWebSocketsHandshakeAuto) {
+        ret = rfbPeekExactTimeout(cl, bbuf, 1, timeout);
+        if ((ret < 0) && (errno == ETIMEDOUT)) {
+            rfbLog("Normal socket connection\n");
+            return TRUE;
+        } else if (ret <= 0) {
+            rfbErr("webSocketsHandshake: unknown connection error\n");
+            return FALSE;
+        }
+
+        if (bbuf[0] != 'G' && bbuf[0] != '\x16' && bbuf[0] != '\x80') {
+            rfbLog("Normal socket connection\n");
+            return TRUE;
+        }
+    }
+
+    ret = rfbPeekExactTimeout(cl, bbuf, 4, timeout);
     if ((ret < 0) && (errno == ETIMEDOUT)) {
-      rfbLog("Normal socket connection\n");
-      return TRUE;
+      if (handshakeMode == rfbWebSocketsHandshakeWebSockets) {
+        rfbErr("webSocketsHandshake: timed out waiting for WebSockets header\n");
+        return FALSE;
+      }
+      rfbErr("webSocketsHandshake: timed out waiting for WebSockets header\n");
+      return FALSE;
     } else if (ret <= 0) {
       rfbErr("webSocketsHandshake: unknown connection error\n");
       return FALSE;
@@ -138,7 +175,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";

test/wscheckmode_test.c

@@ -0,0 +1,82 @@
+#include <rfb/rfb.h>
+
+#include <errno.h>
+#include <string.h>
+
+static int peekCalled;
+static int longPeekCalled;
+
+static int failIfPeeked(rfbClientPtr cl, char *buf, int len)
+{
+  (void)cl;
+  (void)buf;
+  (void)len;
+  peekCalled = 1;
+  errno = EAGAIN;
+  return -1;
+}
+
+static int slowRfbPeek(rfbClientPtr cl, char *buf, int len)
+{
+  (void)cl;
+
+  if (len == 1) {
+    buf[0] = 'R';
+    return 1;
+  }
+
+  longPeekCalled = 1;
+  errno = EAGAIN;
+  return -1;
+}
+
+static int testRfbModeSkipsPeek(void)
+{
+  rfbScreenInfo screen;
+  rfbClientRec client;
+
+  memset(&screen, 0, sizeof(screen));
+  memset(&client, 0, sizeof(client));
+
+  screen.webSocketsHandshakeMode = rfbWebSocketsHandshakeRfb;
+  client.screen = &screen;
+  client.peekAtSocket = failIfPeeked;
+  client.sock = RFB_INVALID_SOCKET;
+
+  peekCalled = 0;
+  if (!webSocketsCheck(&client))
+    return 1;
+
+  return peekCalled ? 1 : 0;
+}
+
+static int testAutoModeAcceptsSlowRfbPrefix(void)
+{
+  rfbScreenInfo screen;
+  rfbClientRec client;
+
+  memset(&screen, 0, sizeof(screen));
+  memset(&client, 0, sizeof(client));
+
+  screen.webSocketsHandshakeMode = rfbWebSocketsHandshakeAuto;
+  client.screen = &screen;
+  client.peekAtSocket = slowRfbPeek;
+  client.sock = 0;
+
+  longPeekCalled = 0;
+  if (!webSocketsCheck(&client))
+    return 1;
+
+  return longPeekCalled ? 1 : 0;
+}
+
+int main(void)
+{
+  if (testRfbModeSkipsPeek())
+    return 1;
+
+  if (testAutoModeAcceptsSlowRfbPrefix())
+    return 1;
+
+  return 0;
+}