codekingpro/portable-devtools
115k
1#define PY_SSIZE_T_CLEAN
2
3#include <Python.h>
4#include <stdint.h>
5
6#define MODULE_NAME "aioquic._buffer"
7
8static PyObject *BufferReadError;
9static PyObject *BufferWriteError;
10
11typedef struct {
12 PyObject_HEAD
13 uint8_t *base;
14 uint8_t *end;
15 uint8_t *pos;
16} BufferObject;
17
18static PyObject *BufferType;
19
20#define CHECK_READ_BOUNDS(self, len) \
21 if (len < 0 || self->pos + len > self->end) { \
22 PyErr_SetString(BufferReadError, "Read out of bounds"); \
23 return NULL; \
24 }
25
26#define CHECK_WRITE_BOUNDS(self, len) \
27 if (self->pos + len > self->end) { \
28 PyErr_SetString(BufferWriteError, "Write out of bounds"); \
29 return NULL; \
30 }
31
32static int
33Buffer_init(BufferObject *self, PyObject *args, PyObject *kwargs)
34{
35 const char *kwlist[] = {"capacity", "data", NULL};
36 Py_ssize_t capacity = 0;
37 const unsigned char *data = NULL;
38 Py_ssize_t data_len = 0;
39
40 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|ny#", (char**)kwlist, &capacity, &data, &data_len))
41 return -1;
42
43 if (data != NULL) {
44 self->base = malloc(data_len);
45 self->end = self->base + data_len;
46 memcpy(self->base, data, data_len);
47 } else {
48 self->base = malloc(capacity);
49 self->end = self->base + capacity;
50 }
51 self->pos = self->base;
52 return 0;
53}
54
55static void
56Buffer_dealloc(BufferObject *self)
57{
58 free(self->base);
59 PyTypeObject *tp = Py_TYPE(self);
60 freefunc free = PyType_GetSlot(tp, Py_tp_free);
61 free(self);
62 Py_DECREF(tp);
63}
64
65static PyObject *
66Buffer_data_slice(BufferObject *self, PyObject *args)
67{
68 Py_ssize_t start, stop;
69 if (!PyArg_ParseTuple(args, "nn", &start, &stop))
70 return NULL;
71
72 if (start < 0 || self->base + start > self->end ||
73 stop < 0 || self->base + stop > self->end ||
74 stop < start) {
75 PyErr_SetString(BufferReadError, "Read out of bounds");
76 return NULL;
77 }
78
79 return PyBytes_FromStringAndSize((const char*)(self->base + start), (stop - start));
80}
81
82static PyObject *
83Buffer_eof(BufferObject *self, PyObject *args)
84{
85 if (self->pos == self->end)
86 Py_RETURN_TRUE;
87 Py_RETURN_FALSE;
88}
89
90static PyObject *
91Buffer_pull_bytes(BufferObject *self, PyObject *args)
92{
93 Py_ssize_t len;
94 if (!PyArg_ParseTuple(args, "n", &len))
95 return NULL;
96
97 CHECK_READ_BOUNDS(self, len);
98
99 PyObject *o = PyBytes_FromStringAndSize((const char*)self->pos, len);
100 self->pos += len;
101 return o;
102}
103
104static PyObject *
105Buffer_pull_uint8(BufferObject *self, PyObject *args)
106{
107 CHECK_READ_BOUNDS(self, 1)
108
109 return PyLong_FromUnsignedLong(
110 (uint8_t)(*(self->pos++))
111 );
112}
113
114static PyObject *
115Buffer_pull_uint16(BufferObject *self, PyObject *args)
116{
117 CHECK_READ_BOUNDS(self, 2)
118
119 uint16_t value = (uint16_t)(*(self->pos)) << 8 |
120 (uint16_t)(*(self->pos + 1));
121 self->pos += 2;
122 return PyLong_FromUnsignedLong(value);
123}
124
125static PyObject *
126Buffer_pull_uint32(BufferObject *self, PyObject *args)
127{
128 CHECK_READ_BOUNDS(self, 4)
129
130 uint32_t value = (uint32_t)(*(self->pos)) << 24 |
131 (uint32_t)(*(self->pos + 1)) << 16 |
132 (uint32_t)(*(self->pos + 2)) << 8 |
133 (uint32_t)(*(self->pos + 3));
134 self->pos += 4;
135 return PyLong_FromUnsignedLong(value);
136}
137
138static PyObject *
139Buffer_pull_uint64(BufferObject *self, PyObject *args)
140{
141 CHECK_READ_BOUNDS(self, 8)
142
143 uint64_t value = (uint64_t)(*(self->pos)) << 56 |
144 (uint64_t)(*(self->pos + 1)) << 48 |
145 (uint64_t)(*(self->pos + 2)) << 40 |
146 (uint64_t)(*(self->pos + 3)) << 32 |
147 (uint64_t)(*(self->pos + 4)) << 24 |
148 (uint64_t)(*(self->pos + 5)) << 16 |
149 (uint64_t)(*(self->pos + 6)) << 8 |
150 (uint64_t)(*(self->pos + 7));
151 self->pos += 8;
152 return PyLong_FromUnsignedLongLong(value);
153}
154
155static PyObject *
156Buffer_pull_uint_var(BufferObject *self, PyObject *args)
157{
158 uint64_t value;
159 CHECK_READ_BOUNDS(self, 1)
160 switch (*(self->pos) >> 6) {
161 case 0:
162 value = *(self->pos++) & 0x3F;
163 break;
164 case 1:
165 CHECK_READ_BOUNDS(self, 2)
166 value = (uint16_t)(*(self->pos) & 0x3F) << 8 |
167 (uint16_t)(*(self->pos + 1));
168 self->pos += 2;
169 break;
170 case 2:
171 CHECK_READ_BOUNDS(self, 4)
172 value = (uint32_t)(*(self->pos) & 0x3F) << 24 |
173 (uint32_t)(*(self->pos + 1)) << 16 |
174 (uint32_t)(*(self->pos + 2)) << 8 |
175 (uint32_t)(*(self->pos + 3));
176 self->pos += 4;
177 break;
178 default:
179 CHECK_READ_BOUNDS(self, 8)
180 value = (uint64_t)(*(self->pos) & 0x3F) << 56 |
181 (uint64_t)(*(self->pos + 1)) << 48 |
182 (uint64_t)(*(self->pos + 2)) << 40 |
183 (uint64_t)(*(self->pos + 3)) << 32 |
184 (uint64_t)(*(self->pos + 4)) << 24 |
185 (uint64_t)(*(self->pos + 5)) << 16 |
186 (uint64_t)(*(self->pos + 6)) << 8 |
187 (uint64_t)(*(self->pos + 7));
188 self->pos += 8;
189 break;
190 }
191 return PyLong_FromUnsignedLongLong(value);
192}
193
194static PyObject *
195Buffer_push_bytes(BufferObject *self, PyObject *args)
196{
197 const unsigned char *data;
198 Py_ssize_t data_len;
199 if (!PyArg_ParseTuple(args, "y#", &data, &data_len))
200 return NULL;
201
202 CHECK_WRITE_BOUNDS(self, data_len)
203
204 memcpy(self->pos, data, data_len);
205 self->pos += data_len;
206 Py_RETURN_NONE;
207}
208
209static PyObject *
210Buffer_push_uint8(BufferObject *self, PyObject *args)
211{
212 uint8_t value;
213 if (!PyArg_ParseTuple(args, "B", &value))
214 return NULL;
215
216 CHECK_WRITE_BOUNDS(self, 1)
217
218 *(self->pos++) = value;
219 Py_RETURN_NONE;
220}
221
222static PyObject *
223Buffer_push_uint16(BufferObject *self, PyObject *args)
224{
225 uint16_t value;
226 if (!PyArg_ParseTuple(args, "H", &value))
227 return NULL;
228
229 CHECK_WRITE_BOUNDS(self, 2)
230
231 *(self->pos++) = (value >> 8);
232 *(self->pos++) = value;
233 Py_RETURN_NONE;
234}
235
236static PyObject *
237Buffer_push_uint32(BufferObject *self, PyObject *args)
238{
239 uint32_t value;
240 if (!PyArg_ParseTuple(args, "I", &value))
241 return NULL;
242
243 CHECK_WRITE_BOUNDS(self, 4)
244 *(self->pos++) = (value >> 24);
245 *(self->pos++) = (value >> 16);
246 *(self->pos++) = (value >> 8);
247 *(self->pos++) = value;
248 Py_RETURN_NONE;
249}
250
251static PyObject *
252Buffer_push_uint64(BufferObject *self, PyObject *args)
253{
254 uint64_t value;
255 if (!PyArg_ParseTuple(args, "K", &value))
256 return NULL;
257
258 CHECK_WRITE_BOUNDS(self, 8)
259 *(self->pos++) = (value >> 56);
260 *(self->pos++) = (value >> 48);
261 *(self->pos++) = (value >> 40);
262 *(self->pos++) = (value >> 32);
263 *(self->pos++) = (value >> 24);
264 *(self->pos++) = (value >> 16);
265 *(self->pos++) = (value >> 8);
266 *(self->pos++) = value;
267 Py_RETURN_NONE;
268}
269
270static PyObject *
271Buffer_push_uint_var(BufferObject *self, PyObject *args)
272{
273 uint64_t value;
274 if (!PyArg_ParseTuple(args, "K", &value))
275 return NULL;
276
277 if (value <= 0x3F) {
278 CHECK_WRITE_BOUNDS(self, 1)
279 *(self->pos++) = value;
280 Py_RETURN_NONE;
281 } else if (value <= 0x3FFF) {
282 CHECK_WRITE_BOUNDS(self, 2)
283 *(self->pos++) = (value >> 8) | 0x40;
284 *(self->pos++) = value;
285 Py_RETURN_NONE;
286 } else if (value <= 0x3FFFFFFF) {
287 CHECK_WRITE_BOUNDS(self, 4)
288 *(self->pos++) = (value >> 24) | 0x80;
289 *(self->pos++) = (value >> 16);
290 *(self->pos++) = (value >> 8);
291 *(self->pos++) = value;
292 Py_RETURN_NONE;
293 } else if (value <= 0x3FFFFFFFFFFFFFFF) {
294 CHECK_WRITE_BOUNDS(self, 8)
295 *(self->pos++) = (value >> 56) | 0xC0;
296 *(self->pos++) = (value >> 48);
297 *(self->pos++) = (value >> 40);
298 *(self->pos++) = (value >> 32);
299 *(self->pos++) = (value >> 24);
300 *(self->pos++) = (value >> 16);
301 *(self->pos++) = (value >> 8);
302 *(self->pos++) = value;
303 Py_RETURN_NONE;
304 } else {
305 PyErr_SetString(PyExc_ValueError, "Integer is too big for a variable-length integer");
306 return NULL;
307 }
308}
309
310static PyObject *
311Buffer_seek(BufferObject *self, PyObject *args)
312{
313 Py_ssize_t pos;
314 if (!PyArg_ParseTuple(args, "n", &pos))
315 return NULL;
316
317 if (pos < 0 || self->base + pos > self->end) {
318 PyErr_SetString(BufferReadError, "Seek out of bounds");
319 return NULL;
320 }
321
322 self->pos = self->base + pos;
323 Py_RETURN_NONE;
324}
325
326static PyObject *
327Buffer_tell(BufferObject *self, PyObject *args)
328{
329 return PyLong_FromSsize_t(self->pos - self->base);
330}
331
332static PyMethodDef Buffer_methods[] = {
333 {"data_slice", (PyCFunction)Buffer_data_slice, METH_VARARGS, ""},
334 {"eof", (PyCFunction)Buffer_eof, METH_VARARGS, ""},
335 {"pull_bytes", (PyCFunction)Buffer_pull_bytes, METH_VARARGS, "Pull bytes."},
336 {"pull_uint8", (PyCFunction)Buffer_pull_uint8, METH_VARARGS, "Pull an 8-bit unsigned integer."},
337 {"pull_uint16", (PyCFunction)Buffer_pull_uint16, METH_VARARGS, "Pull a 16-bit unsigned integer."},
338 {"pull_uint32", (PyCFunction)Buffer_pull_uint32, METH_VARARGS, "Pull a 32-bit unsigned integer."},
339 {"pull_uint64", (PyCFunction)Buffer_pull_uint64, METH_VARARGS, "Pull a 64-bit unsigned integer."},
340 {"pull_uint_var", (PyCFunction)Buffer_pull_uint_var, METH_VARARGS, "Pull a QUIC variable-length unsigned integer."},
341 {"push_bytes", (PyCFunction)Buffer_push_bytes, METH_VARARGS, "Push bytes."},
342 {"push_uint8", (PyCFunction)Buffer_push_uint8, METH_VARARGS, "Push an 8-bit unsigned integer."},
343 {"push_uint16", (PyCFunction)Buffer_push_uint16, METH_VARARGS, "Push a 16-bit unsigned integer."},
344 {"push_uint32", (PyCFunction)Buffer_push_uint32, METH_VARARGS, "Push a 32-bit unsigned integer."},
345 {"push_uint64", (PyCFunction)Buffer_push_uint64, METH_VARARGS, "Push a 64-bit unsigned integer."},
346 {"push_uint_var", (PyCFunction)Buffer_push_uint_var, METH_VARARGS, "Push a QUIC variable-length unsigned integer."},
347 {"seek", (PyCFunction)Buffer_seek, METH_VARARGS, ""},
348 {"tell", (PyCFunction)Buffer_tell, METH_VARARGS, ""},
349 {NULL}
350};
351
352static PyObject*
353Buffer_capacity_getter(BufferObject* self, void *closure) {
354 return PyLong_FromSsize_t(self->end - self->base);
355}
356
357static PyObject*
358Buffer_data_getter(BufferObject* self, void *closure) {
359 return PyBytes_FromStringAndSize((const char*)self->base, self->pos - self->base);
360}
361
362static PyGetSetDef Buffer_getset[] = {
363 {"capacity", (getter) Buffer_capacity_getter, NULL, "", NULL },
364 {"data", (getter) Buffer_data_getter, NULL, "", NULL },
365 {NULL}
366};
367
368static PyType_Slot BufferType_slots[] = {
369 {Py_tp_dealloc, Buffer_dealloc},
370 {Py_tp_methods, Buffer_methods},
371 {Py_tp_doc, "Buffer objects"},
372 {Py_tp_getset, Buffer_getset},
373 {Py_tp_init, Buffer_init},
374 {0, 0},
375};
376
377static PyType_Spec BufferType_spec = {
378 MODULE_NAME ".Buffer",
379 sizeof(BufferObject),
380 0,
381 Py_TPFLAGS_DEFAULT,
382 BufferType_slots
383};
384
385
386static struct PyModuleDef moduledef = {
387 PyModuleDef_HEAD_INIT,
388 MODULE_NAME, /* m_name */
389 "Serialization utilities.", /* m_doc */
390 -1, /* m_size */
391 NULL, /* m_methods */
392 NULL, /* m_reload */
393 NULL, /* m_traverse */
394 NULL, /* m_clear */
395 NULL, /* m_free */
396};
397
398
399PyMODINIT_FUNC
400PyInit__buffer(void)
401{
402 PyObject* m;
403
404 m = PyModule_Create(&moduledef);
405 if (m == NULL)
406 return NULL;
407
408 BufferReadError = PyErr_NewException(MODULE_NAME ".BufferReadError", PyExc_ValueError, NULL);
409 Py_INCREF(BufferReadError);
410 PyModule_AddObject(m, "BufferReadError", BufferReadError);
411
412 BufferWriteError = PyErr_NewException(MODULE_NAME ".BufferWriteError", PyExc_ValueError, NULL);
413 Py_INCREF(BufferWriteError);
414 PyModule_AddObject(m, "BufferWriteError", BufferWriteError);
415
416 BufferType = PyType_FromSpec(&BufferType_spec);
417 if (BufferType == NULL)
418 return NULL;
419 PyModule_AddObject(m, "Buffer", BufferType);
420
421 return m;
422}
423