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
MyAllocInit(VOID)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
MyAllocatePoolTag__(ULONG type,ULONG len,ULONG tag)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
MyFreePool__(PVOID addr)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
MyReallocPool__(PCHAR addr,ULONG len,PCHAR * pnewaddr,ULONG newlen)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