codekingpro/portable-devtools

1from cpython.bytes cimport PyBytes_FromStringAndSize
2from cpython.exc cimport PyErr_NoMemory
3from cpython.mem cimport PyMem_Free, PyMem_Malloc, PyMem_Realloc
4from cpython.object cimport PyObject_Str
5from libc.stdint cimport uint8_t, uint64_t
6from libc.string cimport memcpy
7
8from multidict import istr
9
10DEF BUF_SIZE = 16 * 1024  # 16KiB
11
12cdef object _istr = istr
13
14
15# ----------------- writer ---------------------------
16
17cdef struct Writer:
18    char *buf
19    Py_ssize_t size
20    Py_ssize_t pos
21    bint heap_allocated
22
23cdef inline void _init_writer(Writer* writer, char *buf):
24    writer.buf = buf
25    writer.size = BUF_SIZE
26    writer.pos = 0
27    writer.heap_allocated = 0
28
29
30cdef inline void _release_writer(Writer* writer):
31    if writer.heap_allocated:
32        PyMem_Free(writer.buf)
33
34
35cdef inline int _write_byte(Writer* writer, uint8_t ch):
36    cdef char * buf
37    cdef Py_ssize_t size
38
39    if writer.pos == writer.size:
40        # reallocate
41        size = writer.size + BUF_SIZE
42        if not writer.heap_allocated:
43            buf = <char*>PyMem_Malloc(size)
44            if buf == NULL:
45                PyErr_NoMemory()
46                return -1
47            memcpy(buf, writer.buf, writer.size)
48        else:
49            buf = <char*>PyMem_Realloc(writer.buf, size)
50            if buf == NULL:
51                PyErr_NoMemory()
52                return -1
53        writer.buf = buf
54        writer.size = size
55        writer.heap_allocated = 1
56    writer.buf[writer.pos] = <char>ch
57    writer.pos += 1
58    return 0
59
60
61cdef inline int _write_utf8(Writer* writer, Py_UCS4 symbol):
62    cdef uint64_t utf = <uint64_t> symbol
63
64    if utf < 0x80:
65        return _write_byte(writer, <uint8_t>utf)
66    elif utf < 0x800:
67        if _write_byte(writer, <uint8_t>(0xc0 | (utf >> 6))) < 0:
68            return -1
69        return _write_byte(writer,  <uint8_t>(0x80 | (utf & 0x3f)))
70    elif 0xD800 <= utf <= 0xDFFF:
71        # surogate pair, ignored
72        return 0
73    elif utf < 0x10000:
74        if _write_byte(writer, <uint8_t>(0xe0 | (utf >> 12))) < 0:
75            return -1
76        if _write_byte(writer, <uint8_t>(0x80 | ((utf >> 6) & 0x3f))) < 0:
77            return -1
78        return _write_byte(writer, <uint8_t>(0x80 | (utf & 0x3f)))
79    elif utf > 0x10FFFF:
80        # symbol is too large
81        return 0
82    else:
83        if _write_byte(writer,  <uint8_t>(0xf0 | (utf >> 18))) < 0:
84            return -1
85        if _write_byte(writer,
86                       <uint8_t>(0x80 | ((utf >> 12) & 0x3f))) < 0:
87           return -1
88        if _write_byte(writer,
89                       <uint8_t>(0x80 | ((utf >> 6) & 0x3f))) < 0:
90            return -1
91        return _write_byte(writer, <uint8_t>(0x80 | (utf & 0x3f)))
92
93
94cdef inline int _write_str(Writer* writer, str s):
95    cdef Py_UCS4 ch
96    for ch in s:
97        if _write_utf8(writer, ch) < 0:
98            return -1
99
100
101cdef inline int _write_str_raise_on_nlcr(Writer* writer, object s):
102    cdef Py_UCS4 ch
103    cdef str out_str
104    if type(s) is str:
105        out_str = <str>s
106    elif type(s) is _istr:
107        out_str = PyObject_Str(s)
108    elif not isinstance(s, str):
109        raise TypeError("Cannot serialize non-str key {!r}".format(s))
110    else:
111        out_str = str(s)
112
113    for ch in out_str:
114        if ch in {0x0D, 0x0A, 0x00}:
115            raise ValueError(
116                "Newline, carriage return, or null byte detected in headers. "
117                "Potential header injection attack."
118            )
119        if _write_utf8(writer, ch) < 0:
120            return -1
121
122
123# --------------- _serialize_headers ----------------------
124
125def _serialize_headers(str status_line, headers):
126    cdef Writer writer
127    cdef object key
128    cdef object val
129    cdef char buf[BUF_SIZE]
130
131    _init_writer(&writer, buf)
132
133    try:
134        if _write_str_raise_on_nlcr(&writer, status_line) < 0:
135            raise
136        if _write_byte(&writer, b'\r') < 0:
137            raise
138        if _write_byte(&writer, b'\n') < 0:
139            raise
140
141        for key, val in headers.items():
142            if _write_str_raise_on_nlcr(&writer, key) < 0:
143                raise
144            if _write_byte(&writer, b':') < 0:
145                raise
146            if _write_byte(&writer, b' ') < 0:
147                raise
148            if _write_str_raise_on_nlcr(&writer, val) < 0:
149                raise
150            if _write_byte(&writer, b'\r') < 0:
151                raise
152            if _write_byte(&writer, b'\n') < 0:
153                raise
154
155        if _write_byte(&writer, b'\r') < 0:
156            raise
157        if _write_byte(&writer, b'\n') < 0:
158            raise
159
160        return PyBytes_FromStringAndSize(writer.buf, writer.pos)
161    finally:
162        _release_writer(&writer)
163