1 #include <stdlib.h>
2 #include <stdio.h>
3 #include <assert.h>
4 #include <string.h>
5 #include <stdint.h>
6 
7 #include <mimalloc.h>
8 #include <mimalloc-override.h>  // redefines malloc etc.
9 
10 
11 #include <stdint.h>
12 #include <stdbool.h>
13 
14 #define MI_INTPTR_SIZE 8
15 #define MI_LARGE_WSIZE_MAX (4*1024*1024 / MI_INTPTR_SIZE)
16 
17 #define MI_BIN_HUGE 100
18 //#define MI_ALIGN2W
19 
20 // Bit scan reverse: return the index of the highest bit.
21 static inline uint8_t mi_bsr32(uint32_t x);
22 
23 #if defined(_MSC_VER)
24 #include <windows.h>
25 #include <intrin.h>
mi_bsr32(uint32_t x)26 static inline uint8_t mi_bsr32(uint32_t x) {
27   uint32_t idx;
28   _BitScanReverse((DWORD*)&idx, x);
29   return idx;
30 }
31 #elif defined(__GNUC__) || defined(__clang__)
mi_bsr32(uint32_t x)32 static inline uint8_t mi_bsr32(uint32_t x) {
33   return (31 - __builtin_clz(x));
34 }
35 #else
mi_bsr32(uint32_t x)36 static inline uint8_t mi_bsr32(uint32_t x) {
37   // de Bruijn multiplication, see <http://supertech.csail.mit.edu/papers/debruijn.pdf>
38   static const uint8_t debruijn[32] = {
39      31,  0, 22,  1, 28, 23, 18,  2, 29, 26, 24, 10, 19,  7,  3, 12,
40      30, 21, 27, 17, 25,  9,  6, 11, 20, 16,  8,  5, 15,  4, 14, 13,
41   };
42   x |= x >> 1;
43   x |= x >> 2;
44   x |= x >> 4;
45   x |= x >> 8;
46   x |= x >> 16;
47   x++;
48   return debruijn[(x*0x076be629) >> 27];
49 }
50 #endif
51 
52 /*
53 // Bit scan reverse: return the index of the highest bit.
54 uint8_t _mi_bsr(uintptr_t x) {
55   if (x == 0) return 0;
56   #if MI_INTPTR_SIZE==8
57   uint32_t hi = (x >> 32);
58   return (hi == 0 ? mi_bsr32((uint32_t)x) : 32 + mi_bsr32(hi));
59   #elif MI_INTPTR_SIZE==4
60   return mi_bsr32(x);
61   #else
62   # error "define bsr for non-32 or 64-bit platforms"
63   #endif
64 }
65 */
66 
67 
_mi_wsize_from_size(size_t size)68 static inline size_t _mi_wsize_from_size(size_t size) {
69   return (size + sizeof(uintptr_t) - 1) / sizeof(uintptr_t);
70 }
71 
72 // Return the bin for a given field size.
73 // Returns MI_BIN_HUGE if the size is too large.
74 // We use `wsize` for the size in "machine word sizes",
75 // i.e. byte size == `wsize*sizeof(void*)`.
_mi_bin8(size_t size)76 extern inline uint8_t _mi_bin8(size_t size) {
77   size_t wsize = _mi_wsize_from_size(size);
78   uint8_t bin;
79   if (wsize <= 1) {
80     bin = 1;
81   }
82   #if defined(MI_ALIGN4W)
83   else if (wsize <= 4) {
84     bin = (uint8_t)((wsize+1)&~1); // round to double word sizes
85   }
86   #elif defined(MI_ALIGN2W)
87   else if (wsize <= 8) {
88     bin = (uint8_t)((wsize+1)&~1); // round to double word sizes
89   }
90   #else
91   else if (wsize <= 8) {
92     bin = (uint8_t)wsize;
93   }
94   #endif
95   else if (wsize > MI_LARGE_WSIZE_MAX) {
96     bin = MI_BIN_HUGE;
97   }
98   else {
99     #if defined(MI_ALIGN4W)
100     if (wsize <= 16) { wsize = (wsize+3)&~3; } // round to 4x word sizes
101     #endif
102     wsize--;
103     // find the highest bit
104     uint8_t b = mi_bsr32((uint32_t)wsize);
105     // and use the top 3 bits to determine the bin (~12.5% worst internal fragmentation).
106     // - adjust with 3 because we use do not round the first 8 sizes
107     //   which each get an exact bin
108     bin = ((b << 2) + (uint8_t)((wsize >> (b - 2)) & 0x03)) - 3;
109   }
110   return bin;
111 }
112 
_mi_bin4(size_t size)113 extern inline uint8_t _mi_bin4(size_t size) {
114   size_t wsize = _mi_wsize_from_size(size);
115   uint8_t bin;
116   if (wsize <= 1) {
117     bin = 1;
118   }
119   #if defined(MI_ALIGN4W)
120   else if (wsize <= 4) {
121     bin = (uint8_t)((wsize+1)&~1); // round to double word sizes
122   }
123   #elif defined(MI_ALIGN2W)
124   else if (wsize <= 8) {
125     bin = (uint8_t)((wsize+1)&~1); // round to double word sizes
126   }
127   #else
128   else if (wsize <= 8) {
129     bin = (uint8_t)wsize;
130   }
131   #endif
132   else if (wsize > MI_LARGE_WSIZE_MAX) {
133     bin = MI_BIN_HUGE;
134   }
135   else {
136     uint8_t b = mi_bsr32((uint32_t)wsize);
137     bin = ((b << 1) + (uint8_t)((wsize >> (b - 1)) & 0x01)) + 3;
138   }
139   return bin;
140 }
141 
_mi_binx4(size_t bsize)142 size_t _mi_binx4(size_t bsize) {
143   if (bsize==0) return 0;
144   uint8_t b = mi_bsr32((uint32_t)bsize);
145   if (b <= 1) return bsize;
146   size_t bin =  ((b << 1) | (bsize >> (b - 1))&0x01);
147   return bin;
148 }
149 
_mi_binx8(size_t bsize)150 size_t _mi_binx8(size_t bsize) {
151   if (bsize<=1) return bsize;
152   uint8_t b = mi_bsr32((uint32_t)bsize);
153   if (b <= 2) return bsize;
154   size_t bin = ((b << 2) | (bsize >> (b - 2))&0x03) - 5;
155   return bin;
156 }
157 
mi_bins()158 void mi_bins() {
159   //printf("  QNULL(1), /* 0 */ \\\n  ");
160   size_t last_bin = 0;
161   size_t min_bsize = 0;
162   size_t last_bsize = 0;
163   for (size_t bsize = 1; bsize < 2*1024; bsize++) {
164     size_t size = bsize * 64 * 1024;
165     size_t bin = _mi_binx8(bsize);
166     if (bin != last_bin) {
167       printf("min bsize: %6zd, max bsize: %6zd, bin: %6zd\n", min_bsize, last_bsize, last_bin);
168       //printf("QNULL(%6zd), ", wsize);
169       //if (last_bin%8 == 0) printf("/* %i */ \\\n  ", last_bin);
170       last_bin = bin;
171       min_bsize = bsize;
172     }
173     last_bsize = bsize;
174   }
175 }
176 
177 static void double_free1();
178 static void double_free2();
179 static void corrupt_free();
180 static void block_overflow1();
181 static void invalid_free();
182 static void test_aslr(void);
183 static void test_process_info(void);
184 static void test_reserved(void);
185 static void negative_stat(void);
186 
187 
main()188 int main() {
189   mi_version();
190   mi_stats_reset();
191   // detect double frees and heap corruption
192   // double_free1();
193   // double_free2();
194   // corrupt_free();
195   // block_overflow1();
196   // test_aslr();
197   // invalid_free();
198   // test_reserved();
199   // negative_stat();
200 
201   void* p1 = malloc(78);
202   void* p2 = malloc(24);
203   free(p1);
204   p1 = mi_malloc(8);
205   char* s = strdup("hello\n");
206   free(p2);
207 
208   p2 = malloc(16);
209   p1 = realloc(p1, 32);
210   free(p1);
211   free(p2);
212   free(s);
213 
214   /* now test if override worked by allocating/freeing across the api's*/
215   //p1 = mi_malloc(32);
216   //free(p1);
217   //p2 = malloc(32);
218   //mi_free(p2);
219   mi_collect(true);
220   mi_stats_print(NULL);
221   // test_process_info();
222   return 0;
223 }
224 
invalid_free()225 static void invalid_free() {
226   free((void*)0xBADBEEF);
227   realloc((void*)0xBADBEEF,10);
228 }
229 
block_overflow1()230 static void block_overflow1() {
231   uint8_t* p = (uint8_t*)mi_malloc(17);
232   p[18] = 0;
233   free(p);
234 }
235 
236 // The double free samples come ArcHeap [1] by Insu Yun (issue #161)
237 // [1]: https://arxiv.org/pdf/1903.00503.pdf
238 
double_free1()239 static void double_free1() {
240   void* p[256];
241   //uintptr_t buf[256];
242 
243   p[0] = mi_malloc(622616);
244   p[1] = mi_malloc(655362);
245   p[2] = mi_malloc(786432);
246   mi_free(p[2]);
247   // [VULN] Double free
248   mi_free(p[2]);
249   p[3] = mi_malloc(786456);
250   // [BUG] Found overlap
251   // p[3]=0x429b2ea2000 (size=917504), p[1]=0x429b2e42000 (size=786432)
252   fprintf(stderr, "p3: %p-%p, p1: %p-%p, p2: %p\n", p[3], (uint8_t*)(p[3]) + 786456, p[1], (uint8_t*)(p[1]) + 655362, p[2]);
253 }
254 
double_free2()255 static void double_free2() {
256   void* p[256];
257   //uintptr_t buf[256];
258   // [INFO] Command buffer: 0x327b2000
259   // [INFO] Input size: 182
260   p[0] = malloc(712352);
261   p[1] = malloc(786432);
262   free(p[0]);
263   // [VULN] Double free
264   free(p[0]);
265   p[2] = malloc(786440);
266   p[3] = malloc(917504);
267   p[4] = malloc(786440);
268   // [BUG] Found overlap
269   // p[4]=0x433f1402000 (size=917504), p[1]=0x433f14c2000 (size=786432)
270   fprintf(stderr, "p1: %p-%p, p2: %p-%p\n", p[4], (uint8_t*)(p[4]) + 917504, p[1], (uint8_t*)(p[1]) + 786432);
271 }
272 
273 
274 // Try to corrupt the heap through buffer overflow
275 #define N   256
276 #define SZ  64
277 
corrupt_free()278 static void corrupt_free() {
279   void* p[N];
280   // allocate
281   for (int i = 0; i < N; i++) {
282     p[i] = malloc(SZ);
283   }
284   // free some
285   for (int i = 0; i < N; i += (N/10)) {
286     free(p[i]);
287     p[i] = NULL;
288   }
289   // try to corrupt the free list
290   for (int i = 0; i < N; i++) {
291     if (p[i] != NULL) {
292       memset(p[i], 0, SZ+8);
293     }
294   }
295   // allocate more.. trying to trigger an allocation from a corrupted entry
296   // this may need many allocations to get there (if at all)
297   for (int i = 0; i < 4096; i++) {
298     malloc(SZ);
299   }
300 }
301 
test_aslr(void)302 static void test_aslr(void) {
303   void* p[256];
304   p[0] = malloc(378200);
305   p[1] = malloc(1134626);
306   printf("p1: %p, p2: %p\n", p[0], p[1]);
307 }
308 
test_process_info(void)309 static void test_process_info(void) {
310   size_t elapsed = 0;
311   size_t user_msecs = 0;
312   size_t system_msecs = 0;
313   size_t current_rss = 0;
314   size_t peak_rss = 0;
315   size_t current_commit = 0;
316   size_t peak_commit = 0;
317   size_t page_faults = 0;
318   for (int i = 0; i < 100000; i++) {
319     void* p = calloc(100,10);
320     free(p);
321   }
322   mi_process_info(&elapsed, &user_msecs, &system_msecs, &current_rss, &peak_rss, &current_commit, &peak_commit, &page_faults);
323   printf("\n\n*** process info: elapsed %3zd.%03zd s, user: %3zd.%03zd s, rss: %zd b, commit: %zd b\n\n", elapsed/1000, elapsed%1000, user_msecs/1000, user_msecs%1000, peak_rss, peak_commit);
324 }
325 
test_reserved(void)326 static void test_reserved(void) {
327 #define KiB 1024ULL
328 #define MiB (KiB*KiB)
329 #define GiB (MiB*KiB)
330   mi_reserve_os_memory(4*GiB, false, true);
331   void* p1 = malloc(100);
332   void* p2 = malloc(100000);
333   void* p3 = malloc(2*GiB);
334   void* p4 = malloc(1*GiB + 100000);
335   free(p1);
336   free(p2);
337   free(p3);
338   p3 = malloc(1*GiB);
339   free(p4);
340 }
341 
342 
343 
negative_stat(void)344 static void negative_stat(void) {
345   int* p = mi_malloc(60000);
346   mi_stats_print_out(NULL, NULL);
347   *p = 100;
348   mi_free(p);
349   mi_stats_print_out(NULL, NULL);
350 }