codekingpro/portable-devtools

1from cpython cimport PyBytes_AsString
2
3
4#from cpython cimport PyByteArray_AsString # cython still not exports that
5cdef extern from "Python.h":
6    char* PyByteArray_AsString(bytearray ba) except NULL
7
8from libc.stdint cimport uint32_t, uint64_t, uintmax_t
9
10
11cpdef void _websocket_mask_cython(bytes mask, bytearray data):
12    """Note, this function mutates its `data` argument
13    """
14    cdef:
15        Py_ssize_t data_len, i
16        # bit operations on signed integers are implementation-specific
17        unsigned char * in_buf
18        const unsigned char * mask_buf
19        uint32_t uint32_msk
20        uint64_t uint64_msk
21
22    assert len(mask) == 4
23
24    data_len = len(data)
25    in_buf = <unsigned char*>PyByteArray_AsString(data)
26    mask_buf = <const unsigned char*>PyBytes_AsString(mask)
27    uint32_msk = (<uint32_t*>mask_buf)[0]
28
29    # TODO: align in_data ptr to achieve even faster speeds
30    # does it need in python ?! malloc() always aligns to sizeof(long) bytes
31
32    if sizeof(size_t) >= 8:
33        uint64_msk = uint32_msk
34        uint64_msk = (uint64_msk << 32) | uint32_msk
35
36        while data_len >= 8:
37            (<uint64_t*>in_buf)[0] ^= uint64_msk
38            in_buf += 8
39            data_len -= 8
40
41
42    while data_len >= 4:
43        (<uint32_t*>in_buf)[0] ^= uint32_msk
44        in_buf += 4
45        data_len -= 4
46
47    for i in range(0, data_len):
48        in_buf[i] ^= mask_buf[i]
49