1 #include "parts.h"
2
3 #include <stddef.h>
4
5
6 typedef struct {
7 PyMemAllocatorEx alloc;
8
9 size_t malloc_size;
10 size_t calloc_nelem;
11 size_t calloc_elsize;
12 void *realloc_ptr;
13 size_t realloc_new_size;
14 void *free_ptr;
15 void *ctx;
16 } alloc_hook_t;
17
18 static void *
hook_malloc(void * ctx,size_t size)19 hook_malloc(void *ctx, size_t size)
20 {
21 alloc_hook_t *hook = (alloc_hook_t *)ctx;
22 hook->ctx = ctx;
23 hook->malloc_size = size;
24 return hook->alloc.malloc(hook->alloc.ctx, size);
25 }
26
27 static void *
hook_calloc(void * ctx,size_t nelem,size_t elsize)28 hook_calloc(void *ctx, size_t nelem, size_t elsize)
29 {
30 alloc_hook_t *hook = (alloc_hook_t *)ctx;
31 hook->ctx = ctx;
32 hook->calloc_nelem = nelem;
33 hook->calloc_elsize = elsize;
34 return hook->alloc.calloc(hook->alloc.ctx, nelem, elsize);
35 }
36
37 static void *
hook_realloc(void * ctx,void * ptr,size_t new_size)38 hook_realloc(void *ctx, void *ptr, size_t new_size)
39 {
40 alloc_hook_t *hook = (alloc_hook_t *)ctx;
41 hook->ctx = ctx;
42 hook->realloc_ptr = ptr;
43 hook->realloc_new_size = new_size;
44 return hook->alloc.realloc(hook->alloc.ctx, ptr, new_size);
45 }
46
47 static void
hook_free(void * ctx,void * ptr)48 hook_free(void *ctx, void *ptr)
49 {
50 alloc_hook_t *hook = (alloc_hook_t *)ctx;
51 hook->ctx = ctx;
52 hook->free_ptr = ptr;
53 hook->alloc.free(hook->alloc.ctx, ptr);
54 }
55
56 /* Most part of the following code is inherited from the pyfailmalloc project
57 * written by Victor Stinner. */
58 static struct {
59 int installed;
60 PyMemAllocatorEx raw;
61 PyMemAllocatorEx mem;
62 PyMemAllocatorEx obj;
63 } FmHook;
64
65 static struct {
66 int start;
67 int stop;
68 Py_ssize_t count;
69 } FmData;
70
71 static int
fm_nomemory(void)72 fm_nomemory(void)
73 {
74 FmData.count++;
75 if (FmData.count > FmData.start &&
76 (FmData.stop <= 0 || FmData.count <= FmData.stop))
77 {
78 return 1;
79 }
80 return 0;
81 }
82
83 static void *
hook_fmalloc(void * ctx,size_t size)84 hook_fmalloc(void *ctx, size_t size)
85 {
86 PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
87 if (fm_nomemory()) {
88 return NULL;
89 }
90 return alloc->malloc(alloc->ctx, size);
91 }
92
93 static void *
hook_fcalloc(void * ctx,size_t nelem,size_t elsize)94 hook_fcalloc(void *ctx, size_t nelem, size_t elsize)
95 {
96 PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
97 if (fm_nomemory()) {
98 return NULL;
99 }
100 return alloc->calloc(alloc->ctx, nelem, elsize);
101 }
102
103 static void *
hook_frealloc(void * ctx,void * ptr,size_t new_size)104 hook_frealloc(void *ctx, void *ptr, size_t new_size)
105 {
106 PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
107 if (fm_nomemory()) {
108 return NULL;
109 }
110 return alloc->realloc(alloc->ctx, ptr, new_size);
111 }
112
113 static void
hook_ffree(void * ctx,void * ptr)114 hook_ffree(void *ctx, void *ptr)
115 {
116 PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
117 alloc->free(alloc->ctx, ptr);
118 }
119
120 static void
fm_setup_hooks(void)121 fm_setup_hooks(void)
122 {
123 if (FmHook.installed) {
124 return;
125 }
126 FmHook.installed = 1;
127
128 PyMemAllocatorEx alloc;
129 alloc.malloc = hook_fmalloc;
130 alloc.calloc = hook_fcalloc;
131 alloc.realloc = hook_frealloc;
132 alloc.free = hook_ffree;
133 PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &FmHook.raw);
134 PyMem_GetAllocator(PYMEM_DOMAIN_MEM, &FmHook.mem);
135 PyMem_GetAllocator(PYMEM_DOMAIN_OBJ, &FmHook.obj);
136
137 alloc.ctx = &FmHook.raw;
138 PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &alloc);
139
140 alloc.ctx = &FmHook.mem;
141 PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &alloc);
142
143 alloc.ctx = &FmHook.obj;
144 PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
145 }
146
147 static void
fm_remove_hooks(void)148 fm_remove_hooks(void)
149 {
150 if (FmHook.installed) {
151 FmHook.installed = 0;
152 PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &FmHook.raw);
153 PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &FmHook.mem);
154 PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &FmHook.obj);
155 }
156 }
157
158 static PyObject *
set_nomemory(PyObject * self,PyObject * args)159 set_nomemory(PyObject *self, PyObject *args)
160 {
161 /* Memory allocation fails after 'start' allocation requests, and until
162 * 'stop' allocation requests except when 'stop' is negative or equal
163 * to 0 (default) in which case allocation failures never stop. */
164 FmData.count = 0;
165 FmData.stop = 0;
166 if (!PyArg_ParseTuple(args, "i|i", &FmData.start, &FmData.stop)) {
167 return NULL;
168 }
169 fm_setup_hooks();
170 Py_RETURN_NONE;
171 }
172
173 static PyObject *
remove_mem_hooks(PyObject * self,PyObject * Py_UNUSED (ignored))174 remove_mem_hooks(PyObject *self, PyObject *Py_UNUSED(ignored))
175 {
176 fm_remove_hooks();
177 Py_RETURN_NONE;
178 }
179
180 static PyObject *
test_setallocators(PyMemAllocatorDomain domain)181 test_setallocators(PyMemAllocatorDomain domain)
182 {
183 PyObject *res = NULL;
184 const char *error_msg;
185 alloc_hook_t hook;
186
187 memset(&hook, 0, sizeof(hook));
188
189 PyMemAllocatorEx alloc;
190 alloc.ctx = &hook;
191 alloc.malloc = &hook_malloc;
192 alloc.calloc = &hook_calloc;
193 alloc.realloc = &hook_realloc;
194 alloc.free = &hook_free;
195 PyMem_GetAllocator(domain, &hook.alloc);
196 PyMem_SetAllocator(domain, &alloc);
197
198 /* malloc, realloc, free */
199 size_t size = 42;
200 hook.ctx = NULL;
201 void *ptr;
202 switch(domain) {
203 case PYMEM_DOMAIN_RAW:
204 ptr = PyMem_RawMalloc(size);
205 break;
206 case PYMEM_DOMAIN_MEM:
207 ptr = PyMem_Malloc(size);
208 break;
209 case PYMEM_DOMAIN_OBJ:
210 ptr = PyObject_Malloc(size);
211 break;
212 default:
213 ptr = NULL;
214 break;
215 }
216
217 #define CHECK_CTX(FUNC) \
218 if (hook.ctx != &hook) { \
219 error_msg = FUNC " wrong context"; \
220 goto fail; \
221 } \
222 hook.ctx = NULL; /* reset for next check */
223
224 if (ptr == NULL) {
225 error_msg = "malloc failed";
226 goto fail;
227 }
228 CHECK_CTX("malloc");
229 if (hook.malloc_size != size) {
230 error_msg = "malloc invalid size";
231 goto fail;
232 }
233
234 size_t size2 = 200;
235 void *ptr2;
236 switch(domain) {
237 case PYMEM_DOMAIN_RAW:
238 ptr2 = PyMem_RawRealloc(ptr, size2);
239 break;
240 case PYMEM_DOMAIN_MEM:
241 ptr2 = PyMem_Realloc(ptr, size2);
242 break;
243 case PYMEM_DOMAIN_OBJ:
244 ptr2 = PyObject_Realloc(ptr, size2);
245 break;
246 default:
247 ptr2 = NULL;
248 break;
249 }
250
251 if (ptr2 == NULL) {
252 error_msg = "realloc failed";
253 goto fail;
254 }
255 CHECK_CTX("realloc");
256 if (hook.realloc_ptr != ptr || hook.realloc_new_size != size2) {
257 error_msg = "realloc invalid parameters";
258 goto fail;
259 }
260
261 switch(domain) {
262 case PYMEM_DOMAIN_RAW:
263 PyMem_RawFree(ptr2);
264 break;
265 case PYMEM_DOMAIN_MEM:
266 PyMem_Free(ptr2);
267 break;
268 case PYMEM_DOMAIN_OBJ:
269 PyObject_Free(ptr2);
270 break;
271 }
272
273 CHECK_CTX("free");
274 if (hook.free_ptr != ptr2) {
275 error_msg = "free invalid pointer";
276 goto fail;
277 }
278
279 /* calloc, free */
280 size_t nelem = 2;
281 size_t elsize = 5;
282 switch(domain) {
283 case PYMEM_DOMAIN_RAW:
284 ptr = PyMem_RawCalloc(nelem, elsize);
285 break;
286 case PYMEM_DOMAIN_MEM:
287 ptr = PyMem_Calloc(nelem, elsize);
288 break;
289 case PYMEM_DOMAIN_OBJ:
290 ptr = PyObject_Calloc(nelem, elsize);
291 break;
292 default:
293 ptr = NULL;
294 break;
295 }
296
297 if (ptr == NULL) {
298 error_msg = "calloc failed";
299 goto fail;
300 }
301 CHECK_CTX("calloc");
302 if (hook.calloc_nelem != nelem || hook.calloc_elsize != elsize) {
303 error_msg = "calloc invalid nelem or elsize";
304 goto fail;
305 }
306
307 hook.free_ptr = NULL;
308 switch(domain) {
309 case PYMEM_DOMAIN_RAW:
310 PyMem_RawFree(ptr);
311 break;
312 case PYMEM_DOMAIN_MEM:
313 PyMem_Free(ptr);
314 break;
315 case PYMEM_DOMAIN_OBJ:
316 PyObject_Free(ptr);
317 break;
318 }
319
320 CHECK_CTX("calloc free");
321 if (hook.free_ptr != ptr) {
322 error_msg = "calloc free invalid pointer";
323 goto fail;
324 }
325
326 res = Py_NewRef(Py_None);
327 goto finally;
328
329 fail:
330 PyErr_SetString(PyExc_RuntimeError, error_msg);
331
332 finally:
333 PyMem_SetAllocator(domain, &hook.alloc);
334 return res;
335
336 #undef CHECK_CTX
337 }
338
339 static PyObject *
test_pyobject_setallocators(PyObject * self,PyObject * Py_UNUSED (ignored))340 test_pyobject_setallocators(PyObject *self, PyObject *Py_UNUSED(ignored))
341 {
342 return test_setallocators(PYMEM_DOMAIN_OBJ);
343 }
344
345 static PyObject *
test_pyobject_new(PyObject * self,PyObject * Py_UNUSED (ignored))346 test_pyobject_new(PyObject *self, PyObject *Py_UNUSED(ignored))
347 {
348 PyObject *obj;
349 PyTypeObject *type = &PyBaseObject_Type;
350 PyTypeObject *var_type = &PyBytes_Type;
351
352 // PyObject_New()
353 obj = PyObject_New(PyObject, type);
354 if (obj == NULL) {
355 goto alloc_failed;
356 }
357 Py_DECREF(obj);
358
359 // PyObject_NEW()
360 obj = PyObject_NEW(PyObject, type);
361 if (obj == NULL) {
362 goto alloc_failed;
363 }
364 Py_DECREF(obj);
365
366 // PyObject_NewVar()
367 obj = PyObject_NewVar(PyObject, var_type, 3);
368 if (obj == NULL) {
369 goto alloc_failed;
370 }
371 Py_DECREF(obj);
372
373 // PyObject_NEW_VAR()
374 obj = PyObject_NEW_VAR(PyObject, var_type, 3);
375 if (obj == NULL) {
376 goto alloc_failed;
377 }
378 Py_DECREF(obj);
379
380 Py_RETURN_NONE;
381
382 alloc_failed:
383 PyErr_NoMemory();
384 return NULL;
385 }
386
387 static PyObject *
test_pymem_alloc0(PyObject * self,PyObject * Py_UNUSED (ignored))388 test_pymem_alloc0(PyObject *self, PyObject *Py_UNUSED(ignored))
389 {
390 void *ptr;
391
392 ptr = PyMem_RawMalloc(0);
393 if (ptr == NULL) {
394 PyErr_SetString(PyExc_RuntimeError,
395 "PyMem_RawMalloc(0) returns NULL");
396 return NULL;
397 }
398 PyMem_RawFree(ptr);
399
400 ptr = PyMem_RawCalloc(0, 0);
401 if (ptr == NULL) {
402 PyErr_SetString(PyExc_RuntimeError,
403 "PyMem_RawCalloc(0, 0) returns NULL");
404 return NULL;
405 }
406 PyMem_RawFree(ptr);
407
408 ptr = PyMem_Malloc(0);
409 if (ptr == NULL) {
410 PyErr_SetString(PyExc_RuntimeError,
411 "PyMem_Malloc(0) returns NULL");
412 return NULL;
413 }
414 PyMem_Free(ptr);
415
416 ptr = PyMem_Calloc(0, 0);
417 if (ptr == NULL) {
418 PyErr_SetString(PyExc_RuntimeError,
419 "PyMem_Calloc(0, 0) returns NULL");
420 return NULL;
421 }
422 PyMem_Free(ptr);
423
424 ptr = PyObject_Malloc(0);
425 if (ptr == NULL) {
426 PyErr_SetString(PyExc_RuntimeError,
427 "PyObject_Malloc(0) returns NULL");
428 return NULL;
429 }
430 PyObject_Free(ptr);
431
432 ptr = PyObject_Calloc(0, 0);
433 if (ptr == NULL) {
434 PyErr_SetString(PyExc_RuntimeError,
435 "PyObject_Calloc(0, 0) returns NULL");
436 return NULL;
437 }
438 PyObject_Free(ptr);
439
440 Py_RETURN_NONE;
441 }
442
443 static PyObject *
test_pymem_setrawallocators(PyObject * self,PyObject * Py_UNUSED (ignored))444 test_pymem_setrawallocators(PyObject *self, PyObject *Py_UNUSED(ignored))
445 {
446 return test_setallocators(PYMEM_DOMAIN_RAW);
447 }
448
449 static PyObject *
test_pymem_setallocators(PyObject * self,PyObject * Py_UNUSED (ignored))450 test_pymem_setallocators(PyObject *self, PyObject *Py_UNUSED(ignored))
451 {
452 return test_setallocators(PYMEM_DOMAIN_MEM);
453 }
454
455 static PyObject *
pyobject_malloc_without_gil(PyObject * self,PyObject * args)456 pyobject_malloc_without_gil(PyObject *self, PyObject *args)
457 {
458 char *buffer;
459
460 /* Deliberate bug to test debug hooks on Python memory allocators:
461 call PyObject_Malloc() without holding the GIL */
462 Py_BEGIN_ALLOW_THREADS
463 buffer = PyObject_Malloc(10);
464 Py_END_ALLOW_THREADS
465
466 PyObject_Free(buffer);
467
468 Py_RETURN_NONE;
469 }
470
471 static PyObject *
pymem_buffer_overflow(PyObject * self,PyObject * args)472 pymem_buffer_overflow(PyObject *self, PyObject *args)
473 {
474 char *buffer;
475
476 /* Deliberate buffer overflow to check that PyMem_Free() detects
477 the overflow when debug hooks are installed. */
478 buffer = PyMem_Malloc(16);
479 if (buffer == NULL) {
480 PyErr_NoMemory();
481 return NULL;
482 }
483 buffer[16] = 'x';
484 PyMem_Free(buffer);
485
486 Py_RETURN_NONE;
487 }
488
489 static PyObject *
pymem_api_misuse(PyObject * self,PyObject * args)490 pymem_api_misuse(PyObject *self, PyObject *args)
491 {
492 char *buffer;
493
494 /* Deliberate misusage of Python allocators:
495 allococate with PyMem but release with PyMem_Raw. */
496 buffer = PyMem_Malloc(16);
497 PyMem_RawFree(buffer);
498
499 Py_RETURN_NONE;
500 }
501
502 static PyObject *
pymem_malloc_without_gil(PyObject * self,PyObject * args)503 pymem_malloc_without_gil(PyObject *self, PyObject *args)
504 {
505 char *buffer;
506
507 /* Deliberate bug to test debug hooks on Python memory allocators:
508 call PyMem_Malloc() without holding the GIL */
509 Py_BEGIN_ALLOW_THREADS
510 buffer = PyMem_Malloc(10);
511 Py_END_ALLOW_THREADS
512
513 PyMem_Free(buffer);
514
515 Py_RETURN_NONE;
516 }
517
518
519 // Tracemalloc tests
520 static PyObject *
tracemalloc_track(PyObject * self,PyObject * args)521 tracemalloc_track(PyObject *self, PyObject *args)
522 {
523 unsigned int domain;
524 PyObject *ptr_obj;
525 Py_ssize_t size;
526 int release_gil = 0;
527
528 if (!PyArg_ParseTuple(args, "IOn|i",
529 &domain, &ptr_obj, &size, &release_gil))
530 {
531 return NULL;
532 }
533 void *ptr = PyLong_AsVoidPtr(ptr_obj);
534 if (PyErr_Occurred()) {
535 return NULL;
536 }
537
538 int res;
539 if (release_gil) {
540 Py_BEGIN_ALLOW_THREADS
541 res = PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
542 Py_END_ALLOW_THREADS
543 }
544 else {
545 res = PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
546 }
547 if (res < 0) {
548 PyErr_SetString(PyExc_RuntimeError, "PyTraceMalloc_Track error");
549 return NULL;
550 }
551
552 Py_RETURN_NONE;
553 }
554
555 static PyObject *
tracemalloc_untrack(PyObject * self,PyObject * args)556 tracemalloc_untrack(PyObject *self, PyObject *args)
557 {
558 unsigned int domain;
559 PyObject *ptr_obj;
560
561 if (!PyArg_ParseTuple(args, "IO", &domain, &ptr_obj)) {
562 return NULL;
563 }
564 void *ptr = PyLong_AsVoidPtr(ptr_obj);
565 if (PyErr_Occurred()) {
566 return NULL;
567 }
568
569 int res = PyTraceMalloc_Untrack(domain, (uintptr_t)ptr);
570 if (res < 0) {
571 PyErr_SetString(PyExc_RuntimeError, "PyTraceMalloc_Untrack error");
572 return NULL;
573 }
574
575 Py_RETURN_NONE;
576 }
577
578 static PyMethodDef test_methods[] = {
579 {"pymem_api_misuse", pymem_api_misuse, METH_NOARGS},
580 {"pymem_buffer_overflow", pymem_buffer_overflow, METH_NOARGS},
581 {"pymem_malloc_without_gil", pymem_malloc_without_gil, METH_NOARGS},
582 {"pyobject_malloc_without_gil", pyobject_malloc_without_gil, METH_NOARGS},
583 {"remove_mem_hooks", remove_mem_hooks, METH_NOARGS,
584 PyDoc_STR("Remove memory hooks.")},
585 {"set_nomemory", (PyCFunction)set_nomemory, METH_VARARGS,
586 PyDoc_STR("set_nomemory(start:int, stop:int = 0)")},
587 {"test_pymem_alloc0", test_pymem_alloc0, METH_NOARGS},
588 {"test_pymem_setallocators", test_pymem_setallocators, METH_NOARGS},
589 {"test_pymem_setrawallocators", test_pymem_setrawallocators, METH_NOARGS},
590 {"test_pyobject_new", test_pyobject_new, METH_NOARGS},
591 {"test_pyobject_setallocators", test_pyobject_setallocators, METH_NOARGS},
592
593 // Tracemalloc tests
594 {"tracemalloc_track", tracemalloc_track, METH_VARARGS},
595 {"tracemalloc_untrack", tracemalloc_untrack, METH_VARARGS},
596 {NULL},
597 };
598
599 int
_PyTestCapi_Init_Mem(PyObject * mod)600 _PyTestCapi_Init_Mem(PyObject *mod)
601 {
602 if (PyModule_AddFunctions(mod, test_methods) < 0) {
603 return -1;
604 }
605
606 PyObject *v;
607 #ifdef WITH_PYMALLOC
608 v = Py_True;
609 #else
610 v = Py_False;
611 #endif
612 if (PyModule_AddObjectRef(mod, "WITH_PYMALLOC", v) < 0) {
613 return -1;
614 }
615
616 #ifdef WITH_MIMALLOC
617 v = Py_True;
618 #else
619 v = Py_False;
620 #endif
621 if (PyModule_AddObjectRef(mod, "WITH_MIMALLOC", v) < 0) {
622 return -1;
623 }
624
625 return 0;
626 }
627