1/*
2 * Copyright (c) 2024 Huawei Device Co., Ltd.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 * http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15#include "dfx_allocator.h"
16
17#include <fcntl.h>
18#include <securec.h>
19#include <stddef.h>
20#include <stdint.h>
21#include <stdlib.h>
22#include <unistd.h>
23#include <sys/cdefs.h>
24#include <sys/mman.h>
25#include <sys/prctl.h>
26#include <sys/stat.h>
27#include <sys/user.h>
28
29#include "dfx_define.h"
30#include "dfx_log.h"
31#include "musl_malloc_dispatch_table.h"
32#define HOOK_ENABLE
33#include "musl_preinit_common.h"
34
35#define ALIGN_SIZE 16
36
37#define DFX_PAGE_SIZE 4096
38#define DFX_PAGE_MASK (~(DFX_PAGE_SIZE - 1))
39
40#define DFX_MEMPOOL_MIN_TYPE 4
41#define DFX_MEMPOOL_MAX_TYPE 10
42#define DFX_MEMPOOL_MAX_BLOCK_SIZE 1024
43#define DFX_MMAP_TYPE 111
44
45static struct MallocDispatchType g_dfxCustomMallocDispatch;
46static const size_t PAGE_INFO_SIZE = ((sizeof(PageInfo) + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1));
47static const char DFX_MEM_PAGE_SIGN[DFX_MEMPOOL_TAG_SIZE] = {'D', 'F', 'X', 1};
48static DfxAllocator g_dfxAllocator = {
49    .initFlag = 0,
50    .pageList = NULL,
51};
52
53static inline uintptr_t PageStart(uintptr_t addr)
54{
55    return (addr & DFX_PAGE_MASK);
56}
57
58static inline uintptr_t PageEnd(uintptr_t addr)
59{
60    return PageStart(addr + (DFX_PAGE_SIZE - 1));
61}
62
63static inline size_t AlignRoundUp(size_t val, size_t align)
64{
65    size_t size = align;
66    if (size == 0) {
67        size = 1;
68    }
69    return ((val + size - 1) & ~(size - 1));
70}
71
72static void AddPage(PageInfo** pageList, PageInfo* page)
73{
74    if (pageList == NULL || page == NULL) {
75        return;
76    }
77    page->next = *pageList;
78    page->prev = NULL;
79    if (*pageList) {
80        (*pageList)->prev = page;
81    }
82    *pageList = page;
83}
84
85static void RemovePage(PageInfo** pageList, PageInfo* page)
86{
87    if (pageList == NULL || page == NULL) {
88        return;
89    }
90    if (page->prev) {
91        page->prev->next = page->next;
92    }
93    if (page->next) {
94        page->next->prev = page->prev;
95    }
96    if (*pageList == page) {
97        *pageList = page->next;
98    }
99    page->prev = NULL;
100    page->next = NULL;
101}
102
103static void MempoolAddPage(DfxMempool* mempool, PageInfo* page)
104{
105    if (mempool == NULL) {
106        DFXLOGE("MempoolAddPage Invalid mempool!");
107        return;
108    }
109    return AddPage(&(mempool->pageList), page);
110}
111
112static void MempoolRemovePage(DfxMempool* mempool, PageInfo* page)
113{
114    if (mempool == NULL) {
115        DFXLOGE("MempoolRemovePage Invalid mempool!");
116        return;
117    }
118    return RemovePage(&(mempool->pageList), page);
119}
120
121static void MempoolAllocPage(DfxMempool* mempool)
122{
123    void* mptr = mmap(NULL, DFX_PAGE_SIZE, PROT_READ | PROT_WRITE,
124        MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
125    if (mptr == MAP_FAILED) {
126        DFXLOGE("Mempool AllocPage mmap failed!");
127        return;
128    }
129    PageInfo* page = (PageInfo*)(mptr);
130    // fill PageInfo
131    if (memcpy_s(page->tag.tagInfo, sizeof(page->tag.tagInfo),
132        DFX_MEM_PAGE_SIGN, sizeof(DFX_MEM_PAGE_SIGN)) != EOK) {
133        munmap(mptr, DFX_PAGE_SIZE);
134        DFXLOGE("Mempool AllocPage fill tag failed!");
135        return;
136    }
137    page->tag.type = mempool->type;
138    page->tag.mempool = mempool;
139    page->freeBlocksCnt = mempool->blocksPerPage;
140    // Page firstblockaddr = page start addr + PageInfo size round up by blocksize
141    uintptr_t firstBlockAddr = AlignRoundUp((uintptr_t)(page + 1), mempool->blockSize);
142    BlockInfo* firstBlock = (BlockInfo*)(firstBlockAddr);
143    // here first block is not alloced block，freeblockList point it
144    firstBlock->freeBlocksCnt = mempool->blocksPerPage;
145    page->freeBlockList = firstBlock;
146    MempoolAddPage(mempool, page);
147    mempool->freePagesCnt++;
148}
149
150static void MempoolFreePage(DfxMempool* mempool, PageInfo* page)
151{
152    if (page->freeBlocksCnt != mempool->blocksPerPage) {
153        DFXLOGE("MempoolFreePage Invalid Page free cnt!");
154        return;
155    }
156    MempoolRemovePage(mempool, page);
157    munmap(page, DFX_PAGE_SIZE);
158    mempool->freePagesCnt--;
159}
160
161static void* MempoolAlloc(DfxMempool* mempool)
162{
163    if (mempool == NULL || mempool->blockSize == 0) {
164        return NULL;
165    }
166    // first alloc memory block, or all page memory blocks have been allocated
167    if (mempool->pageList == NULL) {
168        MempoolAllocPage(mempool);
169    }
170    PageInfo* page = mempool->pageList;
171    if (page == NULL || page->freeBlockList == NULL) {
172        DFXLOGE("MempoolAlloc Alloc Page Failed or Invalid blocklist!");
173        return NULL;
174    }
175    BlockInfo* block = page->freeBlockList;
176    if (block->freeBlocksCnt > 1) {
177        // freeBlocksCnt > 1 stand for page have free block's room
178        // when a block will be allocated,
179        // freeblocklist have no node, need alloc a new block in page room ,then add to freeblocklist
180        BlockInfo* nextfree = (BlockInfo*)((uint8_t*)(block) + mempool->blockSize);
181        nextfree->next = block->next;
182        nextfree->freeBlocksCnt = block->freeBlocksCnt - 1;
183        page->freeBlockList = nextfree;
184    } else {
185        // last freed block in freeblocklist or the last one block in page
186        page->freeBlockList = block->next;
187    }
188
189    // new alloc page, free cnt -1 when used, free cnt +1 when allocated
190    if (page->freeBlocksCnt == mempool->blocksPerPage) {
191        mempool->freePagesCnt--;
192    }
193    if (page->freeBlocksCnt > 0) {
194        page->freeBlocksCnt--;
195    }
196    (void)memset_s(block, mempool->blockSize, 0, mempool->blockSize);
197    // when page's blocks all allocated, remove from pagelist but not free
198    // then pagelist will be point page which have unused blocks
199    // or pagelist be NULL will alloc new page
200    if (page->freeBlocksCnt == 0) {
201        MempoolRemovePage(mempool, page);
202    }
203    return (void*)block;
204}
205
206static void MempoolFree(DfxMempool* mempool, void* ptr)
207{
208    PageInfo * const page = (PageInfo*)(PageStart((uintptr_t)(ptr)));
209
210    if (mempool == NULL || ptr == NULL || mempool->blockSize == 0 ||
211        ((uintptr_t)(ptr)) % (mempool->blockSize) != 0) {
212        DFXLOGE("MempoolFree Invalid mempool or address!");
213        return;
214    }
215    // find ptr's page,and page's freeblocklist
216    (void)memset_s(ptr, mempool->blockSize, 0, mempool->blockSize);
217    BlockInfo* block = (BlockInfo*)(ptr);
218    block->next = page->freeBlockList;
219    block->freeBlocksCnt = 1;
220    page->freeBlockList = block;
221    page->freeBlocksCnt++;
222
223    // all page's blocks have been freed, unmap the page
224    // page have only one block be freed, need add page to pageList
225    if (page->freeBlocksCnt == mempool->blocksPerPage) {
226        mempool->freePagesCnt++;
227        if (mempool->freePagesCnt >= 1) {
228            MempoolFreePage(mempool, page);
229        }
230    } else if (page->freeBlocksCnt == 1) {
231        MempoolAddPage(mempool, page);
232    }
233}
234
235static inline uint32_t SelectMempoolType(size_t num)
236{
237    uint32_t res = 0;
238    size_t n = num - 1;
239    // alloc size（1~1024）, use diffrent block size mempool
240    // The 16-byte size range is the smallest.
241    // Each subsequent level is twice that of the previous level. The maximum is 1024 bytes
242    while (n != 0) {
243        res++;
244        n >>= 1;
245    }
246    return res;
247}
248
249static void InitDfxAllocator(void)
250{
251    for (uint32_t i = 0; i < DFX_MEMPOOLS_NUM; i++) {
252        g_dfxAllocator.dfxMempoolBuf[i].type = i + DFX_MEMPOOL_MIN_TYPE;
253        g_dfxAllocator.dfxMempoolBuf[i].blockSize = (1UL << g_dfxAllocator.dfxMempoolBuf[i].type);
254        g_dfxAllocator.dfxMempoolBuf[i].blocksPerPage =
255            (DFX_PAGE_SIZE - sizeof(PageInfo)) / (g_dfxAllocator.dfxMempoolBuf[i].blockSize);
256        g_dfxAllocator.dfxMempoolBuf[i].freePagesCnt = 0;
257        g_dfxAllocator.dfxMempoolBuf[i].pageList = NULL;
258    }
259    g_dfxAllocator.initFlag = 1;
260}
261
262static inline PageInfo* GetPageUnchecked(void* ptr)
263{
264    uintptr_t pageHead = PageStart((uintptr_t)(ptr) - PAGE_INFO_SIZE);
265    return (PageInfo*)(pageHead);
266}
267
268static inline PageInfo* GetPage(void* ptr)
269{
270    PageInfo* page = GetPageUnchecked(ptr);
271    if (memcmp(page->tag.tagInfo, DFX_MEM_PAGE_SIGN, sizeof(DFX_MEM_PAGE_SIGN)) != 0) {
272        DFXLOGE("GetPage untagged address!");
273        return NULL;
274    }
275    return page;
276}
277
278static void* AllocMmap(size_t align, size_t size)
279{
280    const size_t headSize = AlignRoundUp(PAGE_INFO_SIZE, align);
281    size_t allocSize = 0;
282
283    // mmap size page allign
284    if (__builtin_add_overflow(headSize, size, &allocSize) ||
285        PageEnd(allocSize) < allocSize) {
286        DFXLOGE("Invalid mmap size!");
287        return NULL;
288    }
289    allocSize = PageEnd(allocSize);
290    void* mptr = mmap(NULL, allocSize, PROT_READ | PROT_WRITE,
291        MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
292    if (mptr == MAP_FAILED) {
293        DFXLOGE("AllocMmap failed!");
294        return NULL;
295    }
296    void* result = (void*)((char*)(mptr) + headSize);
297    PageInfo* page = GetPageUnchecked(result);
298    if (memcpy_s(page->tag.tagInfo, sizeof(page->tag.tagInfo),
299        DFX_MEM_PAGE_SIGN, sizeof(DFX_MEM_PAGE_SIGN)) != EOK) {
300        munmap(mptr, allocSize);
301        DFXLOGE("AllocMmap fill tag failed!");
302        return NULL;
303    }
304    page->tag.type = DFX_MMAP_TYPE;
305    page->tag.mMapAllocSize = allocSize;
306    page->prev = NULL;
307    page->next = NULL;
308    page->freeBlockList = NULL;
309    page->freeBlocksCnt = 0;
310    AddPage(&(g_dfxAllocator.pageList), page);
311    return result;
312}
313
314static void* AllocImpl(size_t align, size_t size)
315{
316    uint32_t type = 0;
317    DfxMempool* mempool = NULL;
318
319    if (g_dfxAllocator.initFlag == 0) {
320        InitDfxAllocator();
321    }
322    if (size > DFX_MEMPOOL_MAX_BLOCK_SIZE) {
323        return AllocMmap(align, size);
324    }
325
326    type = SelectMempoolType(size);
327    if (type < DFX_MEMPOOL_MIN_TYPE) {
328        type = DFX_MEMPOOL_MIN_TYPE;
329    }
330    mempool = &(g_dfxAllocator.dfxMempoolBuf[type - DFX_MEMPOOL_MIN_TYPE]);
331    return MempoolAlloc(mempool);
332}
333
334static void* DfxAlloc(size_t size)
335{
336    size_t realSize = size;
337    if (size == 0) {
338        realSize = 1;
339    }
340    return AllocImpl(ALIGN_SIZE, realSize);
341}
342
343static size_t GetChunkSize(void* ptr)
344{
345    if (ptr == NULL) {
346        DFXLOGE("GetChunkSize Invalid ptr!");
347        return 0;
348    }
349    PageInfo* page = GetPage(ptr);
350    if (page == NULL || (page->tag.type != DFX_MMAP_TYPE &&
351        (page->tag.type < DFX_MEMPOOL_MIN_TYPE || page->tag.type > DFX_MEMPOOL_MAX_TYPE))) {
352        DFXLOGE("GetChunkSize Invalid page!");
353        return 0;
354    }
355    if (page->tag.type == DFX_MMAP_TYPE) {
356        return page->tag.mMapAllocSize - (size_t)((uintptr_t)(ptr) - (uintptr_t)(page));
357    }
358    return g_dfxAllocator.dfxMempoolBuf[page->tag.type - DFX_MEMPOOL_MIN_TYPE].blockSize;
359}
360
361static void DfxFree(void* ptr)
362{
363    if (g_dfxAllocator.initFlag == 0) {
364        return;
365    }
366    if (ptr == NULL) {
367        return;
368    }
369    PageInfo* page = GetPage(ptr);
370    if (page == NULL) {
371        DFXLOGE("DfxFree Invalid page!");
372        return;
373    }
374    if (page->tag.type == DFX_MMAP_TYPE) {
375        RemovePage(&(g_dfxAllocator.pageList), page);
376        munmap(page, page->tag.mMapAllocSize);
377    } else if (page->tag.type <= DFX_MEMPOOL_MAX_TYPE && page->tag.type >= DFX_MEMPOOL_MIN_TYPE) {
378        DfxMempool* mempool = &(g_dfxAllocator.dfxMempoolBuf[page->tag.type - DFX_MEMPOOL_MIN_TYPE]);
379        MempoolFree(mempool, ptr);
380    }
381}
382
383static void* DfxRealloc(void* ptr, size_t size)
384{
385    if (g_dfxAllocator.initFlag == 0) {
386        InitDfxAllocator();
387    }
388    if (ptr == NULL) {
389        return DfxAlloc(size);
390    }
391    if (size == 0) {
392        return NULL;
393    }
394    size_t oldsize = GetChunkSize(ptr);
395    if (oldsize == 0) {
396        return NULL;
397    }
398    if (oldsize < size) {
399        void* res = DfxAlloc(size);
400        if (res) {
401            if (memcpy_s(res, size, ptr, oldsize) != EOK) {
402                DFXLOGE("DfxRealloc memcpy fail");
403            }
404            DfxFree(ptr);
405            return res;
406        }
407    }
408    return ptr;
409}
410
411static void* HookMalloc(size_t size)
412{
413    return DfxAlloc(size);
414}
415
416static void* HookCalloc(size_t count, size_t size)
417{
418    return DfxAlloc(count * size);
419}
420
421static void* HookRealloc(void* ptr, size_t size)
422{
423    return DfxRealloc(ptr, size);
424}
425
426static void HookFree(void* ptr)
427{
428    return DfxFree(ptr);
429}
430
431void RegisterAllocator(void)
432{
433#ifndef DFX_ALLOCATE_ASAN
434    if (memcpy_s(&g_dfxCustomMallocDispatch, sizeof(g_dfxCustomMallocDispatch),
435        &(__libc_malloc_default_dispatch), sizeof(__libc_malloc_default_dispatch)) != EOK) {
436        DFXLOGE("RegisterAllocator memcpy fail");
437    }
438#endif
439    g_dfxCustomMallocDispatch.malloc = HookMalloc;
440    g_dfxCustomMallocDispatch.calloc = HookCalloc;
441    g_dfxCustomMallocDispatch.realloc = HookRealloc;
442    g_dfxCustomMallocDispatch.free = HookFree;
443#ifndef DFX_ALLOCATE_ASAN
444    atomic_store_explicit(&ohos_malloc_hook_shared_library, -1, memory_order_seq_cst);
445    atomic_store_explicit(&__hook_enable_hook_flag, true, memory_order_seq_cst);
446    atomic_store_explicit(&__musl_libc_globals.current_dispatch_table,
447        (volatile const long long)&g_dfxCustomMallocDispatch, memory_order_seq_cst);
448#endif
449}
450
451void UnregisterAllocator(void)
452{
453#ifndef DFX_ALLOCATE_ASAN
454    atomic_store_explicit(&ohos_malloc_hook_shared_library, 0, memory_order_seq_cst);
455    atomic_store_explicit(&__hook_enable_hook_flag, false, memory_order_seq_cst);
456    atomic_store_explicit(&__musl_libc_globals.current_dispatch_table, 0, memory_order_seq_cst);
457#endif
458}
459
460DfxAllocator* GetDfxAllocator(void)
461{
462    return &g_dfxAllocator;
463}
464
465int IsDfxAllocatorMem(void* ptr)
466{
467    if (!ptr) {
468        return 0;
469    }
470    if (GetPage(ptr)) {
471        return 1;
472    }
473    return 0;
474}
475