Skip to content

Instantly share code, notes, and snippets.

@Vizonex
Last active July 3, 2025 01:55
Show Gist options
  • Select an option

  • Save Vizonex/44761bae90cbac5e1642a5ee097cef6c to your computer and use it in GitHub Desktop.

Select an option

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.
// 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