From fa4086736a8b9de4d5d3260b58e46dbb49429aba Mon Sep 17 00:00:00 2001
From: Nicolas Morais <nmorais.st@gmail.com>
Date: Sun, 12 May 2024 10:39:21 -0300
Subject: [PATCH] Adds lastError to rfbClientRec, adds 'error changed' callback

New rfbClientSetErr() function is only being called in auth.c, when simple password authentication fails.
---
 include/rfb/rfb.h            | 14 ++++++++++++++
 src/libvncserver/auth.c      |  1 +
 src/libvncserver/main.c      | 28 ++++++++++++++++++++++++++++
 src/libvncserver/rfbserver.c |  4 +++-
 4 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/include/rfb/rfb.h b/include/rfb/rfb.h
index 79a446f1b..bdd403411 100644
--- a/include/rfb/rfb.h
+++ b/include/rfb/rfb.h
@@ -138,6 +138,9 @@ typedef int  (*rfbFileTransferPermitted) (struct _rfbClientRec* cl);
 /** Handle the textchat messages */
 typedef void (*rfbSetTextChat) (struct _rfbClientRec* cl, int length, char *string);
 
+/* error handling (server side) */
+typedef void (*ClientErrorChangedProc)(struct _rfbClientRec *client);
+
 typedef struct {
   uint32_t count;
   rfbBool is16; /**< is the data format short? */
@@ -302,6 +305,8 @@ typedef struct _rfbScreenInfo
     rfbFileTransferPermitted  getFileTransferPermission;
     rfbSetTextChat            setTextChat;
 
+    ClientErrorChangedProc clientErrorChanged;
+
     /** newClientHook is called just after a new client is created */
     rfbNewClientHookPtr newClientHook;
     /** displayHook is called just before a frame buffer update */
@@ -440,6 +445,13 @@ typedef struct _rfbClientRec {
      * This is useful if the IO functions have to behave client specific.
      */
     void* clientData;
+    
+    /**
+     * error handling
+    */
+
+    char* lastError;
+
     ClientGoneHookPtr clientGoneHook;
 
     rfbSocket sock;
@@ -1017,6 +1029,8 @@ extern rfbBool rfbProcessSizeArguments(int* width,int* height,int* bpp,int* argc
 
 extern void rfbLogEnable(int enabled);
 typedef void (*rfbLogProc)(const char *format, ...);
+typedef void (*rfbClientSetErrProc)(rfbClientPtr cl, const char *format, ...);
+extern rfbClientSetErrProc rfbClientSetErr;
 extern rfbLogProc rfbLog, rfbErr;
 extern void rfbLogPerror(const char *str);
 
diff --git a/src/libvncserver/auth.c b/src/libvncserver/auth.c
index 748027b1f..5731efc65 100644
--- a/src/libvncserver/auth.c
+++ b/src/libvncserver/auth.c
@@ -382,6 +382,7 @@ rfbAuthProcessClientMessage(rfbClientPtr cl)
 
     if(!cl->screen->passwordCheck(cl,(const char*)response,CHALLENGESIZE)) {
         rfbErr("rfbAuthProcessClientMessage: password check failed\n");
+        rfbClientSetErr(cl, "rfbAuthProcessClientMessage: password check failed");
         authResult = Swap32IfLE(rfbVncAuthFailed);
         if (rfbWriteExact(cl, (char *)&authResult, 4) < 0) {
             rfbLogPerror("rfbAuthProcessClientMessage: write");
diff --git a/src/libvncserver/main.c b/src/libvncserver/main.c
index 1efa83879..577ab33b2 100644
--- a/src/libvncserver/main.c
+++ b/src/libvncserver/main.c
@@ -260,8 +260,36 @@ rfbDefaultLog(const char *format, ...)
     UNLOCK(logMutex);
 }
 
+static void rfbDefaultClientSetErr(rfbClientPtr cl, const char *format, ...) {
+    // props to josemr goudetalvim (@josealvim) for this logic
+    va_list ap;
+    va_start(ap, format);
+    size_t len = (size_t)vsnprintf(NULL, 0, format, ap);
+    va_end(ap);
+
+    char *buffer = malloc(len + 1);
+    if (!buffer)
+        return;
+
+    va_start(ap, format);
+    vsprintf(buffer, format, ap);
+
+    va_end(ap);
+
+    char *oldError = cl->lastError;
+    cl->lastError = buffer;
+
+    // call the callback if it exists
+    if (cl->screen->clientErrorChanged != NULL)
+        cl->screen->clientErrorChanged(cl);
+
+    if (oldError)
+        free(oldError);
+}
+
 rfbLogProc rfbLog=rfbDefaultLog;
 rfbLogProc rfbErr=rfbDefaultLog;
+rfbClientSetErrProc rfbClientSetErr = rfbDefaultClientSetErr;
 
 void rfbLogPerror(const char *str)
 {
diff --git a/src/libvncserver/rfbserver.c b/src/libvncserver/rfbserver.c
index a80fd1aee..eefd1280e 100644
--- a/src/libvncserver/rfbserver.c
+++ b/src/libvncserver/rfbserver.c
@@ -715,7 +715,9 @@ rfbProcessClientProtocolVersion(rfbClientPtr cl)
     pv[sz_rfbProtocolVersionMsg] = 0;
     if (sscanf(pv,rfbProtocolVersionFormat,&major_,&minor_) != 2) {
 	rfbErr("rfbProcessClientProtocolVersion: not a valid RFB client: %s\n", pv);
-	rfbCloseClient(cl);
+        rfbClientSetErr(
+            cl, "rfbProcessClientProtocolVersion: not a valid RFB client: %s\n", pv);
+        rfbCloseClient(cl);
 	return;
     }
     rfbLog("Client Protocol Version %d.%d\n", major_, minor_);