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 ULONG inline MyReallocPool__(PCHAR addr, ULONG len, PCHAR *pnewaddr, ULONG newlen) {
225     ULONG _len, _newlen;
226     _newlen = MyAlignSize__(newlen);
227     _len = MyAlignSize__(len);
228     PCHAR newaddr;
229 
230     ASSERT(len && newlen);
231 
232 #ifdef MY_MEM_BOUNDS_CHECK
233     ULONG i;
234 
235     for(i=0; i<MY_HEAP_ALIGN+1; i++) {
236         if((UCHAR)(addr[len+i]) != (UCHAR)('A'+i)) {
237             __asm int 3;
238             break;
239         }
240     }
241 #endif //MY_MEM_BOUNDS_CHECK
242 
243     if ((_newlen != _len)
244 #ifdef MY_MEM_BOUNDS_CHECK
245         || TRUE
246 #endif //MY_MEM_BOUNDS_CHECK
247     ) {
248 #ifdef TRACK_SYS_ALLOC_CALLERS
249         newaddr = (PCHAR)DebugAllocatePool(NonPagedPool,_newlen, 0x202, __LINE__);
250 #else //TRACK_SYS_ALLOC_CALLERS
251         newaddr = (PCHAR)MyAllocatePool__(NonPagedPool,_newlen);
252 #endif //TRACK_SYS_ALLOC_CALLERS
253         if (!newaddr) {
254             __debugbreak();
255             *pnewaddr = addr;
256             return 0;
257         }
258 #ifdef MY_MEM_BOUNDS_CHECK
259         for(i=0; i<MY_HEAP_ALIGN+1; i++) {
260             newaddr[newlen+i] = (UCHAR)('A'+i);
261         }
262 #endif //MY_MEM_BOUNDS_CHECK
263         *pnewaddr = newaddr;
264         if(_newlen <= _len) {
265             RtlCopyMemory(newaddr, addr, newlen);
266         } else {
267             RtlCopyMemory(newaddr, addr, len);
268             RtlZeroMemory(newaddr+len, newlen - len);
269         }
270 #ifdef MY_MEM_BOUNDS_CHECK
271         for(i=0; i<MY_HEAP_ALIGN+1; i++) {
272             if((UCHAR)(newaddr[newlen+i]) != (UCHAR)('A'+i)) {
273                 __asm int 3;
274                 break;
275             }
276         }
277 #endif //MY_MEM_BOUNDS_CHECK
278 
279         MyFreePool__(addr);
280     } else {
281         *pnewaddr = addr;
282     }
283     if(newlen > len) {
284         //RtlZeroMemory(newaddr+len, newlen - len);
285     }
286 /*
287 #ifdef MY_MEM_BOUNDS_CHECK
288     for(i=0; i<MY_HEAP_ALIGN+1; i++) {
289         newaddr[newlen+i] = (UCHAR)('A'+i);
290     }
291 #endif //MY_MEM_BOUNDS_CHECK
292 */
293     return newlen;
294 }
295 
296 #ifndef MY_USE_ALIGN
297 #undef  MyAlignSize__
298 #define MyAlignSize__(size) (((size)+MY_HEAP_ALIGN)&(~MY_HEAP_ALIGN))
299 #endif
300 
301 #define MyCheckArray(base, index)
302 
303 #endif // MY_USE_INTERNAL_MEMMANAGER
304 
305 #endif // __MY_MEM_TOOLS_H__
306