diff --git a/CMakeLists.txt b/CMakeLists.txt index 0e391f422..00342b9cf 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/include/rfb/rfbclient.h b/include/rfb/rfbclient.h index 97c893824..95e5b2463 100644 --- a/include/rfb/rfbclient.h +++ b/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. diff --git a/include/rfb/rfbconfig.h.cmakein b/include/rfb/rfbconfig.h.cmakein index d50c3c965..1271282ac 100644 --- a/include/rfb/rfbconfig.h.cmakein +++ b/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 header file. */ #cmakedefine LIBVNCSERVER_HAVE_NETINET_IN_H 1 +/* Define to 1 if you have the header file. */ +#cmakedefine LIBVNCSERVER_HAVE_POLL_H 1 + /* Define to 1 if you have the header file. */ #cmakedefine LIBVNCSERVER_HAVE_SYS_ENDIAN_H 1 diff --git a/src/common/sockets.c b/src/common/sockets.c index de66fa768..9ffa47b96 100644 --- a/src/common/sockets.c +++ b/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); diff --git a/src/common/sockets.h b/src/common/sockets.h index 99a2f15e8..b7f41ca2e 100644 --- a/src/common/sockets.h +++ b/src/common/sockets.h @@ -64,6 +64,11 @@ #include #endif +#include "rfb/rfbconfig.h" +#if LIBVNCSERVER_HAVE_POLL_H +#include +#endif + /* Common internal socket functions */ diff --git a/src/libvncclient/rfbclient.c b/src/libvncclient/rfbclient.c index 3d488ba9d..f9f8ccc58 100644 --- a/src/libvncclient/rfbclient.c +++ b/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) diff --git a/src/libvncclient/sockets.c b/src/libvncclient/sockets.c index 1bd41818f..944c092f4 100644 --- a/src/libvncclient/sockets.c +++ b/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,6 +310,18 @@ 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); @@ -303,6 +329,7 @@ WriteToRFBServer(rfbClient* client, const char *buf, unsigned int n) rfbClientErr("select\n"); return FALSE; } +#endif j = 0; } else { rfbClientErr("write\n"); @@ -847,14 +874,25 @@ 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); @@ -862,6 +900,7 @@ int WaitForMessage(rfbClient* client,unsigned int usecs) FD_SET(client->sock,&fds); num=select(client->sock+1, &fds, NULL, NULL, &timeout); +#endif if(num<0) { #ifdef WIN32 errno=WSAGetLastError(); diff --git a/src/libvncclient/vncviewer.c b/src/libvncclient/vncviewer.c index 9df57d99d..2e31a8b58 100644 --- a/src/libvncclient/vncviewer.c +++ b/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);