source: trunk/Utilities/VirtualHeapPtr/VirtualHeapPtr.h @ 35

Last change on this file since 35 was 35, checked in by roman, 11 years ago
  • Property svn:keywords set to Id
File size: 8.4 KB
Line 
1////////////////////////////////////////////////////////////
2// Copyright (C) Roman Ryltsov, 2008-2011
3// Created by Roman Ryltsov roman@alax.info
4//
5// $Id: VirtualHeapPtr.h 35 2011-11-05 19:23:11Z roman $
6
7#pragma once
8
9#include <psapi.h>
10
11#pragma comment(lib, "psapi.lib")
12
13////////////////////////////////////////////////////////////
14// CGlobalVirtualAllocator
15
16class CGlobalVirtualAllocator
17{
18private:
19        SIZE_T m_nPageSize;
20
21public:
22// CGlobalVirtualAllocator
23        CGlobalVirtualAllocator() throw() :
24                m_nPageSize(4 << 10) // 4L
25        {
26                PERFORMANCE_INFORMATION Information;
27                ZeroMemory(&Information, sizeof Information);
28                ATLVERIFY(GetPerformanceInfo(&Information, sizeof Information));
29                m_nPageSize = Information.PageSize;
30                ATLASSERT(m_nPageSize);
31                ATLASSERT(!(m_nPageSize & (m_nPageSize - 1)));
32        }
33        SIZE_T GetPageSize() const throw()
34        {
35                return m_nPageSize;
36        }
37        SIZE_T Align(SIZE_T nDataSize) throw()
38        {
39                const SIZE_T nPageSize = GetPageSize();
40                ATLASSERT(nPageSize);
41                ATLASSERT(!(nPageSize & (nPageSize - 1)));
42                return (nDataSize + nPageSize - 1) & ~(nPageSize - 1);
43        }
44        static SIZE_T StaticAlign(SIZE_T nDataSize) throw();
45};
46
47__declspec(selectany) CGlobalVirtualAllocator g_GlobalVirtualAllocator;
48
49inline SIZE_T CGlobalVirtualAllocator::StaticAlign(SIZE_T nDataSize) throw()
50{
51        return g_GlobalVirtualAllocator.Align(nDataSize);
52}
53
54////////////////////////////////////////////////////////////
55// CVirtualAllocator
56
57class CVirtualAllocator
58{
59public:
60// CVirtualAllocator
61        static VOID* Allocate(_In_ SIZE_T nDataSize)
62        {
63                if(!nDataSize)
64                        return NULL;
65                VOID* pvData = VirtualAlloc(NULL, CGlobalVirtualAllocator::StaticAlign(nDataSize), MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
66                ATLENSURE_THROW(pvData, AtlHresultFromLastError());
67                return pvData;
68        }
69        static VOID* Reallocate(_In_opt_ VOID* pvData, _In_ SIZE_T nDataSize)
70        {
71                MEMORY_BASIC_INFORMATION DataInformation;
72                if(pvData)
73                {
74                        ATLVERIFY(VirtualQuery(pvData, &DataInformation, sizeof DataInformation));
75                        if(nDataSize <= DataInformation.RegionSize)
76                                return pvData;
77                }
78                VOID* pvNewData = NULL;
79                if(nDataSize)
80                {
81                        pvNewData = VirtualAlloc(NULL, CGlobalVirtualAllocator::StaticAlign(nDataSize), MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
82                        ATLENSURE_THROW(pvNewData, AtlHresultFromLastError());
83                        _ATLTRY
84                        {
85                                if(pvNewData && pvData)
86                                {
87                                        ATLASSERT(DataInformation.AllocationProtect == PAGE_READWRITE);
88                                        SIZE_T nCopyDataSize = nDataSize;
89                                        if(nCopyDataSize > DataInformation.RegionSize)
90                                                nCopyDataSize = DataInformation.RegionSize;
91                                        Checked::memcpy_s(pvNewData, nDataSize, pvData, nCopyDataSize);
92                                }
93                        }
94                        _ATLCATCHALL()
95                        {
96                                Free(pvData);
97                                _ATLRETHROW;
98                        }
99                }
100                if(pvData)
101                        Free(pvData);
102                return pvNewData;
103        }
104        static VOID Free(_In_opt_ VOID* pvData) throw()
105        {
106                if(!pvData)
107                        return;
108                ATLVERIFY(VirtualFree(pvData, 0, MEM_RELEASE));
109        }
110};
111
112////////////////////////////////////////////////////////////
113// CVirtualHeapPtr
114
115template <typename T>
116class CVirtualHeapPtr :
117        public CHeapPtr<T, CVirtualAllocator>
118{
119public:
120// CVirtualHeapPtr
121        CVirtualHeapPtr() throw()
122        {
123        }
124        explicit CVirtualHeapPtr(_In_ T* pData) throw() :
125                CHeapPtr<T, CVirtualAllocator>(pData)
126        {
127        }
128        VOID SetProtection(DWORD nProtection)
129        {
130                if(!m_pData)
131                        return;
132                MEMORY_BASIC_INFORMATION DataInformation;
133                ATLENSURE_THROW(VirtualQuery(m_pData, &DataInformation, sizeof DataInformation), AtlHresultFromLastError());
134                DWORD nCurrentProtection;
135                ATLENSURE_THROW(VirtualProtect(m_pData, DataInformation.RegionSize, nProtection, &nCurrentProtection), AtlHresultFromLastError());
136        }
137};
138
139////////////////////////////////////////////////////////////
140// CDebugAllocatorTraits
141
142class CDebugAllocatorTraits
143{
144public:
145        static SIZE_T g_nHeadSanityPageCount;
146        static SIZE_T g_nTailSanityPageCount;
147
148public:
149// CDebugAllocatorTraits
150};
151
152__declspec(selectany) SIZE_T CDebugAllocatorTraits::g_nHeadSanityPageCount = 1;
153__declspec(selectany) SIZE_T CDebugAllocatorTraits::g_nTailSanityPageCount = 1;
154
155////////////////////////////////////////////////////////////
156// CDebugAllocatorT, CDebugAllocator
157
158template <typename _Traits = CDebugAllocatorTraits>
159class CDebugAllocatorT
160{
161public:
162        typedef _Traits CTraits;
163
164        ////////////////////////////////////////////////////////
165        // CDescriptor
166
167        class CDescriptor
168        {
169        public:
170                SIZE_T m_nAllocationSize;
171                VOID* m_pvData;
172                SIZE_T m_nDataSize;
173
174        public:
175        // CDescriptor
176                static CDescriptor* Allocate(SIZE_T nDataSize)
177                {
178                        ATLASSERT(nDataSize > 0);
179                        const SIZE_T nPageSize = g_GlobalVirtualAllocator.GetPageSize();
180                        const SIZE_T nAllocationSize = 
181                                g_GlobalVirtualAllocator.Align(sizeof (CDescriptor)) +
182                                CTraits::g_nHeadSanityPageCount * nPageSize +
183                                g_GlobalVirtualAllocator.Align(nDataSize) +
184                                CTraits::g_nTailSanityPageCount * nPageSize +
185                                0;
186                        VOID* pvData = VirtualAlloc(NULL, nAllocationSize, MEM_COMMIT | MEM_RESERVE, PAGE_NOACCESS);
187                        ATLENSURE_THROW(pvData, AtlHresultFromLastError());
188                        CDescriptor* pDescriptor = (CDescriptor*) pvData;
189                        DWORD nCurrentProtection;
190                        ATLVERIFY(VirtualProtect(pDescriptor, sizeof *pDescriptor, PAGE_READWRITE, &nCurrentProtection));
191                        pDescriptor->Initialize(nAllocationSize, nDataSize);
192                        return pDescriptor;
193                }
194                static CDescriptor* FromData(VOID* pvData)
195                {
196                        ATLASSERT(pvData);
197                        CDescriptor* pDescriptor = (CDescriptor*) ((BYTE*) pvData - CTraits::g_nHeadSanityPageCount * g_GlobalVirtualAllocator.GetPageSize() - g_GlobalVirtualAllocator.Align(sizeof (CDescriptor)));
198                        return pDescriptor;
199                }
200                VOID Initialize(SIZE_T nAllocationSize, SIZE_T nDataSize)
201                {
202                        const SIZE_T nPageSize = g_GlobalVirtualAllocator.GetPageSize();
203                        m_nAllocationSize = nAllocationSize;
204                        m_pvData = (BYTE*) this + g_GlobalVirtualAllocator.Align(sizeof *this) + CTraits::g_nHeadSanityPageCount * nPageSize;
205                        m_nDataSize = nDataSize;
206                        DWORD nCurrentProtection;
207                        ATLVERIFY(VirtualProtect(m_pvData, m_nDataSize, PAGE_READWRITE, &nCurrentProtection));
208                        //ATLASSERT(!IsBadWritePtr(m_pvData, m_nDataSize));
209                        ATLVERIFY(VirtualProtect(this, sizeof *this, PAGE_READONLY, &nCurrentProtection));
210                }
211                VOID Terminate()
212                {
213                        ATLVERIFY(VirtualFree(this, 0, MEM_RELEASE));
214                }
215                VOID* GetData() const throw()
216                {
217                        return m_pvData;
218                }
219                SIZE_T GetDataSize() const throw()
220                {
221                        return m_nDataSize;
222                }
223        };
224
225public:
226// CDebugAllocator
227        static VOID* Allocate(_In_ SIZE_T nDataSize)
228        {
229                if(!nDataSize)
230                        return NULL;
231                CDescriptor* pDescriptor = CDescriptor::Allocate(nDataSize);
232                ATLASSERT(pDescriptor);
233                return pDescriptor->GetData();
234        }
235        static VOID* Reallocate(_In_opt_ VOID* pvData, _In_ SIZE_T nDataSize)
236        {
237                CDescriptor* pDescriptor;
238                if(pvData)
239                {
240                        pDescriptor = CDescriptor::FromData(pvData);
241                        if(nDataSize == pDescriptor->GetDataSize())
242                                return pvData;
243                } else
244                        pDescriptor = NULL;
245                VOID* pvNewData = NULL;
246                if(nDataSize)
247                {
248                        CDescriptor* pNewDescriptor = CDescriptor::Allocate(nDataSize);
249                        ATLASSERT(pNewDescriptor);
250                        _ATLTRY
251                        {
252                                if(pNewDescriptor && pDescriptor)
253                                {
254                                        SIZE_T nCopyDataSize = nDataSize;
255                                        if(nCopyDataSize > pDescriptor->GetDataSize())
256                                                nCopyDataSize = pDescriptor->GetDataSize();
257                                        Checked::memcpy_s(pNewDescriptor->GetData(), pNewDescriptor->GetDataSize(), pDescriptor->GetData(), nCopyDataSize);
258                                }
259                        }
260                        _ATLCATCHALL()
261                        {
262                                pNewDescriptor->Terminate();
263                                _ATLRETHROW;
264                        }
265                        pvNewData = pNewDescriptor->GetData();
266                }
267                if(pvData)
268                        Free(pvData);
269                return pvNewData;
270        }
271        static VOID Free(_In_opt_ VOID* pvData) throw()
272        {
273                if(!pvData)
274                        return;
275                CDescriptor* pDescriptor = CDescriptor::FromData(pvData);
276                ATLASSERT(pDescriptor);
277                pDescriptor->Terminate();
278        }
279};
280
281typedef CDebugAllocatorT<> CDebugAllocator;
282
283////////////////////////////////////////////////////////////
284// CVirtualHeapPtr
285
286template <typename T>
287class CDebugHeapPtr :
288        public CHeapPtr<T, CDebugAllocator>
289{
290public:
291// CDebugHeapPtr
292        CDebugHeapPtr() throw()
293        {
294        }
295        explicit CDebugHeapPtr(_In_ T* pData) throw() :
296                CHeapPtr<T, CDebugAllocator>(pData)
297        {
298        }
299        VOID SetProtection(DWORD nProtection)
300        {
301                if(!m_pData)
302                        return;
303                CDebugAllocator::CDescriptor* pDescriptor = CDebugAllocator::CDescriptor::FromData(m_pData);
304                ATLASSERT(pDescriptor);
305                DWORD nCurrentProtection;
306                ATLENSURE_THROW(VirtualProtect(pDescriptor->GetData(), pDescriptor->GetDataSize(), nProtection, &nCurrentProtection), AtlHresultFromLastError());
307        }
308};
309
Note: See TracBrowser for help on using the repository browser.