Send lock LED state from server to client

7 files changed, 161 insertions, 54 deletions

diff --git a/common/rfb/SMsgHandler.cxx b/common/rfb/SMsgHandler.cxx
index 388b21f8..8e48c673 100644
--- a/common/rfb/SMsgHandler.cxx
+++ b/common/rfb/SMsgHandler.cxx
@@ -41,10 +41,11 @@ void SMsgHandler::setPixelFormat(const PixelFormat& pf)
 
 void SMsgHandler::setEncodings(int nEncodings, const rdr::S32* encodings)
 {
-  bool firstFence, firstContinuousUpdates;
+  bool firstFence, firstContinuousUpdates, firstLEDState;
 
   firstFence = !cp.supportsFence;
   firstContinuousUpdates = !cp.supportsContinuousUpdates;
+  firstLEDState = !cp.supportsLEDState;
 
   cp.setEncodings(nEncodings, encodings);
 
@@ -54,6 +55,8 @@ void SMsgHandler::setEncodings(int nEncodings, const rdr::S32* encodings)
     supportsFence();
   if (cp.supportsContinuousUpdates && firstContinuousUpdates)
     supportsContinuousUpdates();
+  if (cp.supportsLEDState && firstLEDState)
+    supportsLEDState();
 }
 
 void SMsgHandler::supportsLocalCursor()
@@ -68,6 +71,10 @@ void SMsgHandler::supportsContinuousUpdates()
 {
 }
 
+void SMsgHandler::supportsLEDState()
+{
+}
+
 void SMsgHandler::setDesktopSize(int fb_width, int fb_height,
                                  const ScreenSet& layout)
 {
diff --git a/common/rfb/SMsgHandler.h b/common/rfb/SMsgHandler.h
index 74509e0b..cf6b6b3b 100644
--- a/common/rfb/SMsgHandler.h
+++ b/common/rfb/SMsgHandler.h
@@ -74,6 +74,12 @@ namespace rfb {
     // this point if it is supported.
     virtual void supportsContinuousUpdates();
 
+    // supportsLEDState() is called the first time we detect that the
+    // client supports the LED state extension. A LEDState message
+    // should be sent back to the client to inform it of the current
+    // server state.
+    virtual void supportsLEDState();
+
     ConnParams cp;
   };
 }
diff --git a/common/rfb/SMsgWriter.cxx b/common/rfb/SMsgWriter.cxx
index bc3f4398..d8adfbc1 100644
--- a/common/rfb/SMsgWriter.cxx
+++ b/common/rfb/SMsgWriter.cxx
@@ -27,6 +27,7 @@
 #include <rfb/Encoder.h>
 #include <rfb/SMsgWriter.h>
 #include <rfb/LogWriter.h>
+#include <rfb/ledStates.h>
 
 using namespace rfb;
 
@@ -37,7 +38,8 @@ SMsgWriter::SMsgWriter(ConnParams* cp_, rdr::OutStream* os_)
     nRectsInUpdate(0), nRectsInHeader(0),
     needSetDesktopSize(false), needExtendedDesktopSize(false),
     needSetDesktopName(false), needSetCursor(false),
-    needSetXCursor(false), needSetCursorWithAlpha(false)
+    needSetXCursor(false), needSetCursorWithAlpha(false),
+    needLEDState(false)
 {
 }
 
@@ -193,12 +195,26 @@ bool SMsgWriter::writeSetCursorWithAlpha()
   return true;
 }
 
+bool SMsgWriter::writeLEDState()
+{
+  if (!cp->supportsLEDState)
+    return false;
+  if (cp->ledState() == ledUnknown)
+    return false;
+
+  needLEDState = true;
+
+  return true;
+}
+
 bool SMsgWriter::needFakeUpdate()
 {
   if (needSetDesktopName)
     return true;
   if (needSetCursor || needSetXCursor || needSetCursorWithAlpha)
     return true;
+  if (needLEDState)
+    return true;
   if (needNoDataUpdate())
     return true;
 
@@ -247,6 +263,8 @@ void SMsgWriter::writeFramebufferUpdateStart(int nRects)
       nRects++;
     if (needSetCursorWithAlpha)
       nRects++;
+    if (needLEDState)
+      nRects++;
   }
 
   os->writeU16(nRects);
@@ -362,6 +380,11 @@ void SMsgWriter::writePseudoRects()
     writeSetDesktopNameRect(cp->name());
     needSetDesktopName = false;
   }
+
+  if (needLEDState) {
+    writeLEDStateRect(cp->ledState());
+    needLEDState = false;
+  }
 }
 
 void SMsgWriter::writeNoDataRects()
@@ -525,3 +548,20 @@ void SMsgWriter::writeSetCursorWithAlphaRect(int width, int height,
     data += 4;
   }
 }
+
+void SMsgWriter::writeLEDStateRect(rdr::U8 state)
+{
+  if (!cp->supportsLEDState)
+    throw Exception("Client does not support LED state updates");
+  if (cp->ledState() == ledUnknown)
+    throw Exception("Server does not support LED state updates");
+  if (++nRectsInUpdate > nRectsInHeader && nRectsInHeader)
+    throw Exception("SMsgWriter::writeLEDStateRect: nRects out of sync");
+
+  os->writeS16(0);
+  os->writeS16(0);
+  os->writeU16(0);
+  os->writeU16(0);
+  os->writeU32(pseudoEncodingLEDState);
+  os->writeU8(state);
+}
diff --git a/common/rfb/SMsgWriter.h b/common/rfb/SMsgWriter.h
index 548b8e8e..890b2b5b 100644
--- a/common/rfb/SMsgWriter.h
+++ b/common/rfb/SMsgWriter.h
@@ -82,6 +82,9 @@ namespace rfb {
     bool writeSetXCursor();
     bool writeSetCursorWithAlpha();
 
+    // Same for LED state message
+    bool writeLEDState();
+
     // needFakeUpdate() returns true when an immediate update is needed in
     // order to flush out pseudo-rectangles to the client.
     bool needFakeUpdate();
@@ -131,6 +134,7 @@ namespace rfb {
     void writeSetCursorWithAlphaRect(int width, int height,
                                      int hotspotX, int hotspotY,
                                      const rdr::U8* data);
+    void writeLEDStateRect(rdr::U8 state);
 
     ConnParams* cp;
     rdr::OutStream* os;
@@ -145,6 +149,7 @@ namespace rfb {
     bool needSetCursor;
     bool needSetXCursor;
     bool needSetCursorWithAlpha;
+    bool needLEDState;
 
     typedef struct {
       rdr::U16 reason, result;
diff --git a/common/rfb/VNCSConnectionST.cxx b/common/rfb/VNCSConnectionST.cxx
index de41e8fc..232776fc 100644
--- a/common/rfb/VNCSConnectionST.cxx
+++ b/common/rfb/VNCSConnectionST.cxx
@@ -313,6 +313,16 @@ void VNCSConnectionST::setCursorOrClose()
 }
 
 
+void VNCSConnectionST::setLEDStateOrClose(unsigned int state)
+{
+  try {
+    setLEDState(state);
+  } catch(rdr::Exception& e) {
+    close(e.str());
+  }
+}
+
+
 int VNCSConnectionST::checkIdleTimeout()
 {
   int idleTimeout = rfb::Server::idleTimeout;
@@ -418,6 +428,7 @@ void VNCSConnectionST::authSuccess()
   cp.height = server->pb->height();
   cp.screenLayout = server->screenLayout;
   cp.setName(server->getName());
+  cp.setLEDState(server->ledState);
   
   // - Set the default pixel format
   cp.setPF(server->pb->getPF());
@@ -570,61 +581,66 @@ void VNCSConnectionST::keyEvent(rdr::U32 key, bool down) {
     vlog.debug("Ignoring lock key (e.g. caps lock)");
     return;
   }
-  // Always ignore ScrollLock though as we don't have a heuristic
-  // for that
-  if (key == XK_Scroll_Lock) {
-    vlog.debug("Ignoring lock key (e.g. caps lock)");
-    return;
-  }
 
-  if (down && (server->ledState != ledUnknown)) {
-    // CapsLock synchronisation heuristic
-    // (this assumes standard interaction between CapsLock the Shift
-    // keys and normal characters)
-    if (((key >= XK_A) && (key <= XK_Z)) ||
-        ((key >= XK_a) && (key <= XK_z))) {
-      bool uppercase, shift, lock;
-
-      uppercase = (key >= XK_A) && (key <= XK_Z);
-      shift = pressedKeys.find(XK_Shift_L) != pressedKeys.end() ||
-              pressedKeys.find(XK_Shift_R) != pressedKeys.end();
-      lock = server->ledState & ledCapsLock;
-
-      if (lock == (uppercase == shift)) {
-        vlog.debug("Inserting fake CapsLock to get in sync with client");
-        server->desktop->keyEvent(XK_Caps_Lock, true);
-        server->desktop->keyEvent(XK_Caps_Lock, false);
-      }
+  // Lock key heuristics
+  // (only for clients that do not support the LED state extension)
+  if (!cp.supportsLEDState) {
+    // Always ignore ScrollLock as we don't have a heuristic
+    // for that
+    if (key == XK_Scroll_Lock) {
+      vlog.debug("Ignoring lock key (e.g. caps lock)");
+      return;
     }
 
-    // NumLock synchronisation heuristic
-    // (this is more cautious because of the differences between Unix,
-    // Windows and macOS)
-    if (((key >= XK_KP_Home) && (key <= XK_KP_Delete)) ||
-        ((key >= XK_KP_0) && (key <= XK_KP_9)) ||
-        (key == XK_KP_Separator) || (key == XK_KP_Decimal)) {
-      bool number, shift, lock;
-
-      number = ((key >= XK_KP_0) && (key <= XK_KP_9)) ||
-                (key == XK_KP_Separator) || (key == XK_KP_Decimal);
-      shift = pressedKeys.find(XK_Shift_L) != pressedKeys.end() ||
-              pressedKeys.find(XK_Shift_R) != pressedKeys.end();
-      lock = server->ledState & ledNumLock;
-
-      if (shift) {
-        // We don't know the appropriate NumLock state for when Shift
-        // is pressed as it could be one of:
-        //
-        // a) A Unix client where Shift negates NumLock
-        //
-        // b) A Windows client where Shift only cancels NumLock
-        //
-        // c) A macOS client where Shift doesn't have any effect
-        //
-      } else if (lock == (number == shift)) {
-        vlog.debug("Inserting fake NumLock to get in sync with client");
-        server->desktop->keyEvent(XK_Num_Lock, true);
-        server->desktop->keyEvent(XK_Num_Lock, false);
+    if (down && (server->ledState != ledUnknown)) {
+      // CapsLock synchronisation heuristic
+      // (this assumes standard interaction between CapsLock the Shift
+      // keys and normal characters)
+      if (((key >= XK_A) && (key <= XK_Z)) ||
+          ((key >= XK_a) && (key <= XK_z))) {
+        bool uppercase, shift, lock;
+
+        uppercase = (key >= XK_A) && (key <= XK_Z);
+        shift = pressedKeys.find(XK_Shift_L) != pressedKeys.end() ||
+                pressedKeys.find(XK_Shift_R) != pressedKeys.end();
+        lock = server->ledState & ledCapsLock;
+
+        if (lock == (uppercase == shift)) {
+          vlog.debug("Inserting fake CapsLock to get in sync with client");
+          server->desktop->keyEvent(XK_Caps_Lock, true);
+          server->desktop->keyEvent(XK_Caps_Lock, false);
+        }
+      }
+
+      // NumLock synchronisation heuristic
+      // (this is more cautious because of the differences between Unix,
+      // Windows and macOS)
+      if (((key >= XK_KP_Home) && (key <= XK_KP_Delete)) ||
+          ((key >= XK_KP_0) && (key <= XK_KP_9)) ||
+          (key == XK_KP_Separator) || (key == XK_KP_Decimal)) {
+        bool number, shift, lock;
+
+        number = ((key >= XK_KP_0) && (key <= XK_KP_9)) ||
+                  (key == XK_KP_Separator) || (key == XK_KP_Decimal);
+        shift = pressedKeys.find(XK_Shift_L) != pressedKeys.end() ||
+                pressedKeys.find(XK_Shift_R) != pressedKeys.end();
+        lock = server->ledState & ledNumLock;
+
+        if (shift) {
+          // We don't know the appropriate NumLock state for when Shift
+          // is pressed as it could be one of:
+          //
+          // a) A Unix client where Shift negates NumLock
+          //
+          // b) A Windows client where Shift only cancels NumLock
+          //
+          // c) A macOS client where Shift doesn't have any effect
+          //
+        } else if (lock == (number == shift)) {
+          vlog.debug("Inserting fake NumLock to get in sync with client");
+          server->desktop->keyEvent(XK_Num_Lock, true);
+          server->desktop->keyEvent(XK_Num_Lock, false);
+        }
       }
     }
   }
@@ -818,6 +834,11 @@ void VNCSConnectionST::supportsContinuousUpdates()
   writer()->writeEndOfContinuousUpdates();
 }
 
+void VNCSConnectionST::supportsLEDState()
+{
+  writer()->writeLEDState();
+}
+
 
 bool VNCSConnectionST::handleTimeout(Timer* t)
 {
@@ -1230,6 +1251,21 @@ void VNCSConnectionST::setDesktopName(const char *name)
   writeFramebufferUpdate();
 }
 
+void VNCSConnectionST::setLEDState(unsigned int ledstate)
+{
+  if (state() != RFBSTATE_NORMAL)
+    return;
+
+  cp.setLEDState(ledstate);
+
+  if (!writer()->writeLEDState()) {
+    // No client support
+    return;
+  }
+
+  writeFramebufferUpdate();
+}
+
 void VNCSConnectionST::setSocketTimeouts()
 {
   int timeoutms = rfb::Server::clientWaitTimeMillis;
diff --git a/common/rfb/VNCSConnectionST.h b/common/rfb/VNCSConnectionST.h
index 74a6946d..8f33962d 100644
--- a/common/rfb/VNCSConnectionST.h
+++ b/common/rfb/VNCSConnectionST.h
@@ -78,6 +78,7 @@ namespace rfb {
     void bellOrClose();
     void serverCutTextOrClose(const char *str, int len);
     void setDesktopNameOrClose(const char *name);
+    void setLEDStateOrClose(unsigned int state);
 
     // checkIdleTimeout() returns the number of milliseconds left until the
     // idle timeout expires.  If it has expired, the connection is closed and
@@ -146,6 +147,7 @@ namespace rfb {
     virtual void supportsLocalCursor();
     virtual void supportsFence();
     virtual void supportsContinuousUpdates();
+    virtual void supportsLEDState();
 
     // setAccessRights() allows a security package to limit the access rights
     // of a VNCSConnectioST to the server.  These access rights are applied
@@ -174,6 +176,7 @@ namespace rfb {
     void screenLayoutChange(rdr::U16 reason);
     void setCursor();
     void setDesktopName(const char *name);
+    void setLEDState(unsigned int state);
     void setSocketTimeouts();
 
     network::Socket* sock;
diff --git a/common/rfb/VNCServerST.cxx b/common/rfb/VNCServerST.cxx
index 28f6a62b..43e8f3ed 100644
--- a/common/rfb/VNCServerST.cxx
+++ b/common/rfb/VNCServerST.cxx
@@ -461,7 +461,17 @@ void VNCServerST::setCursorPos(const Point& pos)
 
 void VNCServerST::setLEDState(unsigned int state)
 {
+  std::list<VNCSConnectionST*>::iterator ci, ci_next;
+
+  if (state == ledState)
+    return;
+
   ledState = state;
+
+  for (ci = clients.begin(); ci != clients.end(); ci = ci_next) {
+    ci_next = ci; ci_next++;
+    (*ci)->setLEDStateOrClose(state);
+  }
 }
 
 // Other public methods