websockets: Move Hixie disconnect hack to websockets.c

Move the hixie disconnect hack to websockets.c. Removed
the remaining websockets vars from rfbClientPtr, so all
websockets stuff is hidden behind an opaque pointer.

libvncserver/rfbserver.c

@@ -362,11 +362,6 @@ rfbNewTCPOrUDPClient(rfbScreenInfoPtr rfbScreen,
       rfbScreen->clientHead = cl;
       UNLOCK(rfbClientListMutex);
 
-#ifdef LIBVNCSERVER_WITH_WEBSOCKETS
-      cl->webSockets       = FALSE;
-      cl->webSocketsBase64 = FALSE;
-#endif
-
 #if defined(LIBVNCSERVER_HAVE_LIBZ) || defined(LIBVNCSERVER_HAVE_LIBPNG)
       cl->tightQualityLevel = -1;
 #if defined(LIBVNCSERVER_HAVE_LIBJPEG) || defined(LIBVNCSERVER_HAVE_LIBPNG)
@@ -1841,66 +1836,8 @@ rfbProcessClientNormalMessage(rfbClientPtr cl)
     char encBuf2[64];
 
 #ifdef LIBVNCSERVER_WITH_WEBSOCKETS
-    if (cl->webSockets) {
-	if (cl->sslctx)
-	    n = rfbssl_peek(cl, encBuf, 4);
-	else
-	    n = recv(cl->sock, encBuf, 4, MSG_PEEK);
-
-	if (n <= 0) {
-	    if (n != 0)
-		rfbLogPerror("rfbProcessClientNormalMessage: peek");
-	    rfbCloseClient(cl);
-	    return;
-	}
-
-	if (cl->webSocketsBase64) {
-            /* With Base64 encoding we need at least 4 bytes */
-            if ((n > 0) && (n < 4)) {
-                if (encBuf[0] == '\xff') {
-		    int doclose = 0;
-                    /* Make sure we don't miss a client disconnect on an end frame
-                     * marker. Because we use a peek buffer in some cases it is not
-		     * applicable to wait for more data per select(). */
-		    switch (n) {
-			case 3:
-			    if (encBuf[1] == '\xff' && encBuf[2] == '\x00')
-				doclose = 1;
-			    break;
-			case 2:
-			    if (encBuf[1] == '\x00')
-				doclose = 1;
-			    break;
-			default:
-			    ;
-		    }
-
-		    if (cl->sslctx)
-			n = rfbssl_read(cl, encBuf, n);
-		    else
-			n = read(cl->sock, encBuf, n);
-
-		    if (doclose) {
-			rfbErr("rfbProcessClientNormalMessage: websocket close frame received\n");
-			rfbCloseClient(cl);
-		    }
-		    return;
-                }
-            }
-        } else {
-            /* With UTF-8 encoding we need at least 3 bytes (framing + 1) */
-            if ((n == 1) || (n == 2)) {
-                if (encBuf[0] == '\xff') {
-                    /* Make sure we don't miss a client disconnect on an end frame
-                    * marker */
-		    if (cl->sslctx)
-			    n = rfbssl_read(cl, encBuf, 1);
-		    else
-			    n = read(cl->sock, encBuf, 1);
-                }
-            }
-        }
-    }
+    if (cl->wsctx && webSocketCheckDisconnect(cl))
+      return;
 #endif
 
     if ((n = rfbReadExact(cl, (char *)&msg, 1)) <= 0) {

libvncserver/sockets.c

@@ -467,7 +467,7 @@ rfbReadExactTimeout(rfbClientPtr cl, char* buf, int len, int timeout)
 
     while (len > 0) {
 #ifdef LIBVNCSERVER_WITH_WEBSOCKETS
-        if (cl->webSockets) {
+        if (cl->wsctx) {
             n = webSocketsDecode(cl, buf, len);
         } else if (cl->sslctx) {
 	    n = rfbssl_read(cl, buf, len);
@@ -646,7 +646,7 @@ rfbWriteExact(rfbClientPtr cl,
 #endif
 
 #ifdef LIBVNCSERVER_WITH_WEBSOCKETS
-    if (cl->webSockets) {
+    if (cl->wsctx) {
         char *tmp = NULL;
         if ((len = webSocketsEncode(cl, buf, len, &tmp)) < 0) {
             rfbErr("WriteExact: WebSockets encode error\n");

libvncserver/websockets.c

@@ -73,6 +73,7 @@ typedef struct ws_ctx_s {
     char carryBuf[3];                      /* For base64 carry-over */
     int carrylen;
     int version;
+    int base64;
 } ws_ctx_t;
 
 typedef union ws_mask_s {
@@ -218,7 +219,7 @@ webSocketsCheck (rfbClientPtr cl)
     if (!webSocketsHandshake(cl, scheme)) {
         return FALSE;
     }
-    cl->webSockets    = TRUE;   /* Start WebSockets framing */
+    /* Start WebSockets framing */
     return TRUE;
 }
 
@@ -226,7 +227,7 @@ static rfbBool
 webSocketsHandshake(rfbClientPtr cl, char *scheme)
 {
     char *buf, *response, *line;
-    int n, linestart = 0, len = 0, llen;
+    int n, linestart = 0, len = 0, llen, base64 = 0;
     char prefix[5], trailer[17];
     char *path = NULL, *host = NULL, *origin = NULL, *protocol = NULL;
     char *key1 = NULL, *key2 = NULL, *key3 = NULL;
@@ -286,7 +287,7 @@ webSocketsHandshake(rfbClientPtr cl, char *scheme)
                 /* 16 = 4 ("GET ") + 1 ("/.*") + 11 (" HTTP/1.1\r\n") */
                 path = line+4;
                 buf[len-11] = '\0'; /* Trim trailing " HTTP/1.1\r\n" */
-                cl->webSocketsBase64 = TRUE;
+                base64 = TRUE;
                 cl->wspath = strdup(path);
                 /* rfbLog("Got path: %s\n", path); */
             } else if ((strncasecmp("host: ", line, min(llen,6))) == 0) {
@@ -381,6 +382,7 @@ webSocketsHandshake(rfbClientPtr cl, char *scheme)
     free(buf);
     cl->wsctx = (wsCtx *)calloc(1, sizeof(ws_ctx_t));
     ((ws_ctx_t *)cl->wsctx)->version = sec_ws_version ? WEBSOCKETS_VERSION_HYBI : WEBSOCKETS_VERSION_HIXIE;
+    ((ws_ctx_t *)cl->wsctx)->base64 = base64;
     return TRUE;
 }
 
@@ -438,7 +440,7 @@ webSocketsEncodeHixie(rfbClientPtr cl, const char *src, int len, char **dst)
     ws_ctx_t *wsctx = (ws_ctx_t *)cl->wsctx;
 
     wsctx->encodeBuf[sz++] = '\x00';
-    if (cl->webSocketsBase64) {
+    if (wsctx->base64) {
         len = __b64_ntop((unsigned char *)src, len, wsctx->encodeBuf+sz, sizeof(wsctx->encodeBuf) - (sz + 1));
         if (len < 0) {
             return len;
@@ -489,7 +491,10 @@ ws_peek(rfbClientPtr cl, char *buf, int len)
     if (cl->sslctx) {
 	n = rfbssl_peek(cl, buf, len);
     } else {
-	n = recv(cl->sock, buf, len, MSG_PEEK);
+	while (-1 == (n = recv(cl->sock, buf, len, MSG_PEEK))) {
+	    if (errno != EAGAIN)
+		break;
+	}
     }
     return n;
 }
@@ -507,12 +512,12 @@ webSocketsDecodeHixie(rfbClientPtr cl, char *dst, int len)
     n = ws_peek(cl, buf, len*2+2);
 
     if (n <= 0) {
-        rfbErr("%s: peek of %d\n", __func__, n);
+        rfbErr("%s: peek (%d) %m\n", __func__, errno);
         return n;
     }
 
 
-    if (cl->webSocketsBase64) {
+    if (wsctx->base64) {
         /* Base64 encoded WebSockets stream */
 
         if (buf[0] == '\xff') {
@@ -799,7 +804,7 @@ webSocketsEncodeHybi(rfbClientPtr cl, const char *src, int len, char **dst)
 
     header = (ws_header_t *)wsctx->encodeBuf;
 
-    if (cl->webSocketsBase64) {
+    if (wsctx->base64) {
 	opcode = WS_OPCODE_TEXT_FRAME;
 	/* calculate the resulting size */
 	blen = B64LEN(len);
@@ -821,7 +826,7 @@ webSocketsEncodeHybi(rfbClientPtr cl, const char *src, int len, char **dst)
       sz = 10;
     }
 
-    if (cl->webSocketsBase64) {
+    if (wsctx->base64) {
         if (-1 == (ret = __b64_ntop((unsigned char *)src, len, wsctx->encodeBuf + sz, sizeof(wsctx->encodeBuf) - sz))) {
 	  rfbErr("%s: Base 64 encode failed\n", __func__);
 	} else {
@@ -857,3 +862,64 @@ webSocketsDecode(rfbClientPtr cl, char *dst, int len)
     else
 	return webSocketsDecodeHybi(cl, dst, len);
 }
+
+/* returns TRUE if client sent an close frame or a single end of marker
+ * was received, FALSE otherwise
+ *
+ * Note: This is a Hixie-only hack!
+ **/
+rfbBool
+webSocketCheckDisconnect(rfbClientPtr cl)
+{
+    ws_ctx_t *wsctx = (ws_ctx_t *)cl->wsctx;
+    /* With Base64 encoding we need at least 4 bytes */
+    char peekbuf[4];
+    int n;
+
+    if (wsctx->version == WEBSOCKETS_VERSION_HYBI)
+	return FALSE;
+
+    if (cl->sslctx)
+	n = rfbssl_peek(cl, peekbuf, 4);
+    else
+	n = recv(cl->sock, peekbuf, 4, MSG_PEEK);
+
+    if (n <= 0) {
+	if (n != 0)
+	    rfbErr("%s: peek; %m", __func__);
+	rfbCloseClient(cl);
+	return TRUE;
+    }
+
+    if (peekbuf[0] == '\xff') {
+	int doclose = 0;
+	/* Make sure we don't miss a client disconnect on an end frame
+	 * marker. Because we use a peek buffer in some cases it is not
+	 * applicable to wait for more data per select(). */
+	switch (n) {
+	    case 3:
+		if (peekbuf[1] == '\xff' && peekbuf[2] == '\x00')
+		    doclose = 1;
+		break;
+	    case 2:
+		if (peekbuf[1] == '\x00')
+		    doclose = 1;
+		break;
+	    default:
+		;
+	}
+
+	if (cl->sslctx)
+	    n = rfbssl_read(cl, peekbuf, n);
+	else
+	    n = read(cl->sock, peekbuf, n);
+
+	if (doclose) {
+	    rfbErr("%s: websocket close frame received\n", __func__);
+	    rfbCloseClient(cl);
+	}
+	return TRUE;
+    }
+    return FALSE;
+}
+

rfb/rfb.h

@@ -640,8 +640,6 @@ typedef struct _rfbClientRec {
 #endif
 
 #ifdef LIBVNCSERVER_WITH_WEBSOCKETS
-    rfbBool webSockets;
-    rfbBool webSocketsBase64;
     rfbSslCtx *sslctx;
     wsCtx     *wsctx;
     char *wspath;                          /* Requests path component */
@@ -712,6 +710,7 @@ extern rfbBool rfbSetNonBlocking(int sock);
 /* websockets.c */
 
 extern rfbBool webSocketsCheck(rfbClientPtr cl);
+extern rfbBool webSocketCheckDisconnect(rfbClientPtr cl);
 extern int webSocketsEncode(rfbClientPtr cl, const char *src, int len, char **dst);
 extern int webSocketsDecode(rfbClientPtr cl, char *dst, int len);
 #endif