WebSocket: mask 4 bytes at a time

src/ws_decoder.cpp

@@ -131,6 +131,7 @@ int zmq::ws_decoder_t::long_size_ready (unsigned char const *read_from_)
 int zmq::ws_decoder_t::mask_ready (unsigned char const *read_from_)
 {
     memcpy (_mask, _tmpbuf, 4);
+    memcpy (_mask + 4, _tmpbuf, 4);
 
     if (_opcode == ws_protocol_t::opcode_binary) {
         if (_size == 0)
@@ -238,8 +239,8 @@ int zmq::ws_decoder_t::message_ready (unsigned char const *)
 
         unsigned char *data =
           static_cast<unsigned char *> (_in_progress.data ());
-        for (size_t i = 0; i < _size; ++i, mask_index++)
-            data[i] = data[i] ^ _mask[mask_index % 4];
+
+        ws_mask_payload (data, data, _size, _mask + mask_index);
     }
 
     //  Message is completely read. Signal this to the caller

src/ws_decoder.hpp

@@ -45,7 +45,7 @@ class ws_decoder_t ZMQ_FINAL
     const bool _must_mask;
     uint64_t _size;
     zmq::ws_protocol_t::opcode_t _opcode;
-    unsigned char _mask[4];
+    unsigned char _mask[8]; // repeated twice to simplify mask rotation
 
     ZMQ_NON_COPYABLE_NOR_MOVABLE (ws_decoder_t)
 };

src/ws_encoder.cpp

@@ -61,10 +61,12 @@ void zmq::ws_encoder_t::message_ready ()
         offset += 8;
     }
 
+    unsigned char mask_buf[8];
     if (_must_mask) {
         const uint32_t random = generate_random ();
         put_uint32 (_tmp_buf + offset, random);
-        put_uint32 (_mask, random);
+        put_uint32 (mask_buf, random);
+        put_uint32 (mask_buf + 4, random);
         offset += 4;
     }
 
@@ -78,15 +80,18 @@ void zmq::ws_encoder_t::message_ready ()
             protocol_flags |= ws_protocol_t::command_flag;
 
         _tmp_buf[offset++] =
-          _must_mask ? protocol_flags ^ _mask[mask_index++] : protocol_flags;
+          _must_mask ? protocol_flags ^ mask_buf[mask_index++] : protocol_flags;
     }
 
     //  Encode the subscribe/cancel byte.
     //  TODO: remove once there is an opcode for subscribe/cancel
     if (in_progress ()->is_subscribe ())
-        _tmp_buf[offset++] = _must_mask ? 1 ^ _mask[mask_index++] : 1;
+        _tmp_buf[offset++] = _must_mask ? 1 ^ mask_buf[mask_index++] : 1;
     else if (in_progress ()->is_cancel ())
-        _tmp_buf[offset++] = _must_mask ? 0 ^ _mask[mask_index++] : 0;
+        _tmp_buf[offset++] = _must_mask ? 0 ^ mask_buf[mask_index++] : 0;
+
+    // Store the rotated mask, so ws_mask_payload can use it directly
+    memcpy(_mask, mask_buf + mask_index, 4);
 
     next_step (_tmp_buf, offset, &ws_encoder_t::size_ready, false);
 }
@@ -109,14 +114,7 @@ void zmq::ws_encoder_t::size_ready ()
             dest = static_cast<unsigned char *> (_masked_msg.data ());
         }
 
-        int mask_index = 0;
-        if (_is_binary)
-            ++mask_index;
-        //  TODO: remove once there is an opcode for subscribe/cancel
-        if (in_progress ()->is_subscribe () || in_progress ()->is_cancel ())
-            ++mask_index;
-        for (size_t i = 0; i < size; ++i, mask_index++)
-            dest[i] = src[i] ^ _mask[mask_index % 4];
+        ws_mask_payload (dest, src, size, _mask);
 
         next_step (dest, size, &ws_encoder_t::message_ready, true);
     } else {

src/ws_protocol.hpp

@@ -3,8 +3,32 @@
 #ifndef __ZMQ_WS_PROTOCOL_HPP_INCLUDED__
 #define __ZMQ_WS_PROTOCOL_HPP_INCLUDED__
 
+#include "stdint.hpp"
+
+#include <string.h>
+
 namespace zmq
 {
+inline void ws_mask_payload (unsigned char *dest_,
+                             const unsigned char *src_,
+                             size_t size_,
+                             const unsigned char *mask_)
+{
+    uint32_t mask32;
+    memcpy (&mask32, mask_, sizeof (mask32));
+
+    size_t i = 0;
+    for (; i + 4 <= size_; i += 4) {
+        uint32_t chunk;
+        memcpy (&chunk, src_ + i, sizeof (chunk));
+        chunk ^= mask32;
+        memcpy (dest_ + i, &chunk, sizeof (chunk));
+    }
+
+    for (; i < size_; ++i)
+        dest_[i] = src_[i] ^ mask_[i & 3];
+}
+
 //  Definition of constants for WS transport protocol.
 class ws_protocol_t
 {