Merge pull request #1426 from rdkcentral/development/websocket-masking

WebSocket: add masking bytes also while enabling MASKING_FRAME flag
rdkcentral · Oct 11, 2023 · 51da54b · 51da54b
2 parents c1fe3a3 + dc9ae03
commit 51da54b
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/Source/websocket/WebSocketLink.cpp b/Source/websocket/WebSocketLink.cpp
@@ -91,14 +91,8 @@ namespace Web {
                         ::memmove(&dataFrame[2], &(dataFrame[4]), usedSize);
                     }
                 } else {
-                    uint32_t value;
-                    // Generate a new mask value
                     uint8_t maskKey[4];
-                    Crypto::Random(value);
-                    maskKey[0] = value & 0xFF;
-                    maskKey[1] = (value >> 8) & 0xFF;
-                    maskKey[2] = (value >> 16) & 0xFF;
-                    maskKey[3] = (value >> 24) & 0xFF;
+                    GenerateMaskKey(maskKey);
 
                     // Mask and insert the bytes on the right spots
                     uint8_t* source = &dataFrame[usedSize - 1 + 4];
@@ -138,7 +132,7 @@ namespace Web {
                 result += usedSize;
             }
 
-            if (((_controlStatus & (REQUEST_CLOSE | REQUEST_PING | REQUEST_PONG)) != 0) && ((result + 1) < maxSendSize)) {
+            if (((_controlStatus & (REQUEST_CLOSE | REQUEST_PING | REQUEST_PONG)) != 0) && (result + (((_setFlags & MASKING_FRAME) != 0) ? 6 : 2) < maxSendSize)) {
                 if ((_controlStatus & REQUEST_CLOSE) != 0) {
                     dataFrame[result++] = FINISHING_FRAME | Protocol::CLOSE;
                     _controlStatus &= (~REQUEST_CLOSE);
@@ -154,6 +148,14 @@ namespace Web {
                     _controlStatus &= (~REQUEST_PONG);
                     dataFrame[result++] = (_setFlags & MASKING_FRAME);
                 }
+
+                if ((_setFlags & MASKING_FRAME) != 0) {
+                  // Now it seems only control message, hence append with masking keys
+                  uint8_t maskKey[4];
+                  GenerateMaskKey(maskKey);
+                  ::memcpy(&dataFrame[result], &maskKey, 4);
+                  result += 4;
+                }
             }
 
             return (result);

diff --git a/Source/websocket/WebSocketLink.h b/Source/websocket/WebSocketLink.h
@@ -122,6 +122,18 @@ namespace Web {
             uint16_t Encoder(uint8_t* dataFrame, const uint16_t maxSendSize, const uint16_t usedSize);
             uint16_t Decoder(uint8_t* dataFrame, uint16_t& receivedSize);
 
+        private:
+            inline void GenerateMaskKey(uint8_t *maskKey)
+            {
+                uint32_t value;
+                // Generate a new mask value
+                Crypto::Random(value);
+                maskKey[0] = value & 0xFF;
+                maskKey[1] = (value >> 8) & 0xFF;
+                maskKey[2] = (value >> 16) & 0xFF;
+                maskKey[3] = (value >> 24) & 0xFF;
+            }
+
         private:
             uint8_t _setFlags;
             uint8_t _progressInfo;