libvncclient: add hooks for custom socket I/O

CMakeLists.txt

@@ -183,6 +183,7 @@ check_include_file("dirent.h"      LIBVNCSERVER_HAVE_DIRENT_H)
 check_include_file("endian.h"      LIBVNCSERVER_HAVE_ENDIAN_H)
 check_include_file("fcntl.h"       LIBVNCSERVER_HAVE_FCNTL_H)
 check_include_file("netinet/in.h"  LIBVNCSERVER_HAVE_NETINET_IN_H)
+check_include_file("poll.h"        LIBVNCSERVER_HAVE_POLL_H)
 check_include_file("sys/endian.h"  LIBVNCSERVER_HAVE_SYS_ENDIAN_H)
 check_include_file("sys/socket.h"  LIBVNCSERVER_HAVE_SYS_SOCKET_H)
 check_include_file("sys/stat.h"    LIBVNCSERVER_HAVE_SYS_STAT_H)
@@ -218,6 +219,7 @@ check_function_exists(inet_ntoa       LIBVNCSERVER_HAVE_INET_NTOA)
 check_function_exists(memmove         LIBVNCSERVER_HAVE_MEMMOVE)
 check_function_exists(memset          LIBVNCSERVER_HAVE_MEMSET)
 check_function_exists(mkfifo          LIBVNCSERVER_HAVE_MKFIFO)
+check_function_exists(poll            LIBVNCSERVER_HAVE_POLL)
 check_function_exists(select          LIBVNCSERVER_HAVE_SELECT)
 check_function_exists(socket          LIBVNCSERVER_HAVE_SOCKET)
 check_function_exists(strchr          LIBVNCSERVER_HAVE_STRCHR)

include/rfb/rfbclient.h

@@ -233,6 +233,10 @@ typedef void (*GotBitmapProc)(struct _rfbClient* client, const uint8_t* buffer,
 typedef rfbBool (*GotJpegProc)(struct _rfbClient* client, const uint8_t* buffer, int length, int x, int y, int w, int h);
 typedef rfbBool (*LockWriteToTLSProc)(struct _rfbClient* client);   /** @deprecated */
 typedef rfbBool (*UnlockWriteToTLSProc)(struct _rfbClient* client); /** @deprecated */
+typedef rfbSocket (*ConnectToRFBServerProc)(struct _rfbClient* client, const char* hostname, int port);
+typedef int (*ReadFromSocketProc)(struct _rfbClient* client, char* buf, unsigned int len);
+typedef int (*WriteToSocketProc)(struct _rfbClient* client, const char* buf, unsigned int len);
+typedef void (*CloseSocketProc)(struct _rfbClient* client);
 
 #ifdef LIBVNCSERVER_HAVE_SASL
 typedef char* (*GetUserProc)(struct _rfbClient* client);
@@ -467,6 +471,12 @@ typedef struct _rfbClient {
 	 * ReadFromRFBServer() - keep at 0 to disable timeout detection and handling */
 	unsigned int readTimeout;
 
+	/** hooks for custom socket I/O */
+	ConnectToRFBServerProc ConnectToRFBServer;
+	ReadFromSocketProc ReadFromSocket;
+	WriteToSocketProc WriteToSocket;
+	CloseSocketProc CloseSocket;
+
 	/**
 	 * Mutex to protect concurrent TLS read/write.
 	 * For internal use only.

include/rfb/rfbconfig.h.cmakein

@@ -42,6 +42,9 @@
 /* Define to 1 if you have the `mkfifo' function. */
 #cmakedefine LIBVNCSERVER_HAVE_MKFIFO  1 
 
+/* Define to 1 if you have the `poll' function. */
+#cmakedefine LIBVNCSERVER_HAVE_POLL 1
+
 /* Define to 1 if you have the `select' function. */
 #cmakedefine LIBVNCSERVER_HAVE_SELECT  1 
 
@@ -84,6 +87,9 @@
 /* Define to 1 if you have the <netinet/in.h> header file. */
 #cmakedefine LIBVNCSERVER_HAVE_NETINET_IN_H  1 
 
+/* Define to 1 if you have the <poll.h> header file. */
+#cmakedefine LIBVNCSERVER_HAVE_POLL_H  1
+
 /* Define to 1 if you have the <sys/endian.h> header file. */
 #cmakedefine LIBVNCSERVER_HAVE_SYS_ENDIAN_H 1
 

src/common/sockets.c

@@ -53,6 +53,15 @@ rfbBool sock_set_nonblocking(rfbSocket sock, rfbBool non_blocking, void (*log)(c
 
 rfbBool sock_wait_for_connected(int socket, unsigned int timeout_seconds)
 {
+#ifdef LIBVNCSERVER_HAVE_POLL
+  struct pollfd pfd;
+  pfd.fd = socket;
+  pfd.events = POLLIN | POLLPRI;
+
+  if (poll(&pfd, 1, timeout_seconds*1000)==1) {
+    if ((pfd.revents & POLLERR) || (pfd.revents & POLLHUP))
+      return FALSE;
+#else
   fd_set writefds;
   fd_set exceptfds;
   struct timeval timeout;
@@ -68,7 +77,9 @@ rfbBool sock_wait_for_connected(int socket, unsigned int timeout_seconds)
 #ifdef WIN32
     if (FD_ISSET(socket, &exceptfds))
       return FALSE;
-#else
+#endif
+#endif
+#ifndef WIN32
     int so_error;
     socklen_t len = sizeof so_error;
     getsockopt(socket, SOL_SOCKET, SO_ERROR, &so_error, &len);

src/common/sockets.h

@@ -64,6 +64,11 @@
 #include <netdb.h>
 #endif
 
+#include "rfb/rfbconfig.h"
+#if LIBVNCSERVER_HAVE_POLL_H
+#include <poll.h>
+#endif
+
 /*
   Common internal socket functions
  */

src/libvncclient/rfbclient.c

@@ -324,6 +324,11 @@ ConnectToRFBServer(rfbClient* client,const char *hostname, int port)
     return TRUE;
   }
 
+  if(client->ConnectToRFBServer)
+  {
+      client->sock = client->ConnectToRFBServer(client, hostname, port);
+  }
+  else
 #ifndef WIN32
   if(IsUnixSocket(hostname))
     /* serverHost is a UNIX socket. */
@@ -1424,8 +1429,6 @@ SetFormatAndEncodings(rfbClient* client)
   /* New Frame Buffer Size */
   if (se->nEncodings < MAX_ENCODINGS && client->canHandleNewFBSize)
     encs[se->nEncodings++] = rfbClientSwap32IfLE(rfbEncodingNewFBSize);
-  if (se->nEncodings < MAX_ENCODINGS)
-    encs[se->nEncodings++] = rfbClientSwap32IfLE(rfbEncodingExtDesktopSize);
 
   /* Last Rect */
   if (se->nEncodings < MAX_ENCODINGS && requestLastRectEncoding)

src/libvncclient/sockets.c

@@ -134,7 +134,9 @@ ReadFromRFBServer(rfbClient* client, char *out, unsigned int n)
 
     while (client->buffered < n) {
       int i;
-      if (client->tlsSession)
+      if (client->ReadFromSocket)
+          i = client->ReadFromSocket(client, client->buf + client->buffered, RFB_BUF_SIZE - client->buffered);
+      else if (client->tlsSession)
         i = ReadFromTLS(client, client->buf + client->buffered, RFB_BUF_SIZE - client->buffered);
       else
 #ifdef LIBVNCSERVER_HAVE_SASL
@@ -186,7 +188,9 @@ ReadFromRFBServer(rfbClient* client, char *out, unsigned int n)
 
     while (n > 0) {
       int i;
-      if (client->tlsSession)
+      if (client->ReadFromSocket)
+          i = client->ReadFromSocket(client, out, n);
+      else if (client->tlsSession)
         i = ReadFromTLS(client, out, n);
       else
 #ifdef LIBVNCSERVER_HAVE_SASL
@@ -249,7 +253,11 @@ ReadFromRFBServer(rfbClient* client, char *out, unsigned int n)
 rfbBool
 WriteToRFBServer(rfbClient* client, const char *buf, unsigned int n)
 {
+#ifdef LIBVNCSERVER_HAVE_POLL
+  struct pollfd pfd;
+#else
   fd_set fds;
+#endif
   int i = 0;
   int j;
   const char *obuf = buf;
@@ -262,6 +270,12 @@ WriteToRFBServer(rfbClient* client, const char *buf, unsigned int n)
   if (client->serverPort==-1)
     return TRUE; /* vncrec playing */
 
+  if (client->WriteToSocket) {
+      i = client->WriteToSocket(client, buf, n);
+      if (i <= 0) return FALSE;
+
+      return TRUE;
+  }
   if (client->tlsSession) {
     /* WriteToTLS() will guarantee either everything is written, or error/eof returns */
     i = WriteToTLS(client, buf, n);
@@ -296,13 +310,26 @@ WriteToRFBServer(rfbClient* client, const char *buf, unsigned int n)
 		errno == ENOENT ||
 #endif
 		errno == EAGAIN) {
+#ifdef LIBVNCSERVER_HAVE_POLL
+	  struct pollfd pfd;
+	  pfd.fd = client->sock;
+	  pfd.events = POLLOUT;
+
+	  if (poll(&pfd, 1, -1) <= 0) {
+	    if ((pfd.revents & POLLERR) || (pfd.revents & POLLHUP)) {
+	      rfbClientErr("poll\n");
+	      return FALSE;
+	    }
+	  }
+#else
 	  FD_ZERO(&fds);
 	  FD_SET(client->sock,&fds);
 
 	  if (select(client->sock+1, NULL, &fds, NULL, NULL) <= 0) {
 	    rfbClientErr("select\n");
 	    return FALSE;
 	  }
+#endif
 	  j = 0;
 	} else {
 	  rfbClientErr("write\n");
@@ -847,21 +874,33 @@ PrintInHex(char *buf, int len)
 
 int WaitForMessage(rfbClient* client,unsigned int usecs)
 {
+#ifdef LIBVNCSERVER_HAVE_POLL
+  struct pollfd pfd;
+#else
   fd_set fds;
   struct timeval timeout;
+#endif
   int num;
 
   if (client->serverPort==-1)
     /* playing back vncrec file */
     return 1;
 
+#ifdef LIBVNCSERVER_HAVE_POLL
+  pfd.fd = client->sock;
+  pfd.events = POLLIN | POLLPRI;
+  num = poll(&pfd, 1, usecs/1000);
+  if ((pfd.revents & POLLERR) || (pfd.revents & POLLHUP))
+    return -1;
+#else
   timeout.tv_sec=(usecs/1000000);
   timeout.tv_usec=(usecs%1000000);
 
   FD_ZERO(&fds);
   FD_SET(client->sock,&fds);
 
   num=select(client->sock+1, &fds, NULL, NULL, &timeout);
+#endif
   if(num<0) {
 #ifdef WIN32
     errno=WSAGetLastError();

src/libvncclient/vncviewer.c

@@ -554,7 +554,11 @@ void rfbClientCleanup(rfbClient* client) {
 	  free(client->vncRec);
 
   if (client->sock != RFB_INVALID_SOCKET)
+  {
+    if (client->CloseSocket)
+      client->CloseSocket(client);
     rfbCloseSocket(client->sock);
+  }
   if (client->listenSock != RFB_INVALID_SOCKET)
     rfbCloseSocket(client->listenSock);
   free(client->desktopName);