1 files changed, 56 insertions, 0 deletions
diff --git a/third_party/python/aiohttp/aiohttp/_websocket.pyx b/third_party/python/aiohttp/aiohttp/_websocket.pyx
new file mode 100644
index 0000000000..94318d2b1b
--- /dev/null
+++ b/third_party/python/aiohttp/aiohttp/_websocket.pyx
@@ -0,0 +1,56 @@
+from cpython cimport PyBytes_AsString
+
+
+#from cpython cimport PyByteArray_AsString # cython still not exports that
+cdef extern from "Python.h":
+    char* PyByteArray_AsString(bytearray ba) except NULL
+
+from libc.stdint cimport uint32_t, uint64_t, uintmax_t
+
+
+def _websocket_mask_cython(object mask, object data):
+    """Note, this function mutates its `data` argument
+    """
+    cdef:
+        Py_ssize_t data_len, i
+        # bit operations on signed integers are implementation-specific
+        unsigned char * in_buf
+        const unsigned char * mask_buf
+        uint32_t uint32_msk
+        uint64_t uint64_msk
+
+    assert len(mask) == 4
+
+    if not isinstance(mask, bytes):
+        mask = bytes(mask)
+
+    if isinstance(data, bytearray):
+        data = <bytearray>data
+    else:
+        data = bytearray(data)
+
+    data_len = len(data)
+    in_buf = <unsigned char*>PyByteArray_AsString(data)
+    mask_buf = <const unsigned char*>PyBytes_AsString(mask)
+    uint32_msk = (<uint32_t*>mask_buf)[0]
+
+    # TODO: align in_data ptr to achieve even faster speeds
+    # does it need in python ?! malloc() always aligns to sizeof(long) bytes
+
+    if sizeof(size_t) >= 8:
+        uint64_msk = uint32_msk
+        uint64_msk = (uint64_msk << 32) | uint32_msk
+
+        while data_len >= 8:
+            (<uint64_t*>in_buf)[0] ^= uint64_msk
+            in_buf += 8
+            data_len -= 8
+
+
+    while data_len >= 4:
+        (<uint32_t*>in_buf)[0] ^= uint32_msk
+        in_buf += 4
+        data_len -= 4
+
+    for i in range(0, data_len):
+        in_buf[i] ^= mask_buf[i]