v 0. Pasted by Zeux as cpp at 2009-03-01 17:10:30 MSK and set expiration to never.

Paste will expire never. Expiration is locked.

  1. // User-defined global heap prototypes
  2. extern void mem_global_init();
  3. extern void mem_global_term();
  4.  
  5. extern void* mem_global_allocate(size_t size, size_t align);
  6. extern void mem_global_deallocate(void* ptr);
  7. extern size_t mem_global_get_size(void* ptr);
  8.  
  9. // Actual code
  10. #define BREAK() __debugbreak()
  11.  
  12. #include <stdio.h>
  13. #include <stdlib.h>
  14.  
  15. #define WIN32_LEAN_AND_MEAN
  16. #include <windows.h>
  17.  
  18. static inline void patch_with_jump(void* dest, void* address)
  19. {
  20.     // get offset for relative jmp
  21.     unsigned int offset = (unsigned int)((char*)address - (char*)dest - 5);
  22.    
  23.     // unprotect memory
  24.     unsigned long old_protect;
  25.     VirtualProtect(dest, 5, PAGE_READWRITE, &old_protect);
  26.    
  27.     // write jmp
  28.     *(unsigned char*)dest = 0xe9;
  29.     *(unsigned int*)((char*)dest + 1) = offset;
  30.    
  31.     // protect memory
  32.     VirtualProtect(dest, 5, old_protect, &old_protect);
  33. }
  34.  
  35. static void break_stub()
  36. {
  37.     BREAK();
  38. }
  39.  
  40. #define PATCH(target, stub) void target(); patch_with_jump((void*)(target), (void*)(stub))
  41. #define PATCH_STUB(target) PATCH(target, target ## _stub)
  42. #define PATCH_BREAK(target) PATCH(target, break_stub)
  43.  
  44. #ifdef _DEBUG
  45. #define PATCH_DEBUG(target, stub) PATCH(target, stub)
  46. #define PATCH_BREAK_DEBUG(target) PATCH_BREAK(target)
  47. #else
  48. #define PATCH_DEBUG(target, stub) (void)0
  49. #define PATCH_BREAK_DEBUG(target) (void)0
  50. #endif
  51.  
  52. #define PATCH_BREAK_ALL(target) PATCH_BREAK(target); PATCH_BREAK_DEBUG(target ## _base); PATCH_BREAK_DEBUG(target ## _dbg)
  53.  
  54. extern "C"
  55. {
  56.     const size_t malloc_alignment = 8;
  57.    
  58.     // This is a copy of __freeCrtMemory (we can't use it directly since it's not in CRT in release)
  59.     static void freeCrtMemory()
  60.     {
  61.         // code removed to avoid EULA issues (I hate licenses)
  62.     }
  63.    
  64.     static void cleanup_crt_leaks()
  65.     {
  66.         freeCrtMemory();
  67.        
  68.         // turn off buffering for stdout/stderr and release buffers
  69.         extern void* _stdbuf[2];
  70.        
  71.         setvbuf(stdout, 0, _IONBF, 0);
  72.         free(_stdbuf[0]);
  73.         _stdbuf[0] = NULL;
  74.        
  75.         setvbuf(stderr, 0, _IONBF, 0);
  76.         free(_stdbuf[1]);
  77.         _stdbuf[1] = NULL;
  78.        
  79.         // free ptmcbinfo (allocated in _setmbcp)
  80.         extern pthreadmbcinfo __ptmbcinfo;
  81.         free(__ptmbcinfo);
  82.         __ptmbcinfo = NULL;
  83.        
  84.         // free ioinfo (allocated in _ioinit)
  85.         extern void* __pioinfo[];
  86.         free(__pioinfo[0]);
  87.         __pioinfo[0] = NULL;
  88.        
  89.         // free ptd (allocated in _mtinit)
  90.         void* _getptd_noexit();
  91.         free(_getptd_noexit());
  92.        
  93.         // free TLS indices (allocated in _mtinit) and critical sections (allocated in _mtinitlocknum)
  94.         void _mtterm();
  95.         _mtterm();
  96.     }
  97.    
  98.     static int heap_init_stub(int)
  99.     {
  100.         // call user init function
  101.         mem_global_init();
  102.        
  103.         // pretend that we initialized CRT heap (needed for _mtinitlocknum)
  104.         extern HANDLE _crtheap;
  105.         _crtheap = (HANDLE)-1;
  106.        
  107.         return 1;
  108.     }
  109.    
  110.     static void heap_term_stub()
  111.     {
  112.         // cleanup memory that CRT writers are too lazy to free themselves (bastards!)
  113.         cleanup_crt_leaks();
  114.        
  115.         // call user init function
  116.         mem_global_term();
  117.     }
  118.    
  119.     static void* malloc_stub(size_t size)
  120.     {
  121.         return mem_global_allocate(size, malloc_alignment);
  122.     }
  123.    
  124.     static void* calloc_stub(size_t count, size_t size)
  125.     {
  126.         void* ptr = mem_global_allocate(count * size, malloc_alignment);
  127.  
  128.         if (ptr) memset(ptr, 0, mem_global_get_size(ptr));
  129.  
  130.         return ptr;
  131.     }
  132.    
  133.     static void* realloc_stub(void* ptr, size_t size)
  134.     {
  135.         if (!ptr)
  136.         {
  137.             return mem_global_allocate(size, malloc_alignment);
  138.         }
  139.  
  140.         if (!size)
  141.         {
  142.             mem_global_deallocate(ptr);
  143.             return 0;
  144.         }
  145.  
  146.         size_t old_size = mem_global_get_size(ptr);
  147.  
  148.         void* result = mem_global_allocate(size, malloc_alignment);
  149.  
  150.         if (result)
  151.         {
  152.             memcpy(result, ptr, min(size, old_size));
  153.             mem_global_deallocate(ptr);
  154.         }
  155.  
  156.         return result;
  157.     }
  158.    
  159.     // Note: this function has a bug - since it calculates size to clear using
  160.     // mem_global_get_size, and it can return inaccurate values (i.e. more than
  161.     // allocated size), it's possible that not everything is cleared (i.e. malloc(0),
  162.     // followed by recalloc(100) will not clear first 12 bytes for current implementation).
  163.     // However, there is no problem if only calloc/_recalloc were used for the block in question,
  164.     // and this is the case in CRT, so I won't worry about it.
  165.     static void* recalloc_stub(void* ptr, size_t count, size_t size)
  166.     {
  167.         if (!ptr)
  168.         {
  169.             return calloc_stub(count, size);
  170.         }
  171.  
  172.         if (!size)
  173.         {
  174.             mem_global_deallocate(ptr);
  175.             return 0;
  176.         }
  177.  
  178.         size_t old_size = mem_global_get_size(ptr);
  179.  
  180.         void* result = calloc_stub(count, size);
  181.  
  182.         if (result)
  183.         {
  184.             memcpy(result, ptr, min(size, old_size));
  185.             mem_global_deallocate(ptr);
  186.         }
  187.  
  188.         return result;
  189.     }
  190.    
  191.     static void free_stub(void* ptr)
  192.     {
  193.         mem_global_deallocate(ptr);
  194.     }
  195.    
  196.     static size_t msize_stub(void* ptr)
  197.     {
  198.         return mem_global_get_size(ptr);
  199.     }
  200.    
  201.     static void patch_memory_management_functions()
  202.     {
  203.         // heap init/term
  204.         // note: we patch __crtCorExitProcess because _heap_term is called only in DLL CRT
  205.         PATCH(_heap_init, heap_init_stub);
  206.         PATCH(__crtCorExitProcess, heap_term_stub);
  207.        
  208.         // malloc
  209.         PATCH(malloc, malloc_stub);
  210.         PATCH(_malloc_crt, malloc_stub);
  211.         PATCH_DEBUG(_malloc_base, malloc_stub);
  212.         PATCH_DEBUG(_malloc_dbg, malloc_stub);
  213.        
  214.         // calloc
  215.         PATCH(calloc, calloc_stub);
  216.         PATCH(_calloc_crt, calloc_stub);
  217.         PATCH_DEBUG(_calloc_base, calloc_stub);
  218.         PATCH_DEBUG(_calloc_dbg, calloc_stub);
  219.        
  220.         // realloc
  221.         PATCH(realloc, realloc_stub);
  222.         PATCH(_realloc_crt, realloc_stub);
  223.         PATCH_DEBUG(_realloc_base, realloc_stub);
  224.         PATCH_DEBUG(_realloc_dbg, realloc_stub);
  225.        
  226.         // recalloc
  227.         PATCH(_recalloc, recalloc_stub);
  228.         PATCH(_recalloc_crt, recalloc_stub);
  229.         PATCH_DEBUG(_recalloc_base, recalloc_stub);
  230.         PATCH_DEBUG(_recalloc_dbg, recalloc_stub);
  231.        
  232.         // free
  233.         PATCH(free, free_stub);
  234.         PATCH_DEBUG(_free_nolock, free_stub);
  235.         PATCH_DEBUG(_free_dbg_nolock, free_stub);
  236.         PATCH_DEBUG(_free_base, free_stub);
  237.         PATCH_DEBUG(_free_dbg, free_stub);
  238.        
  239.         // msize
  240.         PATCH(_msize, msize_stub);
  241.         PATCH_DEBUG(_msize_base, msize_stub);
  242.         PATCH_DEBUG(_msize_dbg, msize_stub);
  243.        
  244.         // aligned functions
  245.         PATCH_BREAK_ALL(_aligned_malloc);
  246.         PATCH_BREAK_ALL(_aligned_offset_malloc);
  247.         PATCH_BREAK_ALL(_aligned_realloc);
  248.         PATCH_BREAK_ALL(_aligned_recalloc);
  249.         PATCH_BREAK_ALL(_aligned_offset_realloc);
  250.         PATCH_BREAK_ALL(_aligned_msize);
  251.         PATCH_BREAK_ALL(_aligned_offset_recalloc);
  252.         PATCH_BREAK_ALL(_aligned_free);
  253.         PATCH_BREAK_ALL(_aligned_malloc);
  254.        
  255.         // supporting functions
  256.         PATCH_BREAK_DEBUG(_nh_malloc);
  257.         PATCH_BREAK_DEBUG(_nh_malloc_dbg);
  258.         PATCH_BREAK_ALL(_heap_alloc);
  259.         PATCH_BREAK_ALL(_expand);
  260.     }
  261.    
  262.     int entrypoint()
  263.     {
  264.         int mainCRTStartup();
  265.        
  266.         patch_memory_management_functions();
  267.        
  268.         return mainCRTStartup();
  269.     }
  270. };


Editing is locked.