Last active
July 3, 2025 01:55
-
-
Save Vizonex/44761bae90cbac5e1642a5ee097cef6c to your computer and use it in GitHub Desktop.
If Aiohttp's _http_writer were moved to C and utilized multidict's upcoming C-API it would behave something like this. I know the iterator is still not mature yet because there's 4 of them but maybe it could be optimized to just 1 in the future.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| // If Aiohttp's _http_writer were moved to C and utilized multidict's upcoming C-API it would behave something like this. | |
| // It's a little bit bigger than cython's version because more care is required to work with it but the higher control is well worth it. | |
| // - Vizonex | |
| // Maybe the next step might be to start looking at making everything threadsafe so that multiple threads could share different | |
| // http writers? | |
| #include <Python.h> | |
| #include <string.h> | |
| #include <stdint.h> | |
| #include <multidict_api.h> | |
| #define BUF_SIZE 16 * 1024 // 16KiB | |
| typedef struct { | |
| MultiDict_CAPI* capi; | |
| // Didn't feel right to make this into a global in C | |
| // it wouldn't be threadsafe if done like so. | |
| char BUFFER[BUF_SIZE]; | |
| } mod_state; | |
| static inline mod_state * | |
| get_mod_state(PyObject *mod) | |
| { | |
| mod_state *state = (mod_state *)PyModule_GetState(mod); | |
| assert(state != NULL); | |
| return state; | |
| } | |
| typedef struct Writer_s { | |
| char* buf; | |
| Py_ssize_t size; | |
| Py_ssize_t pos; | |
| } Writer; | |
| static inline void | |
| _init_writer(mod_state* state, Writer* writer){ | |
| writer->buf = &state->BUFFER[0]; | |
| writer->size = BUF_SIZE; | |
| writer->pos = 0; | |
| } | |
| static inline void | |
| _release_writer(mod_state* state, Writer* writer){ | |
| if (writer->buf != state->BUFFER){ | |
| PyMem_Free(writer->buf); | |
| } | |
| } | |
| static inline int | |
| _write_byte(mod_state* state, Writer* writer, uint8_t ch){ | |
| char* buf; | |
| Py_ssize_t size; | |
| if (writer->pos == writer->size){ | |
| // reallocate | |
| size = writer->size + (BUF_SIZE); | |
| if (writer->buf == state->BUFFER){ | |
| buf = (char*)PyMem_Malloc(size); | |
| if (buf == NULL){ | |
| goto no_memory; | |
| } | |
| memcpy(buf, writer->buf, writer->size); | |
| } else { | |
| buf = (char*)PyMem_Realloc(writer->buf, size); | |
| if (buf == NULL){ | |
| goto no_memory; | |
| } | |
| } | |
| writer->buf = buf; | |
| writer->size = size; | |
| } | |
| writer->buf[writer->pos] = (char)ch; | |
| writer->pos++; | |
| return 0; | |
| no_memory: | |
| PyErr_NoMemory(); | |
| return -1; | |
| } | |
| // this is where the fun begins... | |
| static inline int | |
| _write_utf8(mod_state* state, Writer* writer, Py_UCS4 symbol){ | |
| uint64_t utf = (uint64_t)symbol; | |
| if (utf < 0x80){ | |
| return _write_byte(state, writer, (uint8_t)utf); | |
| } | |
| else if (utf < 0x800){ | |
| if (_write_byte(state, writer, (uint8_t)(0xc0 | (utf >> 6)) < 0)) | |
| return -1; | |
| return _write_byte(state, writer, (uint8_t)(0x80 | (utf & 0x3f))); | |
| } | |
| else if (0xD800 <= utf <= 0xDFFF){ | |
| // Surrogate pair, ignored | |
| return 0; | |
| } | |
| else if (utf < 0x10000){ | |
| if ((_write_byte(state, writer, (uint8_t)(0xe0 | (utf >> 12))) < 0)) | |
| return -1; | |
| if (_write_byte(state, writer, (uint8_t)(0x80 | ((utf >> 6) & 0x3f))) < 0) | |
| return -1; | |
| return _write_byte(state, writer, (uint8_t)(0x80 | (utf & 0x3f))); | |
| } | |
| else if (utf > 0x10FFFF) { | |
| // symbol is too large | |
| return 0; | |
| } | |
| else { | |
| if (_write_byte(state, writer, (uint8_t)(0xf0 | (utf >> 18))) < 0) | |
| return -1; | |
| if (_write_byte(state, writer, | |
| (uint8_t)(0x80 | ((utf >> 12) & 0x3f))) < 0) | |
| return -1; | |
| if (_write_byte(state, writer, | |
| (uint8_t)(0x80 | ((utf >> 6) & 0x3f))) < 0) | |
| return -1; | |
| return _write_byte(state, writer, (uint8_t)(0x80 | (utf & 0x3f))); | |
| } | |
| } | |
| static inline int | |
| _write_str(mod_state* state, Writer* writer, PyObject* s){ | |
| Py_UCS4* ch = PyUnicode_4BYTE_DATA(s); | |
| if (ch == NULL){ | |
| return -1; | |
| } | |
| Py_ssize_t len = PyUnicode_GET_LENGTH(s); | |
| if (!len){ | |
| return 0; | |
| } | |
| for (Py_ssize_t i = 0; i < len; i++){ | |
| if (_write_utf8(state, writer, ch[i]) < 0){ | |
| return -1; | |
| } | |
| } | |
| return 0; | |
| } | |
| static inline int | |
| _write_str_raise_on_nlcr(mod_state* state, Writer* writer, PyObject* s){ | |
| Py_UCS4* chars, ch; | |
| PyObject* out_str; | |
| Py_ssize_t len; | |
| if (PyUnicode_CheckExact(s)){ | |
| out_str = s; | |
| } | |
| else if (IStr_CheckExact(state->capi, s)){ | |
| out_str = PyObject_Str(s); | |
| } | |
| else if (!PyUnicode_Check(s)){ | |
| PyErr_Format( | |
| PyExc_TypeError, "Cannot serialize non-str key %R", s | |
| ); | |
| return -1; | |
| } | |
| else { | |
| out_str = PyObject_Str(s); | |
| } | |
| // Null Check | |
| if (out_str == NULL){ | |
| return -1; | |
| } | |
| len = PyUnicode_GET_LENGTH(out_str); | |
| chars = PyUnicode_4BYTE_DATA(s); | |
| for (Py_ssize_t i = 0; i < len; i++){ | |
| ch = chars[i]; | |
| if ((ch == 0x0D) || (ch == 0x0A)){ | |
| PyErr_SetString(PyExc_ValueError, | |
| "Newline or carriage return detected in headers. " | |
| "Potential header injection attack."); | |
| return -1; | |
| } | |
| if (_write_utf8(state, writer, ch) < 0){ | |
| return -1; | |
| } | |
| } | |
| return 0; | |
| } | |
| // Cython worries too much about the positional arguments in the _http_writer.pyx version | |
| // This version throws all of that out the window assuming that keywords aren't required. | |
| static PyObject* | |
| _serialize_headers(PyObject* self, PyObject *const *args, Py_ssize_t nargs){ | |
| if (nargs != 2){ | |
| PyErr_Format( | |
| PyExc_TypeError, | |
| "_serlize_headers should be called with 2 positional arguments got %d", | |
| nargs | |
| ); | |
| return NULL; | |
| } | |
| mod_state *state = get_mod_state(self); | |
| PyObject *key, *val, *iter, *data = NULL; | |
| iter = MultiDictIter_New(state->capi, args[1]); | |
| if (iter != NULL){ | |
| return NULL; | |
| } | |
| Writer writer; | |
| _init_writer(state, &writer); | |
| if (_write_str(state, &writer, args[0]) < 0){ | |
| goto finish; | |
| }; | |
| if (_write_byte(state, &writer, '\r') < 0){ | |
| goto finish; | |
| } | |
| if (_write_byte(state, &writer, '\n') < 0){ | |
| goto finish; | |
| } | |
| int status = 1; | |
| while (status){ | |
| status = MultiDictIter_Next(state->capi, iter, &key, &val); | |
| // check if something internally bad happened... | |
| if (status < 0){ | |
| goto finish; | |
| } | |
| if (_write_str_raise_on_nlcr(state, &writer, key)){ | |
| goto finish; | |
| } | |
| if (_write_byte(state, &writer, ':') < 0){ | |
| goto finish; | |
| } | |
| if (_write_byte(state, &writer, ' ') < 0){ | |
| goto finish; | |
| } | |
| if (_write_str_raise_on_nlcr(state, &writer, key)){ | |
| goto finish; | |
| } | |
| if (_write_byte(state, &writer, '\r') < 0){ | |
| goto finish; | |
| } | |
| if (_write_byte(state, &writer, '\n') < 0){ | |
| goto finish; | |
| } | |
| } | |
| if (_write_byte(state, &writer, '\r') < 0){ | |
| goto finish; | |
| } | |
| if (_write_byte(state, &writer, '\n') < 0){ | |
| goto finish; | |
| } | |
| PyObject* data = PyBytes_FromStringAndSize(writer.buf, writer.pos); | |
| finish: | |
| Py_XDECREF(iter); | |
| _release_writer(state, &writer); | |
| return data; | |
| } | |
| /* module slots */ | |
| static int | |
| module_traverse(PyObject *mod, visitproc visit, void *arg) | |
| { | |
| return 0; | |
| } | |
| static int | |
| module_clear(PyObject *mod) | |
| { | |
| return 0; | |
| } | |
| static void | |
| module_free(void *mod) | |
| { | |
| (void)module_clear((PyObject *)mod); | |
| } | |
| static PyMethodDef module_methods[] = { | |
| {"_serialize_headers", (PyCFunction)_serialize_headers, METH_FASTCALL}, | |
| {NULL, NULL} | |
| }; | |
| static int | |
| module_exec(PyObject *mod) | |
| { | |
| mod_state *state = get_mod_state(mod); | |
| state->capi = MultiDict_Import(); | |
| if (state->capi == NULL) { | |
| return -1; | |
| } | |
| return 0; | |
| } | |
| static struct PyModuleDef_Slot module_slots[] = { | |
| {Py_mod_exec, module_exec}, | |
| #if PY_VERSION_HEX >= 0x030c00f0 | |
| {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED}, | |
| #endif | |
| #if PY_VERSION_HEX >= 0x030d00f0 | |
| {Py_mod_gil, Py_MOD_GIL_NOT_USED}, | |
| #endif | |
| {0, NULL}, | |
| }; | |
| static PyModuleDef http_writer_module = { | |
| .m_base = PyModuleDef_HEAD_INIT, | |
| .m_name = "_http_writer", | |
| .m_size = sizeof(mod_state), | |
| .m_methods = module_methods, | |
| .m_slots = module_slots, | |
| .m_traverse = module_traverse, | |
| .m_clear = module_clear, | |
| .m_free = (freefunc)module_free, | |
| }; | |
| PyMODINIT_FUNC | |
| PyInit__http_writer(void) | |
| { | |
| return PyModuleDef_Init(&http_writer_module); | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment