1 ////////////////////////////////////////////////////////////////////
2 // Copyright (C) Alexander Telyatnikov, Ivan Keliukh, Yegor Anchishkin, SKIF Software, 1999-2013. Kiev, Ukraine
3 // All rights reserved
4 // This file was released under the GPLv2 on June 2015.
5 ////////////////////////////////////////////////////////////////////
6 
7 #ifndef __MY_MEM_TOOLS_H__
8 #define __MY_MEM_TOOLS_H__
9 
10 #define MY_HEAP_FLAG_USED       0x00000001
11 #define MY_HEAP_FLAG_LEN_MASK   0xfffffffe
12 
13 #define MyFreeMemoryAndPointer(ptr) \
14     if(ptr) {                   \
15         MyFreePool__(ptr);      \
16         ptr = NULL;             \
17     }
18 
19 #define MY_USE_ALIGN
20 //#define MY_MEM_BOUNDS_CHECK
21 
22 typedef struct _MEM_ALLOC_DESC {
23     ULONG Addr;
24     ULONG Len;
25 #ifdef MY_HEAP_TRACK_OWNERS
26     USHORT Src;
27     USHORT Line;
28 #endif
29 #ifdef MY_HEAP_TRACK_REF
30 //    PCHAR Ref;
31     PCHAR Tag;
32 #endif
33 } MEM_ALLOC_DESC, *PMEM_ALLOC_DESC;
34 
35 typedef struct _MEM_FRAME_ALLOC_DESC {
36     PMEM_ALLOC_DESC Frame;
37     ULONG LastUsed;
38     ULONG FirstFree;
39     ULONG Type;
40 } MEM_FRAME_ALLOC_DESC, *PMEM_FRAME_ALLOC_DESC;
41 
42 extern PCHAR BreakAddr;
43 extern ULONG MemTotalAllocated;
44 
45 #define MY_HEAP_FRAME_SIZE          (256*1024)
46 #define MY_HEAP_MAX_FRAMES          512
47 #define MY_HEAP_MAX_BLOCKS          4*1024  // blocks per frame
48 
49 #ifdef USE_THREAD_HEAPS
50 //extern HANDLE MemLock;
51 extern "C" VOID ExInitThreadPools();
52 extern "C" VOID ExDeInitThreadPools();
53 extern "C" VOID ExFreeThreadPool();
54 #endif //USE_THREAD_HEAPS
55 
56 // Mem
57 BOOLEAN MyAllocInit(VOID);
58 VOID MyAllocRelease(VOID);
59 #ifdef MY_HEAP_TRACK_OWNERS
60 PCHAR MyAllocatePool(ULONG Type, ULONG size, USHORT Src, USHORT Line
61   #ifdef MY_HEAP_TRACK_REF
62                     ,PCHAR Tag
63   #endif //MY_HEAP_TRACK_REF
64       );
65 ULONG MyReallocPool( PCHAR addr, ULONG OldLength, PCHAR* NewBuff, ULONG NewLength, USHORT Src, USHORT Line);
66 #else
67 PCHAR __fastcall MyAllocatePool(ULONG Type, ULONG size
68   #ifdef MY_HEAP_TRACK_REF
69                     ,PCHAR Tag
70   #endif //MY_HEAP_TRACK_REF
71       );
72 ULONG __fastcall MyReallocPool( PCHAR addr, ULONG OldLength, PCHAR* NewBuff, ULONG NewLength);
73 #endif
74 VOID __fastcall MyFreePool(PCHAR addr);
75 
76 #ifdef MY_HEAP_CHECK_BOUNDS
77   #define MY_HEAP_ALIGN             63
78 #else
79   #define MY_HEAP_ALIGN             63
80 #endif
81 #define PAGE_SIZE_ALIGN             (PAGE_SIZE - 1)
82 
83 #define AlignToPageSize(size) (((size)+PAGE_SIZE_ALIGN)&(~PAGE_SIZE_ALIGN))
84 #define MyAlignSize__(size) (((size)+MY_HEAP_ALIGN)&(~MY_HEAP_ALIGN))
85 #ifdef MY_HEAP_FORCE_NONPAGED
86 #define MyFixMemType(type)          NonPagedPool
87 #else //MY_HEAP_FORCE_NONPAGED
88 #define MyFixMemType(type)          (type)
89 #endif //MY_HEAP_FORCE_NONPAGED
90 
91 #ifdef MY_USE_INTERNAL_MEMMANAGER
92 
93 #ifdef MY_HEAP_TRACK_OWNERS
94   #ifdef MY_HEAP_TRACK_REF
95     #define MyAllocatePool__(type,size) MyAllocatePool(MyFixMemType(type), MyAlignSize__(size), UDF_BUG_CHECK_ID, __LINE__, NULL)
96     #define MyAllocatePoolTag__(type,size,tag) MyAllocatePool(MyFixMemType(type), MyAlignSize__(size), UDF_BUG_CHECK_ID, __LINE__, (PCHAR)(tag))
97   #else
98     #define MyAllocatePool__(type,size) MyAllocatePool(MyFixMemType(type), MyAlignSize__(size), UDF_BUG_CHECK_ID, __LINE__)
99     #define MyAllocatePoolTag__(type,size,tag) MyAllocatePool(MyFixMemType(type), MyAlignSize__(size), UDF_BUG_CHECK_ID, __LINE__)
100   #endif //MY_HEAP_TRACK_REF
101 #else //MY_HEAP_TRACK_OWNERS
102   #ifdef MY_HEAP_TRACK_REF
103     #define MyAllocatePool__(type,size) MyAllocatePool(MyFixMemType(type), MyAlignSize__(size), NULL)
104     #define MyAllocatePoolTag__(type,size,tag) MyAllocatePool(MyFixMemType(type), MyAlignSize__(size), (PCHAR)(tag))
105   #else
106     #define MyAllocatePool__(type,size) MyAllocatePool(MyFixMemType(type), MyAlignSize__(size))
107     #define MyAllocatePoolTag__(type,size,tag) MyAllocatePool(MyFixMemType(type), MyAlignSize__(size))
108   #endif //MY_HEAP_TRACK_REF
109 #endif //MY_HEAP_TRACK_OWNERS
110 
111 #define MyFreePool__(addr) MyFreePool((PCHAR)(addr))
112 
113 #ifdef MY_HEAP_TRACK_OWNERS
114 #define MyReallocPool__(addr, len, pnewaddr, newlen) MyReallocPool((PCHAR)(addr), MyAlignSize__(len), pnewaddr, MyAlignSize__(newlen), UDF_BUG_CHECK_ID, __LINE__)
115 #else
116 #define MyReallocPool__(addr, len, pnewaddr, newlen) MyReallocPool((PCHAR)(addr), MyAlignSize__(len), pnewaddr, MyAlignSize__(newlen))
117 #endif
118 
119 #ifdef UDF_DBG
120 LONG
121 MyFindMemBaseByAddr(
122     PCHAR addr
123     );
124 
125 #define MyCheckArray(base, index) \
126     ASSERT(MyFindMemBaseByAddr((PCHAR)(base)) == MyFindMemBaseByAddr((PCHAR)(base+(index))))
127 
128 #else
129 #define MyCheckArray(base, index)
130 #endif //UDF_DBG
131 
132 
133 #else //MY_USE_INTERNAL_MEMMANAGER
134 
135 #ifndef MY_USE_ALIGN
136 #undef  MyAlignSize__
137 #define MyAlignSize__(size) (size)
138 #endif
139 
140 BOOLEAN inline MyAllocInit(VOID) {return TRUE;}
141 #define MyAllocRelease()
142 
143 #ifndef MY_MEM_BOUNDS_CHECK
144 
145   #ifdef TRACK_SYS_ALLOC_CALLERS
146     #define MyAllocatePool__(type,size) DebugAllocatePool(NonPagedPool,MyAlignSize__(size), UDF_BUG_CHECK_ID, __LINE__)
147     #define MyAllocatePoolTag__(type,size,tag) DebugAllocatePool(NonPagedPool,MyAlignSize__(size), UDF_BUG_CHECK_ID, __LINE__)
148   #else //TRACK_SYS_ALLOC_CALLERS
149     #define MyAllocatePool__(type,size) DbgAllocatePoolWithTag(NonPagedPool,MyAlignSize__(size), 'fNWD')
150     #define MyAllocatePoolTag__(type,size,tag) DbgAllocatePoolWithTag(NonPagedPool,MyAlignSize__(size), tag)
151   #endif //TRACK_SYS_ALLOC_CALLERS
152   #define MyFreePool__(addr) DbgFreePool((PCHAR)(addr))
153 
154 #else //MY_MEM_BOUNDS_CHECK
155 
156   #ifdef TRACK_SYS_ALLOC_CALLERS
157     #define MyAllocatePool_(type,size) DebugAllocatePool(NonPagedPool,MyAlignSize__(size), UDF_BUG_CHECK_ID, __LINE__)
158     #define MyAllocatePoolTag_(type,size,tag) DebugAllocatePool(NonPagedPool,MyAlignSize__(size), UDF_BUG_CHECK_ID, __LINE__)
159   #else //TRACK_SYS_ALLOC_CALLERS
160     #define MyAllocatePool_(type,size) DbgAllocatePoolWithTag(NonPagedPool,MyAlignSize__(size), 'mNWD')
161     #define MyAllocatePoolTag_(type,size,tag) DbgAllocatePoolWithTag(NonPagedPool,MyAlignSize__(size), tag)
162   #endif //TRACK_SYS_ALLOC_CALLERS
163   #define MyFreePool_(addr) DbgFreePool((PCHAR)(addr))
164 
165 #define MyAllocatePool__(type,size) MyAllocatePoolTag__(type,size,'mNWD')
166 
167 /*
168 PVOID inline MyAllocatePool__(ULONG type, ULONG len) {
169     PCHAR newaddr;
170     ULONG i;
171 //    newaddr = (PCHAR)MyAllocatePool_(type, len+MY_HEAP_ALIGN+1);
172 #ifdef TRACK_SYS_ALLOC_CALLERS
173     newaddr = (PCHAR)DebugAllocatePool(type,len+MY_HEAP_ALIGN+1, 0x202, __LINE__);
174 #else //TRACK_SYS_ALLOC_CALLERS
175     newaddr = (PCHAR)MyAllocatePool_(type,len+MY_HEAP_ALIGN+1);
176 #endif //TRACK_SYS_ALLOC_CALLERS
177     if(!newaddr)
178         return NULL;
179     for(i=0; i<MY_HEAP_ALIGN+1; i++) {
180         newaddr[len+i] = (UCHAR)('A'+i);
181     }
182     return newaddr;
183 }
184 */
185 
186 PVOID inline MyAllocatePoolTag__(ULONG type, ULONG len, /*PCHAR*/ULONG tag) {
187     PCHAR newaddr;
188     ULONG i;
189 //    newaddr = (PCHAR)MyAllocatePoolTag_(type, len+MY_HEAP_ALIGN+1, tag);
190 #ifdef TRACK_SYS_ALLOC_CALLERS
191     newaddr = (PCHAR)DebugAllocatePool(type,len+MY_HEAP_ALIGN+1, 0x202, __LINE__);
192 #else //TRACK_SYS_ALLOC_CALLERS
193     newaddr = (PCHAR)MyAllocatePoolTag_(type,len+MY_HEAP_ALIGN+1, tag);
194 #endif //TRACK_SYS_ALLOC_CALLERS
195     if(!newaddr)
196         return NULL;
197     for(i=0; i<MY_HEAP_ALIGN+1; i++) {
198         newaddr[len+i] = (UCHAR)('A'+i);
199     }
200     return newaddr;
201 }
202 
203 VOID inline MyFreePool__(PVOID addr) {
204     PCHAR newaddr;
205 //    ULONG i;
206     newaddr = (PCHAR)addr;
207     if(!newaddr) {
208         __asm int 3;
209         return;
210     }
211 /*
212     for(i=0; i<MY_HEAP_ALIGN+1; i++) {
213         if(newaddr[len+i] != (UCHAR)('A'+i)) {
214             __asm int 3;
215             break;
216         }
217     }
218 */
219     MyFreePool_(newaddr);
220 }
221 
222 #endif //MY_MEM_BOUNDS_CHECK
223 
224 /* This function just scares the hell out of GCC */
225 #if defined(__GNUC__) && !defined(__clang__)
226 #pragma GCC diagnostic push
227 #pragma GCC diagnostic ignored "-Wstringop-overflow"
228 #endif
229 
230 ULONG inline MyReallocPool__(PCHAR addr, ULONG len, PCHAR *pnewaddr, ULONG newlen) {
231     ULONG _len, _newlen;
232     _newlen = MyAlignSize__(newlen);
233     _len = MyAlignSize__(len);
234     PCHAR newaddr;
235 
236     ASSERT(len && newlen);
237 
238 #ifdef MY_MEM_BOUNDS_CHECK
239     ULONG i;
240 
241     for(i=0; i<MY_HEAP_ALIGN+1; i++) {
242         if((UCHAR)(addr[len+i]) != (UCHAR)('A'+i)) {
243             __asm int 3;
244             break;
245         }
246     }
247 #endif //MY_MEM_BOUNDS_CHECK
248 
249     if ((_newlen != _len)
250 #ifdef MY_MEM_BOUNDS_CHECK
251         || TRUE
252 #endif //MY_MEM_BOUNDS_CHECK
253     ) {
254 #ifdef TRACK_SYS_ALLOC_CALLERS
255         newaddr = (PCHAR)DebugAllocatePool(NonPagedPool,_newlen, 0x202, __LINE__);
256 #else //TRACK_SYS_ALLOC_CALLERS
257         newaddr = (PCHAR)MyAllocatePool__(NonPagedPool,_newlen);
258 #endif //TRACK_SYS_ALLOC_CALLERS
259         if (!newaddr) {
260             __debugbreak();
261             *pnewaddr = addr;
262             return 0;
263         }
264 #ifdef MY_MEM_BOUNDS_CHECK
265         for(i=0; i<MY_HEAP_ALIGN+1; i++) {
266             newaddr[newlen+i] = (UCHAR)('A'+i);
267         }
268 #endif //MY_MEM_BOUNDS_CHECK
269         *pnewaddr = newaddr;
270         if(_newlen <= _len) {
271             RtlCopyMemory(newaddr, addr, newlen);
272         } else {
273             RtlCopyMemory(newaddr, addr, len);
274             RtlZeroMemory(newaddr+len, _newlen - len);
275         }
276 #ifdef MY_MEM_BOUNDS_CHECK
277         for(i=0; i<MY_HEAP_ALIGN+1; i++) {
278             if((UCHAR)(newaddr[newlen+i]) != (UCHAR)('A'+i)) {
279                 __asm int 3;
280                 break;
281             }
282         }
283 #endif //MY_MEM_BOUNDS_CHECK
284 
285         MyFreePool__(addr);
286     } else {
287         *pnewaddr = addr;
288     }
289     if(newlen > len) {
290         //RtlZeroMemory(newaddr+len, newlen - len);
291     }
292 /*
293 #ifdef MY_MEM_BOUNDS_CHECK
294     for(i=0; i<MY_HEAP_ALIGN+1; i++) {
295         newaddr[newlen+i] = (UCHAR)('A'+i);
296     }
297 #endif //MY_MEM_BOUNDS_CHECK
298 */
299     return newlen;
300 }
301 #if defined(__GNUC__) && !defined(__clang__)
302 #pragma GCC diagnostic pop
303 #endif
304 
305 #ifndef MY_USE_ALIGN
306 #undef  MyAlignSize__
307 #define MyAlignSize__(size) (((size)+MY_HEAP_ALIGN)&(~MY_HEAP_ALIGN))
308 #endif
309 
310 #define MyCheckArray(base, index)
311 
312 #endif // MY_USE_INTERNAL_MEMMANAGER
313 
314 #endif // __MY_MEM_TOOLS_H__
315