1 files changed, 55 insertions, 6 deletions

diff --git a/module/lib/mod_pywebsocket/util.py b/module/lib/mod_pywebsocket/util.py
index d60b53f75..7bb0b5d9e 100644
--- a/module/lib/mod_pywebsocket/util.py
+++ b/module/lib/mod_pywebsocket/util.py
@@ -31,6 +31,8 @@
 """WebSocket utilities.
 """
 
+
+import array
 import errno
 
 # Import hash classes from a module available and recommended for each Python
@@ -160,6 +162,34 @@ class NoopMasker(object):
         return s
 
 
+class RepeatedXorMasker(object):
+    """A masking object that applies XOR on the string given to mask method
+    with the masking bytes given to the constructor repeatedly. This object
+    remembers the position in the masking bytes the last mask method call
+    ended and resumes from that point on the next mask method call.
+    """
+
+    def __init__(self, mask):
+        self._mask = map(ord, mask)
+        self._mask_size = len(self._mask)
+        self._count = 0
+
+    def mask(self, s):
+        result = array.array('B')
+        result.fromstring(s)
+        # Use temporary local variables to eliminate the cost to access
+        # attributes
+        count = self._count
+        mask = self._mask
+        mask_size = self._mask_size
+        for i in xrange(len(result)):
+            result[i] ^= mask[count]
+            count = (count + 1) % mask_size
+        self._count = count
+
+        return result.tostring()
+
+
 class DeflateRequest(object):
     """A wrapper class for request object to intercept send and recv to perform
     deflate compression and decompression transparently.
@@ -202,6 +232,12 @@ class _Deflater(object):
         self._compress = zlib.compressobj(
             zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -window_bits)
 
+    def compress(self, bytes):
+        compressed_bytes = self._compress.compress(bytes)
+        self._logger.debug('Compress input %r', bytes)
+        self._logger.debug('Compress result %r', compressed_bytes)
+        return compressed_bytes
+
     def compress_and_flush(self, bytes):
         compressed_bytes = self._compress.compress(bytes)
         compressed_bytes += self._compress.flush(zlib.Z_SYNC_FLUSH)
@@ -209,6 +245,12 @@ class _Deflater(object):
         self._logger.debug('Compress result %r', compressed_bytes)
         return compressed_bytes
 
+    def compress_and_finish(self, bytes):
+        compressed_bytes = self._compress.compress(bytes)
+        compressed_bytes += self._compress.flush(zlib.Z_FINISH)
+        self._logger.debug('Compress input %r', bytes)
+        self._logger.debug('Compress result %r', compressed_bytes)
+        return compressed_bytes
 
 class _Inflater(object):
 
@@ -288,14 +330,21 @@ class _RFC1979Deflater(object):
         self._window_bits = window_bits
         self._no_context_takeover = no_context_takeover
 
-    def filter(self, bytes):
-        if self._deflater is None or self._no_context_takeover:
+    def filter(self, bytes, flush=True, bfinal=False):
+        if self._deflater is None or (self._no_context_takeover and flush):
             self._deflater = _Deflater(self._window_bits)
 
-        # Strip last 4 octets which is LEN and NLEN field of a non-compressed
-        # block added for Z_SYNC_FLUSH.
-        return self._deflater.compress_and_flush(bytes)[:-4]
-
+        if bfinal:
+            result = self._deflater.compress_and_finish(bytes)
+            # Add a padding block with BFINAL = 0 and BTYPE = 0.
+            result = result + chr(0)
+            self._deflater = None
+            return result
+        if flush:
+            # Strip last 4 octets which is LEN and NLEN field of a
+            # non-compressed block added for Z_SYNC_FLUSH.
+            return self._deflater.compress_and_flush(bytes)[:-4]
+        return self._deflater.compress(bytes)
 
 class _RFC1979Inflater(object):
     """A decompressor class for byte sequence compressed and flushed following