1 #include <stdarg.h>
2 #include <stddef.h>
3 #include <setjmp.h>
4 #include <assert.h>
5 #include <cmocka.h>
6 /* cmocka < 1.0 didn't support these features we need */
7 #ifndef assert_ptr_equal
8 #define assert_ptr_equal(a, b) \
9 _assert_int_equal(cast_ptr_to_largest_integral_type(a), \
10 cast_ptr_to_largest_integral_type(b), \
11 __FILE__, __LINE__)
12 #define CMUnitTest UnitTest
13 #define cmocka_unit_test unit_test
14 #define cmocka_run_group_tests(t, setup, teardown) run_tests(t)
15 #endif
16
17
18 extern void mock_assert(const int result, const char* const expression,
19 const char * const file, const int line);
20 #undef assert
21 #define assert(expression) \
22 mock_assert((int)(expression), #expression, __FILE__, __LINE__);
23
24 #include "afl-prealloc.h"
25
26 /* remap exit -> assert, then use cmocka's mock_assert
27 (compile with `--wrap=exit`) */
28 extern void exit(int status);
29 extern void __real_exit(int status);
30 void __wrap_exit(int status);
__wrap_exit(int status)31 void __wrap_exit(int status) {
32 (void)status;
33 assert(0);
34 }
35
36 /* ignore all printfs */
37 #undef printf
38 extern int printf(const char *format, ...);
39 extern int __real_printf(const char *format, ...);
40 //int __wrap_printf(const char *format, ...);
__wrap_printf(const char * format,...)41 int __wrap_printf(const char *format, ...) {
42 (void)format;
43 return 1;
44 }
45
46 typedef struct prealloc_me
47 {
48 PREALLOCABLE;
49
50 u8 *content[128];
51
52 } element_t;
53
54 #define PREALLOCED_BUF_SIZE (64)
55 element_t prealloc_me_buf[PREALLOCED_BUF_SIZE];
56 s32 prealloc_me_size = 0;
57
test_alloc_free(void ** state)58 static void test_alloc_free(void **state) {
59 (void)state;
60
61 element_t *prealloced = NULL;
62 PRE_ALLOC(prealloced, prealloc_me_buf, PREALLOCED_BUF_SIZE, prealloc_me_size);
63 assert_non_null(prealloced);
64 PRE_FREE(prealloced, prealloc_me_size);
65
66 }
67
test_prealloc_overflow(void ** state)68 static void test_prealloc_overflow(void **state) {
69 (void)state;
70
71 u32 i = 0;
72 element_t *prealloced[PREALLOCED_BUF_SIZE + 10];
73
74 for (i = 0; i < PREALLOCED_BUF_SIZE + 10; i++) {
75
76 PRE_ALLOC(prealloced[i], prealloc_me_buf, PREALLOCED_BUF_SIZE, prealloc_me_size);
77 assert_non_null(prealloced[i]);
78
79 }
80 assert_int_equal(prealloced[0]->pre_status, PRE_STATUS_USED);
81 assert_int_equal(prealloced[PREALLOCED_BUF_SIZE]->pre_status, PRE_STATUS_MALLOC);
82
83 PRE_FREE(prealloced[20], prealloc_me_size);
84 PRE_ALLOC(prealloced[20], prealloc_me_buf, PREALLOCED_BUF_SIZE, prealloc_me_size);
85 assert_non_null(prealloced[20]);
86 assert_int_equal(prealloced[20]->pre_status, PRE_STATUS_USED);
87
88 PRE_FREE(prealloced[PREALLOCED_BUF_SIZE], prealloc_me_size);
89 PRE_FREE(prealloced[0], prealloc_me_size);
90 PRE_ALLOC(prealloced[PREALLOCED_BUF_SIZE], prealloc_me_buf, PREALLOCED_BUF_SIZE, prealloc_me_size);
91 assert_non_null(prealloced[PREALLOCED_BUF_SIZE]);
92 /* there should be space now! */
93 assert_int_equal(prealloced[PREALLOCED_BUF_SIZE]->pre_status, PRE_STATUS_USED);
94
95 PRE_ALLOC(prealloced[0], prealloc_me_buf, PREALLOCED_BUF_SIZE, prealloc_me_size);
96 assert_non_null(prealloced[0]);
97 /* no more space */
98 assert_int_equal(prealloced[0]->pre_status, PRE_STATUS_MALLOC);
99
100 for (i = 0; i < PREALLOCED_BUF_SIZE + 10; i++) {
101
102 PRE_FREE(prealloced[i], prealloc_me_size);
103
104 }
105
106 }
107
main(int argc,char ** argv)108 int main(int argc, char **argv) {
109 (void)argc;
110 (void)argv;
111
112 const struct CMUnitTest tests[] = {
113 cmocka_unit_test(test_alloc_free),
114 cmocka_unit_test(test_prealloc_overflow),
115 };
116
117 //return cmocka_run_group_tests (tests, setup, teardown);
118 __real_exit( cmocka_run_group_tests (tests, NULL, NULL) );
119
120 // fake return for dumb compilers
121 return 0;
122 }
123