/*
    Copyright 2012 Luigi Auriemma

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307 USA

    http://www.gnu.org/licenses/gpl-2.0.txt
*/

#ifdef XDBG_ALLOC_DISABLED
#define XDBG_ALLOC_INCLUDE
#endif

#ifndef XDBG_ALLOC_INCLUDE
#define XDBG_ALLOC_INCLUDE

    /*
    everything remains allocated till the closing of the main program
    the only exceptions are the reallocations that require more memory than
    that available in the previous buffer, in which case it must be
    freed and allocated newly.

    11112222333333333333333334444
    |   |   |                |
    |   |   |                PAGE_GUARD
    |   |   |the requested buffer
    |   |   
    |   PAGE_GUARD
    structure with informations

    notes: stdout is necessary for debugging
    */

    /*
    problems:
    - if I don't use the index (linked lists) there is the increasing of the
      problem about memory unallocable due to VirtualAlloc that doesn't
      manipolate the memory like malloc
    - a secure usage of the memory caged between two PAGE_GUARDS take 16kb
      for each allocation (even if you try to allocate 1 byte) wasting tons
      of memory
    - the usage of linked lists make the program very slow when there are
      many allocations
    */



#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <malloc.h>

#ifdef WIN32
    #include <windows.h>
#else
    #include <sys/mman.h>
    #define MEM_COMMIT      (MAP_PRIVATE | MAP_ANONYMOUS)
    #define MEM_RESERVE     0
    #define MEM_RELEASE     0
    #define PAGE_NOACCESS   PROT_NONE
    #define PAGE_READONLY   PROT_READ
    #define PAGE_GUARD      PROT_NONE
    #define PAGE_READWRITE  (PROT_READ | PROT_WRITE)

    typedef uint32_t    DWORD;

    typedef struct {
        DWORD   RegionSize;
    } MEMORY_BASIC_INFORMATION;
#endif



void *(*real_malloc)(size_t size)                   = (void *)malloc;
void *(*real_calloc)(size_t nelem, size_t elsize)   = (void *)calloc;
void *(*real_realloc)(void *mem, size_t size)       = (void *)realloc;
void (*real_free)(void *mem)                        = (void *)free;



        // what if the pointer has been allocated by another library that doens't use this system?
        // char *p = mydll_func();
        // free(p); // would return an error!
#define XDBG_ALLOC_ALLOW_UNKNOWN_FREE

static int XDBG_ALLOC_ACTIVE    = 1;    // the alternative is using the classical malloc
static int XDBG_ALLOC_INDEX     = 1;    // use linked lists (must be enabled!)
static int XDBG_ALLOC_VERBOSE   = 0;    // debug informations
static int XDBG_HEAPVALIDATE    = 0;    // time consuming!
//#define XDBG_USE_MALLOC     // an experiment, it can't be used with PAGE_GUARD because
                              // memory gets moved and so the PAGE_GUARDs affect other
                              // buffers



                                     // info          // page_guard
#define XDBG_ALLOC_MEM2PTR(X)   (X + (xdbg_pagesize + xdbg_pagesize))
#define XDBG_ALLOC_PTR2MEM(X)   (X - (xdbg_pagesize + xdbg_pagesize))
#define XDBG_ALLOC_MEM2INFO(X)  ((xdbg_alloc_t *)(X))
#define XDBG_ALLOC_PTR2INFO(X)  ((xdbg_alloc_t *)(XDBG_ALLOC_PTR2MEM(X)))
#define XDBG_ALLOC_FREECHR      0xcc
#define XDBG_ALLOC_ERROR        xdbg_alloc_error(__FILE__, __LINE__, __FUNCTION__)
#define xbdg_alloc_dontuse      "Error: do NOT use malloc/calloc/realloc/free!"
#define malloc                  xdbg_malloc     //xbdg_alloc_dontuse
#define calloc                  xdbg_calloc     //xbdg_alloc_dontuse
#define realloc                 xdbg_realloc    //xbdg_alloc_dontuse
#define free                    xdbg_free       //xbdg_alloc_dontuse



#pragma pack(16)
typedef struct {
    size_t  original;   // the real size requested by the user
    size_t  size;       // the amount of data really allocated without PAGE_GUARDs
                        // (so only the structure + the memory)
    size_t  xsize;      // everything included the PAGE_GUARDs
    size_t  active;     // in use
    void    *next;      // the next allocated memory
    void    *prev;      // the previons allocated memory
    void    *pad[3];    // unused, only to align memory
} xdbg_alloc_t;
#pragma pack()



static size_t   xdbg_pagesize       = 0;
static void     *xdbg_alloc_main    = NULL;



#ifndef WIN32
    void *VirtualAlloc(void *lpAddress, size_t dwSize, DWORD flAllocationType, DWORD flProtect) {
        lpAddress = mmap(lpAddress, dwSize, flProtect, flAllocationType, 0, 0);
        if((long)lpAddress == -1) return(NULL);
        return(lpAddress);
    }

    size_t VirtualProtect(void *lpAddress, size_t dwSize, DWORD flNewProtect, DWORD *lpflOldProtect) {
        if(mprotect(lpAddress, dwSize, flNewProtect) == -1) return(0);
        return(1);
    }

    size_t VirtualFree(void *lpAddress, size_t dwSize, DWORD dwFreeType) {
        if(!lpAddress) return(1);
        if(!dwSize) dwSize = XDBG_ALLOC_MEM2INFO(lpAddress)->xsize;
        if(munmap(lpAddress, dwSize) == -1) return(0);
        return(1);
    }

    size_t HeapValidate(size_t hHeap, DWORD dwFlags, void *lpMem) {
        return(1);
    }

    size_t GetProcessHeap(void) {
        return(1);
    }

    #define CopyMemory  memcpy
#endif



void xdbg_alloc_error(const char *fname, DWORD line, const char *func) {
    char    tmp[16];
    char    *message = NULL;

    #ifdef WIN32
    if(!GetLastError()) {
        message = "";
    } else {
        FormatMessage(
          FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
          NULL,
          GetLastError(),
          0,
          (char *)&message,
          0,
          NULL);
    }
    #endif
    if(!message) message = strerror(errno);

    fprintf(stderr, "\n- error in %s line %ld: %s()\n", fname, line, func);
    fprintf(stderr, "\n"
        "Error: memory allocation problem\n"
        "       %s\n"
        "\n"
        "press RETURN to quit\n",
        message ? message : "unknown error");
    // if(message && message[0]) LocalFree(message);  // no!!!
    fgets(tmp, sizeof(tmp), stdin);
    exit(1);
}



size_t xdbg_alloc_align(size_t size) {
    if(!xdbg_pagesize) {
        #ifdef WIN32
        SYSTEM_INFO sSysInfo;
        GetSystemInfo(&sSysInfo);
        xdbg_pagesize = sSysInfo.dwPageSize;
        #else
        xdbg_pagesize = getpagesize();
        #endif
        if(!xdbg_pagesize) xdbg_pagesize = 4096;
        if(XDBG_ALLOC_VERBOSE) fprintf(stdout, "# xdbg_pagesize %u\n", xdbg_pagesize);
    }
    size = (size + (xdbg_pagesize - 1)) & (~(xdbg_pagesize - 1));
    return(size);
}



size_t xdbg_alloc_init(size_t *ret_original, size_t *ret_size, size_t *ret_xsize) {
    size_t  original,
            size,
            xsize;

    if(!ret_size) XDBG_ALLOC_ERROR;
    size = *ret_size;

    original = size;
    size = xdbg_alloc_align(size);
    if(size < original) XDBG_ALLOC_ERROR;

    *ret_size = size;
    if(ret_original) *ret_original = original;
    if(ret_xsize) {
                //info          //page_guard    //data //page_guard
        xsize = xdbg_pagesize + xdbg_pagesize + size + xdbg_pagesize;
        if(xsize < original) XDBG_ALLOC_ERROR;
        *ret_xsize = xsize;
    }
    return(size);
}



/*
mem: pointer to the beginning of the allocated memory block
*/
int xdbg_alloc_remove_mem(void *mem) {
    if(!XDBG_ALLOC_INDEX) return(0);

    void    *prev_mem,
            *next_mem;

    if(!mem) return(-1);
    if(!xdbg_alloc_main) return(-1);
    next_mem = XDBG_ALLOC_MEM2INFO(mem)->next;
    if(xdbg_alloc_main == mem) {
        xdbg_alloc_main = next_mem;
        if(xdbg_alloc_main) XDBG_ALLOC_MEM2INFO(xdbg_alloc_main)->prev = NULL;
    } else {
        prev_mem = XDBG_ALLOC_MEM2INFO(mem)->prev;
        if(prev_mem) XDBG_ALLOC_MEM2INFO(prev_mem)->next = XDBG_ALLOC_MEM2INFO(mem)->next;
        if(next_mem) XDBG_ALLOC_MEM2INFO(next_mem)->prev = XDBG_ALLOC_MEM2INFO(mem)->prev;
    }
    return(0);
}



/*
new_mem: pointer to the beginning of the allocated memory block
*/
int xdbg_alloc_add_mem(void *new_mem) {
    if(!XDBG_ALLOC_INDEX) return(0);

    if(!new_mem) return(-1);
    XDBG_ALLOC_MEM2INFO(new_mem)->prev  = NULL;
    XDBG_ALLOC_MEM2INFO(new_mem)->next  = xdbg_alloc_main;
    if(xdbg_alloc_main) XDBG_ALLOC_MEM2INFO(xdbg_alloc_main)->prev = new_mem;
    xdbg_alloc_main = new_mem;
    return(0);
}



void *xdbg_alloc_reuse(size_t size) {
    if(!XDBG_ALLOC_INDEX) return(NULL);

    xdbg_alloc_t    *xdbg_alloc;
    void    *mem;

    size = xdbg_alloc_align(size);
    for(mem = xdbg_alloc_main; mem; mem = xdbg_alloc->next) {
        xdbg_alloc = XDBG_ALLOC_MEM2INFO(mem);
        if(!xdbg_alloc->active) {
            if(xdbg_alloc->size >= size) return(mem);
        }
    }
    return(NULL);
}



void *xdbg_alloc_check(void *mymem, char check_active) {
    if(!XDBG_ALLOC_INDEX) return((void *)0xffffffff);

    void    *mem;

    for(mem = xdbg_alloc_main; mem; mem = XDBG_ALLOC_MEM2INFO(mem)->next) {
        if(mem == mymem) {
            if(check_active && !XDBG_ALLOC_MEM2INFO(mem)->active) break;    // use-after-free
            return(mem);
        }
    }
    return(NULL);
}



void *xdbg_alloc_check_ptr(void *ptr, char check_active) {
    return(xdbg_alloc_check(XDBG_ALLOC_PTR2MEM(ptr), check_active));
}



size_t xdbg_alloc_size(void *ptr) {
    if(!ptr) return(0);
    xdbg_alloc_check_ptr(ptr, 1);
    return(XDBG_ALLOC_PTR2INFO(ptr)->original);
}



/*
ptr: the buffer visible by the main program
*/
void xdbg_real_free( void *ptr ) {
    if(XDBG_ALLOC_VERBOSE) fprintf(stdout, "# real_free(%p)\n", ptr);

    if(!ptr) return;

    if(XDBG_HEAPVALIDATE) {
        if(!HeapValidate(GetProcessHeap(), 0, NULL)) XDBG_ALLOC_ERROR;  // NO ptr!
    }

    if(!XDBG_ALLOC_ACTIVE) {
        real_free(ptr);
        return;
    }

    void    *mem;

    if(!xdbg_alloc_check_ptr(ptr, 0)) {
        #ifdef XDBG_ALLOC_ALLOW_UNKNOWN_FREE
        real_free(ptr);
        return;
        #else
        XDBG_ALLOC_ERROR;
        #endif
    }
    mem = XDBG_ALLOC_PTR2MEM(ptr);
    xdbg_alloc_remove_mem(mem);

    #ifndef WIN32
    size_t  xsize;  // xsize is a backup of the one contained in the
                    // structure because I clear (memset) it before
                    // the free, and so Linux can't work without it
    xsize = XDBG_ALLOC_MEM2INFO(mem)->xsize;
    #endif

    // may help with use-after-free
    memset(ptr, XDBG_ALLOC_FREECHR, XDBG_ALLOC_MEM2INFO(mem)->size);            // data
    memset(XDBG_ALLOC_MEM2INFO(mem), XDBG_ALLOC_FREECHR, sizeof(xdbg_alloc_t)); // info (xdbg_alloc->size is made before)

    #ifdef XDBG_USE_MALLOC
    real_free(mem);
    #else
    if(!VirtualFree(mem,
        #ifdef WIN32
        0
        #else
        xsize
        #endif
        , MEM_RELEASE)) XDBG_ALLOC_ERROR;
    #endif
}



/*
ptr: the buffer visible by the main program
*/
void xdbg_free( void *ptr ) {
    if(!XDBG_ALLOC_INDEX) {
        xdbg_real_free(ptr);
        return;
    }

    if(XDBG_ALLOC_VERBOSE) fprintf(stdout, "# free(%p)\n", ptr);

    if(!ptr) return;

    if(XDBG_HEAPVALIDATE) {
        if(!HeapValidate(GetProcessHeap(), 0, NULL)) XDBG_ALLOC_ERROR;  // NO ptr!
    }

    if(!XDBG_ALLOC_ACTIVE) {
        real_free(ptr);
        return;
    }

    xdbg_alloc_t    *xdbg_alloc;
    if(!xdbg_alloc_check_ptr(ptr, 1)) {
        // what if the pointer has been allocated by another library that doens't use this system?
        // char *p = mydll_func();
        // free(p); // would return an error!
        #ifdef XDBG_ALLOC_ALLOW_UNKNOWN_FREE
        real_free(ptr);
        return;
        #else
        XDBG_ALLOC_ERROR;
        #endif
    }
    xdbg_alloc = XDBG_ALLOC_PTR2INFO(ptr);
    xdbg_alloc->active = 0;
    memset(ptr, XDBG_ALLOC_FREECHR, xdbg_alloc->size);  // may help with use-after-free
}



void xdbg_freeall(void) {
    if(!XDBG_ALLOC_INDEX) return;

    void    *mem,
            *p;

    for(mem = xdbg_alloc_main; mem; mem = p) {
        p = XDBG_ALLOC_MEM2INFO(mem)->next;  // save it because the structure gets freed
        xdbg_real_free(XDBG_ALLOC_MEM2PTR(mem));
        xdbg_alloc_main = p;
    }
}



/*
ptr: the buffer visible by the main program
new_size: requested size
*/
void xdbg_realloc_page_guard(void *ptr, size_t new_size) {
    // unfortunately a continous usage of VirtualProtect
    // "kills" the performances in a terrible way so my idea
    // is simply to keep only the last PAGE_GUARD active
    // and leaving the rest as PAGE_READWRITE:
    //   1111233333333333333334444  malloc(16)
    //   11112333333..........4444  realloc(buff, 6)
    //   111123333333333......4444  realloc(buff, 10)
    // what really matters is that the end of the allocated
    // memory is guarded and so secure.
    // even the memset kills time!
    return;

    size_t  total_size;

    total_size = XDBG_ALLOC_PTR2INFO(ptr)->size;

    //total_size = xdbg_alloc_align(total_size);
    new_size = xdbg_alloc_align(new_size);
    if(total_size < new_size) XDBG_ALLOC_ERROR;

    #ifndef XDBG_USE_MALLOC
    DWORD   tmp;
    if(!VirtualProtect(
        ptr,
        new_size,
        PAGE_READWRITE, &tmp)) XDBG_ALLOC_ERROR;                // our memory
    #endif
    memset(
        ptr + new_size,
        XDBG_ALLOC_FREECHR,
        total_size - new_size);
    #ifndef XDBG_USE_MALLOC
    if(!VirtualProtect(
        ptr + new_size,
        (total_size - new_size) + xdbg_pagesize,
        PAGE_READONLY | PAGE_GUARD, &tmp)) XDBG_ALLOC_ERROR;    // post PAGE_GUARD
    #endif

    if(XDBG_ALLOC_VERBOSE) fprintf(stdout, "# realloc_page_guard(%p, %u, %u)\n", ptr, new_size, total_size);
}



void *xdbg_malloc( size_t size ) {
    if(!XDBG_ALLOC_ACTIVE) return(real_malloc(size));

    xdbg_alloc_t    *xdbg_alloc;
    size_t  xsize,
            original;
    unsigned char   *mem;

    xdbg_alloc_init(&original, &size, &xsize);

    if(XDBG_ALLOC_VERBOSE) fprintf(stdout, "# malloc(%u) -> %u %u\n", original, size, xsize);

    if(XDBG_HEAPVALIDATE) {
        if(!HeapValidate(GetProcessHeap(), 0, NULL)) XDBG_ALLOC_ERROR;
    }

    mem = xdbg_alloc_reuse(size);
    if(mem) {
        if(XDBG_ALLOC_VERBOSE) fprintf(stdout, "# xdbg_alloc_reuse %p\n", mem);

        xdbg_alloc = XDBG_ALLOC_MEM2INFO(mem);
        xdbg_alloc->original    = original;
        xdbg_alloc->active      = 1;
        memset(XDBG_ALLOC_MEM2PTR(mem), 0, size);   // needed! think to calloc!
        //xdbg_realloc_page_guard(XDBG_ALLOC_MEM2PTR(mem), size);
    } else {
        #ifdef XDBG_USE_MALLOC
        mem = real_calloc(1, xsize);
        #else
        mem = VirtualAlloc(NULL, xsize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
        #endif
        if(!mem) XDBG_ALLOC_ERROR;
        //ZeroMemory(mem, xsize);   // VirtualAlloc/real_calloc automatically put zeroes in the memory

        #ifndef XDBG_USE_MALLOC
        DWORD   tmp;
        //if(!VirtualProtect(mem,                                         xdbg_pagesize,  PAGE_READWRITE,             &tmp)) XDBG_ALLOC_ERROR;    // informations
        if(!VirtualProtect(mem + xdbg_pagesize,                         xdbg_pagesize,  PAGE_READONLY | PAGE_GUARD, &tmp)) XDBG_ALLOC_ERROR;    // pre PAGE_GUARD
        //if(!VirtualProtect(mem + xdbg_pagesize + xdbg_pagesize,         size,           PAGE_READWRITE,             &tmp)) XDBG_ALLOC_ERROR;    // our memory
        if(!VirtualProtect(mem + xdbg_pagesize + xdbg_pagesize + size,  xdbg_pagesize,  PAGE_READONLY | PAGE_GUARD, &tmp)) XDBG_ALLOC_ERROR;    // post PAGE_GUARD
        #endif

        xdbg_alloc = XDBG_ALLOC_MEM2INFO(mem);
        xdbg_alloc->original    = original;
        xdbg_alloc->size        = size;
        xdbg_alloc->xsize       = xsize;
        xdbg_alloc->active      = 1;
        xdbg_alloc->next        = NULL; // set by xdbg_alloc_add_mem
        xdbg_alloc->prev        = NULL; // set by xdbg_alloc_add_mem
        xdbg_alloc_add_mem(mem);
    }

    if(XDBG_ALLOC_VERBOSE)
        fprintf(stdout, "# %p (%p -> %p) = malloc(%u / %u / %u)\n",
            mem, XDBG_ALLOC_MEM2PTR(mem), mem + xsize, original, size, xsize);
    return(XDBG_ALLOC_MEM2PTR(mem));
}



void *xdbg_realloc( void *ptr, size_t size ) {
    if(!XDBG_ALLOC_ACTIVE) return(real_realloc(ptr, size));

    if(XDBG_ALLOC_VERBOSE) fprintf(stdout, "# realloc(%p, %u)\n", ptr, size);

    if(!ptr) return(xdbg_malloc(size));

    if(XDBG_HEAPVALIDATE) {
        if(!HeapValidate(GetProcessHeap(), 0, NULL)) XDBG_ALLOC_ERROR;  // NO ptr!
    }

    xdbg_alloc_t    *xdbg_alloc;
    size_t  original;
    void    *new_ptr;

    xdbg_alloc_init(&original, &size, NULL);
    if(!xdbg_alloc_check_ptr(ptr, 1)) XDBG_ALLOC_ERROR;

    xdbg_alloc = XDBG_ALLOC_PTR2INFO(ptr);
    if(xdbg_alloc->size >= size) {  // aligned_allocated versus aligned_requested
        if(XDBG_ALLOC_VERBOSE) fprintf(stdout, "# realloc reused same buffer %p (%u %u)\n", ptr, size, xdbg_alloc->size);
        xdbg_alloc->original = original;
        //memset(ptr, 0, xdbg_alloc->size - original);  // takes really lot of time!
        //xdbg_realloc_page_guard(ptr, size);
        return(ptr);
    }

    new_ptr = xdbg_malloc(original);
    if(!new_ptr) XDBG_ALLOC_ERROR;
    if(XDBG_ALLOC_VERBOSE) {
        void    *new_mem = XDBG_ALLOC_PTR2MEM(new_ptr);
        void    *old_mem = XDBG_ALLOC_PTR2MEM(ptr);
        size_t  tmp = XDBG_ALLOC_MEM2INFO(new_mem)->xsize;
        fprintf(stdout, "# %p (%p -> %p) = realloc(%p, %u / %u / %u)\n",
            new_mem, new_ptr, new_mem + tmp, old_mem, original, size, tmp);
    }
    CopyMemory(new_ptr, ptr, xdbg_alloc->original);
    xdbg_real_free(ptr);    // real_free because it's not good to leave it allocated wasting space
    return(new_ptr);
}



void *xdbg_calloc( size_t num, size_t size ) {
    if(!XDBG_ALLOC_ACTIVE) return(real_calloc(num, size));

    if(XDBG_ALLOC_VERBOSE) fprintf(stdout, "# calloc(%u, %u)\n", num, size);

    uint64_t    tmp;
    tmp = (uint64_t)num * (uint64_t)size;
    if(tmp > (uint64_t)0xffffffffLL) XDBG_ALLOC_ERROR;
    return(xdbg_malloc(tmp));
}



// replace any call to the original malloc with the new one.
// not needed because almost unused and only for x86
// and would make the XDBG_ALLOC_ACTIVE impossible to use
void xdbg_alloc_extreme(void) {
    #define XDBG_ALLOC_EXTREME_JMP(X) \
        p = (void *)real_##X; \
        if(VirtualProtect(p, 1 + sizeof(void *), PAGE_READWRITE, &tmp)) { \
            p[0] = 0xe9; \
            *((uint32_t *)(p + 1)) = ((void *)X) - (void *)(p + 1 + 4); \
            VirtualProtect(p, 1 + sizeof(void *), tmp, &tmp); \
        }

    #ifdef _X86_
    static int      init = 0;
    DWORD           tmp;
    unsigned char   *p;

    if(!init) {
        XDBG_ALLOC_EXTREME_JMP(malloc)
        XDBG_ALLOC_EXTREME_JMP(calloc)
        XDBG_ALLOC_EXTREME_JMP(realloc)
        XDBG_ALLOC_EXTREME_JMP(free)
        init = 1;
    }
    #endif
}

#else

void xdbg_freeall(void) {
}

#endif

