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 *
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 *
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 *
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
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
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 *
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 *
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 *
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
114 hook_ffree(void *ctx, void *ptr)
115 {
116 PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
117 alloc->free(alloc->ctx, ptr);
118 }
119
120 static 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
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 *
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 *
174 remove_mem_hooks(PyObject *self, PyObject *Py_UNUSED(ignored))
175 {
176 fm_remove_hooks();
177 Py_RETURN_NONE;
178 }
179
180 static PyObject *
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 *
340 test_pyobject_setallocators(PyObject *self, PyObject *Py_UNUSED(ignored))
341 {
342 return test_setallocators(PYMEM_DOMAIN_OBJ);
343 }
344
345 static PyObject *
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 *
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 *
444 test_pymem_getallocatorsname(PyObject *self, PyObject *args)
445 {
446 const char *name = _PyMem_GetCurrentAllocatorName();
447 if (name == NULL) {
448 PyErr_SetString(PyExc_RuntimeError, "cannot get allocators name");
449 return NULL;
450 }
451 return PyUnicode_FromString(name);
452 }
453
454 static PyObject *
455 test_pymem_setrawallocators(PyObject *self, PyObject *Py_UNUSED(ignored))
456 {
457 return test_setallocators(PYMEM_DOMAIN_RAW);
458 }
459
460 static PyObject *
461 test_pymem_setallocators(PyObject *self, PyObject *Py_UNUSED(ignored))
462 {
463 return test_setallocators(PYMEM_DOMAIN_MEM);
464 }
465
466 static PyObject *
467 pyobject_malloc_without_gil(PyObject *self, PyObject *args)
468 {
469 char *buffer;
470
471 /* Deliberate bug to test debug hooks on Python memory allocators:
472 call PyObject_Malloc() without holding the GIL */
473 Py_BEGIN_ALLOW_THREADS
474 buffer = PyObject_Malloc(10);
475 Py_END_ALLOW_THREADS
476
477 PyObject_Free(buffer);
478
479 Py_RETURN_NONE;
480 }
481
482 static PyObject *
483 pymem_buffer_overflow(PyObject *self, PyObject *args)
484 {
485 char *buffer;
486
487 /* Deliberate buffer overflow to check that PyMem_Free() detects
488 the overflow when debug hooks are installed. */
489 buffer = PyMem_Malloc(16);
490 if (buffer == NULL) {
491 PyErr_NoMemory();
492 return NULL;
493 }
494 buffer[16] = 'x';
495 PyMem_Free(buffer);
496
497 Py_RETURN_NONE;
498 }
499
500 static PyObject *
501 pymem_api_misuse(PyObject *self, PyObject *args)
502 {
503 char *buffer;
504
505 /* Deliberate misusage of Python allocators:
506 allococate with PyMem but release with PyMem_Raw. */
507 buffer = PyMem_Malloc(16);
508 PyMem_RawFree(buffer);
509
510 Py_RETURN_NONE;
511 }
512
513 static PyObject *
514 pymem_malloc_without_gil(PyObject *self, PyObject *args)
515 {
516 char *buffer;
517
518 /* Deliberate bug to test debug hooks on Python memory allocators:
519 call PyMem_Malloc() without holding the GIL */
520 Py_BEGIN_ALLOW_THREADS
521 buffer = PyMem_Malloc(10);
522 Py_END_ALLOW_THREADS
523
524 PyMem_Free(buffer);
525
526 Py_RETURN_NONE;
527 }
528
529 static PyObject *
530 test_pyobject_is_freed(const char *test_name, PyObject *op)
531 {
532 if (!_PyObject_IsFreed(op)) {
533 PyErr_SetString(PyExc_AssertionError,
534 "object is not seen as freed");
535 return NULL;
536 }
537 Py_RETURN_NONE;
538 }
539
540 static PyObject *
541 check_pyobject_null_is_freed(PyObject *self, PyObject *Py_UNUSED(args))
542 {
543 PyObject *op = NULL;
544 return test_pyobject_is_freed("check_pyobject_null_is_freed", op);
545 }
546
547
548 static PyObject *
549 check_pyobject_uninitialized_is_freed(PyObject *self,
550 PyObject *Py_UNUSED(args))
551 {
552 PyObject *op = (PyObject *)PyObject_Malloc(sizeof(PyObject));
553 if (op == NULL) {
554 return NULL;
555 }
556 /* Initialize reference count to avoid early crash in ceval or GC */
557 Py_SET_REFCNT(op, 1);
558 /* object fields like ob_type are uninitialized! */
559 return test_pyobject_is_freed("check_pyobject_uninitialized_is_freed", op);
560 }
561
562
563 static PyObject *
564 check_pyobject_forbidden_bytes_is_freed(PyObject *self,
565 PyObject *Py_UNUSED(args))
566 {
567 /* Allocate an incomplete PyObject structure: truncate 'ob_type' field */
568 PyObject *op = (PyObject *)PyObject_Malloc(offsetof(PyObject, ob_type));
569 if (op == NULL) {
570 return NULL;
571 }
572 /* Initialize reference count to avoid early crash in ceval or GC */
573 Py_SET_REFCNT(op, 1);
574 /* ob_type field is after the memory block: part of "forbidden bytes"
575 when using debug hooks on memory allocators! */
576 return test_pyobject_is_freed("check_pyobject_forbidden_bytes_is_freed", op);
577 }
578
579
580 static PyObject *
581 check_pyobject_freed_is_freed(PyObject *self, PyObject *Py_UNUSED(args))
582 {
583 /* This test would fail if run with the address sanitizer */
584 #ifdef _Py_ADDRESS_SANITIZER
585 Py_RETURN_NONE;
586 #else
587 PyObject *op = PyObject_CallNoArgs((PyObject *)&PyBaseObject_Type);
588 if (op == NULL) {
589 return NULL;
590 }
591 Py_TYPE(op)->tp_dealloc(op);
592 /* Reset reference count to avoid early crash in ceval or GC */
593 Py_SET_REFCNT(op, 1);
594 /* object memory is freed! */
595 return test_pyobject_is_freed("check_pyobject_freed_is_freed", op);
596 #endif
597 }
598
599 // Tracemalloc tests
600 static PyObject *
601 tracemalloc_track(PyObject *self, PyObject *args)
602 {
603 unsigned int domain;
604 PyObject *ptr_obj;
605 Py_ssize_t size;
606 int release_gil = 0;
607
608 if (!PyArg_ParseTuple(args, "IOn|i",
609 &domain, &ptr_obj, &size, &release_gil))
610 {
611 return NULL;
612 }
613 void *ptr = PyLong_AsVoidPtr(ptr_obj);
614 if (PyErr_Occurred()) {
615 return NULL;
616 }
617
618 int res;
619 if (release_gil) {
620 Py_BEGIN_ALLOW_THREADS
621 res = PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
622 Py_END_ALLOW_THREADS
623 }
624 else {
625 res = PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
626 }
627 if (res < 0) {
628 PyErr_SetString(PyExc_RuntimeError, "PyTraceMalloc_Track error");
629 return NULL;
630 }
631
632 Py_RETURN_NONE;
633 }
634
635 static PyObject *
636 tracemalloc_untrack(PyObject *self, PyObject *args)
637 {
638 unsigned int domain;
639 PyObject *ptr_obj;
640
641 if (!PyArg_ParseTuple(args, "IO", &domain, &ptr_obj)) {
642 return NULL;
643 }
644 void *ptr = PyLong_AsVoidPtr(ptr_obj);
645 if (PyErr_Occurred()) {
646 return NULL;
647 }
648
649 int res = PyTraceMalloc_Untrack(domain, (uintptr_t)ptr);
650 if (res < 0) {
651 PyErr_SetString(PyExc_RuntimeError, "PyTraceMalloc_Untrack error");
652 return NULL;
653 }
654
655 Py_RETURN_NONE;
656 }
657
658 static PyObject *
659 tracemalloc_get_traceback(PyObject *self, PyObject *args)
660 {
661 unsigned int domain;
662 PyObject *ptr_obj;
663
664 if (!PyArg_ParseTuple(args, "IO", &domain, &ptr_obj)) {
665 return NULL;
666 }
667 void *ptr = PyLong_AsVoidPtr(ptr_obj);
668 if (PyErr_Occurred()) {
669 return NULL;
670 }
671
672 return _PyTraceMalloc_GetTraceback(domain, (uintptr_t)ptr);
673 }
674
675 static PyMethodDef test_methods[] = {
676 {"check_pyobject_forbidden_bytes_is_freed",
677 check_pyobject_forbidden_bytes_is_freed, METH_NOARGS},
678 {"check_pyobject_freed_is_freed", check_pyobject_freed_is_freed, METH_NOARGS},
679 {"check_pyobject_null_is_freed", check_pyobject_null_is_freed, METH_NOARGS},
680 {"check_pyobject_uninitialized_is_freed",
681 check_pyobject_uninitialized_is_freed, METH_NOARGS},
682 {"pymem_api_misuse", pymem_api_misuse, METH_NOARGS},
683 {"pymem_buffer_overflow", pymem_buffer_overflow, METH_NOARGS},
684 {"pymem_getallocatorsname", test_pymem_getallocatorsname, METH_NOARGS},
685 {"pymem_malloc_without_gil", pymem_malloc_without_gil, METH_NOARGS},
686 {"pyobject_malloc_without_gil", pyobject_malloc_without_gil, METH_NOARGS},
687 {"remove_mem_hooks", remove_mem_hooks, METH_NOARGS,
688 PyDoc_STR("Remove memory hooks.")},
689 {"set_nomemory", (PyCFunction)set_nomemory, METH_VARARGS,
690 PyDoc_STR("set_nomemory(start:int, stop:int = 0)")},
691 {"test_pymem_alloc0", test_pymem_alloc0, METH_NOARGS},
692 {"test_pymem_setallocators", test_pymem_setallocators, METH_NOARGS},
693 {"test_pymem_setrawallocators", test_pymem_setrawallocators, METH_NOARGS},
694 {"test_pyobject_new", test_pyobject_new, METH_NOARGS},
695 {"test_pyobject_setallocators", test_pyobject_setallocators, METH_NOARGS},
696
697 // Tracemalloc tests
698 {"tracemalloc_track", tracemalloc_track, METH_VARARGS},
699 {"tracemalloc_untrack", tracemalloc_untrack, METH_VARARGS},
700 {"tracemalloc_get_traceback", tracemalloc_get_traceback, METH_VARARGS},
701 {NULL},
702 };
703
704 int
705 _PyTestCapi_Init_Mem(PyObject *mod)
706 {
707 if (PyModule_AddFunctions(mod, test_methods) < 0) {
708 return -1;
709 }
710
711 PyObject *v;
712 #ifdef WITH_PYMALLOC
713 v = Py_NewRef(Py_True);
714 #else
715 v = Py_NewRef(Py_False);
716 #endif
717 int rc = PyModule_AddObjectRef(mod, "WITH_PYMALLOC", v);
718 Py_DECREF(v);
719 if (rc < 0) {
720 return -1;
721 }
722
723 return 0;
724 }