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
11cdef char BUFFER[BUF_SIZE]
12
13cdef object _istr = istr
14
15
16# ----------------- writer ---------------------------
17
18cdef struct Writer:
19    char *buf
20    Py_ssize_t size
21    Py_ssize_t pos
22
23
24cdef inline void _init_writer(Writer* writer):
25    writer.buf = &BUFFER[0]
26    writer.size = BUF_SIZE
27    writer.pos = 0
28
29
30cdef inline void _release_writer(Writer* writer):
31    if writer.buf != BUFFER:
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 writer.buf == BUFFER:
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.buf[writer.pos] = <char>ch
56    writer.pos += 1
57    return 0
58
59
60cdef inline int _write_utf8(Writer* writer, Py_UCS4 symbol):
61    cdef uint64_t utf = <uint64_t> symbol
62
63    if utf < 0x80:
64        return _write_byte(writer, <uint8_t>utf)
65    elif utf < 0x800:
66        if _write_byte(writer, <uint8_t>(0xc0 | (utf >> 6))) < 0:
67            return -1
68        return _write_byte(writer,  <uint8_t>(0x80 | (utf & 0x3f)))
69    elif 0xD800 <= utf <= 0xDFFF:
70        # surogate pair, ignored
71        return 0
72    elif utf < 0x10000:
73        if _write_byte(writer, <uint8_t>(0xe0 | (utf >> 12))) < 0:
74            return -1
75        if _write_byte(writer, <uint8_t>(0x80 | ((utf >> 6) & 0x3f))) < 0:
76            return -1
77        return _write_byte(writer, <uint8_t>(0x80 | (utf & 0x3f)))
78    elif utf > 0x10FFFF:
79        # symbol is too large
80        return 0
81    else:
82        if _write_byte(writer,  <uint8_t>(0xf0 | (utf >> 18))) < 0:
83            return -1
84        if _write_byte(writer,
85                       <uint8_t>(0x80 | ((utf >> 12) & 0x3f))) < 0:
86           return -1
87        if _write_byte(writer,
88                       <uint8_t>(0x80 | ((utf >> 6) & 0x3f))) < 0:
89            return -1
90        return _write_byte(writer, <uint8_t>(0x80 | (utf & 0x3f)))
91
92
93cdef inline int _write_str(Writer* writer, str s):
94    cdef Py_UCS4 ch
95    for ch in s:
96        if _write_utf8(writer, ch) < 0:
97            return -1
98
99
100# --------------- _serialize_headers ----------------------
101
102cdef str to_str(object s):
103    typ = type(s)
104    if typ is str:
105        return <str>s
106    elif typ is _istr:
107        return PyObject_Str(s)
108    elif not isinstance(s, str):
109        raise TypeError("Cannot serialize non-str key {!r}".format(s))
110    else:
111        return str(s)
112
113
114def _serialize_headers(str status_line, headers):
115    cdef Writer writer
116    cdef object key
117    cdef object val
118    cdef bytes ret
119
120    _init_writer(&writer)
121
122    try:
123        if _write_str(&writer, status_line) < 0:
124            raise
125        if _write_byte(&writer, b'\r') < 0:
126            raise
127        if _write_byte(&writer, b'\n') < 0:
128            raise
129
130        for key, val in headers.items():
131            if _write_str(&writer, to_str(key)) < 0:
132                raise
133            if _write_byte(&writer, b':') < 0:
134                raise
135            if _write_byte(&writer, b' ') < 0:
136                raise
137            if _write_str(&writer, to_str(val)) < 0:
138                raise
139            if _write_byte(&writer, b'\r') < 0:
140                raise
141            if _write_byte(&writer, b'\n') < 0:
142                raise
143
144        if _write_byte(&writer, b'\r') < 0:
145            raise
146        if _write_byte(&writer, b'\n') < 0:
147            raise
148
149        return PyBytes_FromStringAndSize(writer.buf, writer.pos)
150    finally:
151        _release_writer(&writer)
152