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.

4 files changed, 80 insertions, 78 deletions

diff --git a/libvncserver/rfbserver.c b/libvncserver/rfbserver.c
index d6a5da4..63f21db 100644
--- a/libvncserver/rfbserver.c
+++ b/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) {
diff --git a/libvncserver/sockets.c b/libvncserver/sockets.c
index b3d5b59..415f712 100644
--- a/libvncserver/sockets.c
+++ b/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");
diff --git a/libvncserver/websockets.c b/libvncserver/websockets.c
index e3b47e3..88b76a5 100755
--- a/libvncserver/websockets.c
+++ b/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;
+}
+
diff --git a/rfb/rfb.h b/rfb/rfb.h
index 1f29e63..11d1447 100644
--- a/rfb/rfb.h
+++ b/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