1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
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
 
 
cpdef void _websocket_mask_cython(bytes mask, bytearray 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
 
    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]