From 0ee9b2a98edbe9eedf951e645a5bdaaf2206d94d Mon Sep 17 00:00:00 2001
From: Tobias Junghans <tobydox@veyon.io>
Date: Mon, 19 Jul 2021 08:49:25 +0200
Subject: [PATCH] libvncclient: add hooks for custom socket I/O

This allows using libvncclient on any kind of custom transport, e.g.
for TLS tunneling via a special TLS socket implementation.
---
 include/rfb/rfbclient.h      | 10 ++++++++++
 src/libvncclient/rfbclient.c |  5 +++++
 src/libvncclient/sockets.c   | 14 ++++++++++++--
 src/libvncclient/vncviewer.c |  4 ++++
 4 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/include/rfb/rfbclient.h b/include/rfb/rfbclient.h
index 07ef0cf52..5a6fb7dcd 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/src/libvncclient/rfbclient.c b/src/libvncclient/rfbclient.c
index 94a371491..37bc30c67 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. */
diff --git a/src/libvncclient/sockets.c b/src/libvncclient/sockets.c
index 0318391bd..017aa9229 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
@@ -262,6 +266,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);
diff --git a/src/libvncclient/vncviewer.c b/src/libvncclient/vncviewer.c
index 4a0d7ebfd..8ab7e29dc 100644
--- a/src/libvncclient/vncviewer.c
+++ b/src/libvncclient/vncviewer.c
@@ -552,7 +552,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);
   if (client->listen6Sock != RFB_INVALID_SOCKET)