speedups.c 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207
  1. /* C implementation of performance sensitive functions. */
  2. #define PY_SSIZE_T_CLEAN
  3. #include <Python.h>
  4. #include <stdint.h> /* uint32_t, uint64_t */
  5. #if __SSE2__
  6. #include <emmintrin.h>
  7. #endif
  8. static const Py_ssize_t MASK_LEN = 4;
  9. /* Similar to PyBytes_AsStringAndSize, but accepts more types */
  10. static int
  11. _PyBytesLike_AsStringAndSize(PyObject *obj, PyObject **tmp, char **buffer, Py_ssize_t *length)
  12. {
  13. // This supports bytes, bytearrays, and memoryview objects,
  14. // which are common data structures for handling byte streams.
  15. // websockets.framing.prepare_data() returns only these types.
  16. // If *tmp isn't NULL, the caller gets a new reference.
  17. if (PyBytes_Check(obj))
  18. {
  19. *tmp = NULL;
  20. *buffer = PyBytes_AS_STRING(obj);
  21. *length = PyBytes_GET_SIZE(obj);
  22. }
  23. else if (PyByteArray_Check(obj))
  24. {
  25. *tmp = NULL;
  26. *buffer = PyByteArray_AS_STRING(obj);
  27. *length = PyByteArray_GET_SIZE(obj);
  28. }
  29. else if (PyMemoryView_Check(obj))
  30. {
  31. *tmp = PyMemoryView_GetContiguous(obj, PyBUF_READ, 'C');
  32. if (*tmp == NULL)
  33. {
  34. return -1;
  35. }
  36. Py_buffer *mv_buf;
  37. mv_buf = PyMemoryView_GET_BUFFER(*tmp);
  38. *buffer = mv_buf->buf;
  39. *length = mv_buf->len;
  40. }
  41. else
  42. {
  43. PyErr_Format(
  44. PyExc_TypeError,
  45. "expected a bytes-like object, %.200s found",
  46. Py_TYPE(obj)->tp_name);
  47. return -1;
  48. }
  49. return 0;
  50. }
  51. /* C implementation of websockets.utils.apply_mask */
  52. static PyObject *
  53. apply_mask(PyObject *self, PyObject *args, PyObject *kwds)
  54. {
  55. // In order to support various bytes-like types, accept any Python object.
  56. static char *kwlist[] = {"data", "mask", NULL};
  57. PyObject *input_obj;
  58. PyObject *mask_obj;
  59. // A pointer to a char * + length will be extracted from the data and mask
  60. // arguments, possibly via a Py_buffer.
  61. PyObject *input_tmp = NULL;
  62. char *input;
  63. Py_ssize_t input_len;
  64. PyObject *mask_tmp = NULL;
  65. char *mask;
  66. Py_ssize_t mask_len;
  67. // Initialize a PyBytesObject then get a pointer to the underlying char *
  68. // in order to avoid an extra memory copy in PyBytes_FromStringAndSize.
  69. PyObject *result = NULL;
  70. char *output;
  71. // Other variables.
  72. Py_ssize_t i = 0;
  73. // Parse inputs.
  74. if (!PyArg_ParseTupleAndKeywords(
  75. args, kwds, "OO", kwlist, &input_obj, &mask_obj))
  76. {
  77. goto exit;
  78. }
  79. if (_PyBytesLike_AsStringAndSize(input_obj, &input_tmp, &input, &input_len) == -1)
  80. {
  81. goto exit;
  82. }
  83. if (_PyBytesLike_AsStringAndSize(mask_obj, &mask_tmp, &mask, &mask_len) == -1)
  84. {
  85. goto exit;
  86. }
  87. if (mask_len != MASK_LEN)
  88. {
  89. PyErr_SetString(PyExc_ValueError, "mask must contain 4 bytes");
  90. goto exit;
  91. }
  92. // Create output.
  93. result = PyBytes_FromStringAndSize(NULL, input_len);
  94. if (result == NULL)
  95. {
  96. goto exit;
  97. }
  98. // Since we juste created result, we don't need error checks.
  99. output = PyBytes_AS_STRING(result);
  100. // Perform the masking operation.
  101. // Apparently GCC cannot figure out the following optimizations by itself.
  102. // We need a new scope for MSVC 2010 (non C99 friendly)
  103. {
  104. #if __SSE2__
  105. // With SSE2 support, XOR by blocks of 16 bytes = 128 bits.
  106. // Since we cannot control the 16-bytes alignment of input and output
  107. // buffers, we rely on loadu/storeu rather than load/store.
  108. Py_ssize_t input_len_128 = input_len & ~15;
  109. __m128i mask_128 = _mm_set1_epi32(*(uint32_t *)mask);
  110. for (; i < input_len_128; i += 16)
  111. {
  112. __m128i in_128 = _mm_loadu_si128((__m128i *)(input + i));
  113. __m128i out_128 = _mm_xor_si128(in_128, mask_128);
  114. _mm_storeu_si128((__m128i *)(output + i), out_128);
  115. }
  116. #else
  117. // Without SSE2 support, XOR by blocks of 8 bytes = 64 bits.
  118. // We assume the memory allocator aligns everything on 8 bytes boundaries.
  119. Py_ssize_t input_len_64 = input_len & ~7;
  120. uint32_t mask_32 = *(uint32_t *)mask;
  121. uint64_t mask_64 = ((uint64_t)mask_32 << 32) | (uint64_t)mask_32;
  122. for (; i < input_len_64; i += 8)
  123. {
  124. *(uint64_t *)(output + i) = *(uint64_t *)(input + i) ^ mask_64;
  125. }
  126. #endif
  127. }
  128. // XOR the remainder of the input byte by byte.
  129. for (; i < input_len; i++)
  130. {
  131. output[i] = input[i] ^ mask[i & (MASK_LEN - 1)];
  132. }
  133. exit:
  134. Py_XDECREF(input_tmp);
  135. Py_XDECREF(mask_tmp);
  136. return result;
  137. }
  138. static PyMethodDef speedups_methods[] = {
  139. {
  140. "apply_mask",
  141. (PyCFunction)apply_mask,
  142. METH_VARARGS | METH_KEYWORDS,
  143. "Apply masking to the data of a WebSocket message.",
  144. },
  145. {NULL, NULL, 0, NULL}, /* Sentinel */
  146. };
  147. static struct PyModuleDef speedups_module = {
  148. PyModuleDef_HEAD_INIT,
  149. "websocket.speedups", /* m_name */
  150. "C implementation of performance sensitive functions.",
  151. /* m_doc */
  152. -1, /* m_size */
  153. speedups_methods, /* m_methods */
  154. NULL,
  155. NULL,
  156. NULL,
  157. NULL
  158. };
  159. PyMODINIT_FUNC
  160. PyInit_speedups(void)
  161. {
  162. return PyModule_Create(&speedups_module);
  163. }