source: trunk/Utilities/VirtualHeapPtr/VirtualHeapPtr.h

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